GradientTape gives different gradients depending on loss function being decorated by tf.function or not

Question

I find that the gradients computed depend on the interplay of tf.function decorators in the following way.

First I create some synthetic data for a binary classification

tf.random.set_seed(42)
np.random.seed(42)
x=tf.random.normal((2,1))
y=tf.constant(np.random.choice([0,1],2))

Then I define two loss functions that differ only in the tf.function decorator

weights=tf.constant([1.,.1])[tf.newaxis,...]

def customloss1(y_true,y_pred,sample_weight=None):
    y_true_one_hot=tf.one_hot(tf.cast(y_true,tf.uint8),2)
    y_true_scale=tf.multiply(weights,y_true_one_hot)
    return tf.reduce_mean(tf.keras.losses.categorical_crossentropy(y_true_scale,y_pred))

@tf.function
def customloss2(y_true,y_pred,sample_weight=None):
    y_true_one_hot=tf.one_hot(tf.cast(y_true,tf.uint8),2)
    y_true_scale=tf.multiply(weights,y_true_one_hot)
    return tf.reduce_mean(tf.keras.losses.categorical_crossentropy(y_true_scale,y_pred))

Then I make a very simple logistic regression model with all the bells and whistles removed to keep it simple

tf.random.set_seed(42)
np.random.seed(42)
model=tf.keras.Sequential([
    tf.keras.layers.Dense(2,use_bias=False,activation='softmax',input_shape=[1,])
])

and finally define two functions to calculate the gradients of the aforementioned loss functions with one being decorated by tf.function and the other not being decorated by it

def get_gradients1(x,y):
    with tf.GradientTape() as tape1:
        p1=model(x)
        l1=customloss1(y,p1)
    with tf.GradientTape() as tape2:
        p2=model(x)
        l2=customloss2(y,p2)

    gradients1=tape1.gradient(l1,model.trainable_variables)
    gradients2=tape2.gradient(l2,model.trainable_variables)

    return gradients1, gradients2

@tf.function
def get_gradients2(x,y):
    with tf.GradientTape() as tape1:
        p1=model(x)
        l1=customloss1(y,p1)
    with tf.GradientTape() as tape2:
        p2=model(x)
        l2=customloss2(y,p2)

    gradients1=tape1.gradient(l1,model.trainable_variables)
    gradients2=tape2.gradient(l2,model.trainable_variables)

    return gradients1, gradients2

Now when I run

get_gradients1(x,y)

I get

([],
 [])

and the gradients are equal as expected. However when I run

get_gradients2(x,y)

I get

([],
 [])

where only the second answer is correct. Thus, when my outer function is decorated I only get the correct answer from the inner function that is decorated as well. I was under the impression that decorating the outer one (which is the training loop in many applications) is sufficient but here we see its not. I want to understand why and also then how deep does one have to go to decorate the functions being used?

Added some debugging info

I added some debugging info and I show the code only for customloss2 (the other is identical)

@tf.function
def customloss2(y_true,y_pred,sample_weight=None):
    y_true_one_hot=tf.one_hot(tf.cast(y_true,tf.uint8),2)
    y_true_scale=tf.multiply(weights,y_true_one_hot)
    tf.print('customloss2',type(y_true_scale),type(y_pred))
    tf.print('y_true_scale','
',y_true_scale)
    tf.print('y_pred','
',y_pred)
    return tf.reduce_mean(tf.keras.losses.categorical_crossentropy(y_true_scale,y_pred))

and on running get_gradients1 I get

customloss1  
y_true_scale 
 [[1 0]
 [0 0.1]]
y_pred 
 [[0.510775387 0.489224613]
 [0.529191136 0.470808864]]
customloss2  
y_true_scale 
 [[1 0]
 [0 0.1]]
y_pred 
 [[0.510775387 0.489224613]
 [0.529191136 0.470808864]]

we see that the tensors for customloss1 are Eager but for customloss2 are Tensor and yet we get same value for gradients.

On the other hand when I run it on get_gradients2

customloss1  
y_true_scale 
 [[1 0]
 [0 0.1]]
y_pred 
 [[0.510775387 0.489224613]
 [0.529191136 0.470808864]]
customloss2  
y_true_scale 
 [[1 0]
 [0 0.1]]
y_pred 
 [[0.510775387 0.489224613]
 [0.529191136 0.470808864]]

we see everything is identical with no tensors being Eager and yet I get different gradients!

user_1782792 · Answer

This is a somewhat complicated issue, but it has an explanation. The problem lies within the function tf.keras.backend.categorical_crossentropy, which has a different behavior depending on whether you are running on eager or graph (tf.function) mode.

The function considers three possible situations. The first one is that you pass from_logits=True, in which case it just calls tf.nn.softmax_cross_entropy_with_logits:

if from_logits:
  return nn.softmax_cross_entropy_with_logits_v2(
      labels=target, logits=output, axis=axis)

If you give from_logits=False, which is the most common in Keras, since the output layer for categorical classification is generally a softmax, then it considers two possibilities. The first is that, if the given output value comes from a softmax operation, then it can just use the input to that operation and call tf.nn.softmax_cross_entropy_with_logits, which is preferred to compute the actual cross entropy with the softmax values because it prevents "saturated" results. However, this can only be done in graph mode, because eager mode tensors do not keep track of the operation it produced them, nevermind the inputs to that operation.

if not isinstance(output, (ops.EagerTensor, variables_module.Variable)):
  output = _backtrack_identity(output)
  if output.op.type == 'Softmax':
    # When softmax activation function is used for output operation, we
    # use logits from the softmax function directly to compute loss in order
    # to prevent collapsing zero when training.
    # See b/117284466
    assert len(output.op.inputs) == 1
    output = output.op.inputs[0]
    return nn.softmax_cross_entropy_with_logits_v2(
        labels=target, logits=output, axis=axis)

The last case is when you have given from_logits=False and either you are in eager mode or the given output tensor does not directly come from a softmax operation, in which case the only option is to compute the cross entropy from the softmax value.

# scale preds so that the class probas of each sample sum to 1
output = output / math_ops.reduce_sum(output, axis, True)
# Compute cross entropy from probabilities.
epsilon_ = _constant_to_tensor(epsilon(), output.dtype.base_dtype)
output = clip_ops.clip_by_value(output, epsilon_, 1. - epsilon_)
return -math_ops.reduce_sum(target * math_ops.log(output), axis)

The problem is that, even though these are mathematically equivalent ways to compute the cross entropy, they do not have the same precision. They are pretty much the same when logits are small, but if they get big they can diverge a lot. Here is a simple test:

import tensorflow as tf

@tf.function
def test_keras_xent(y, p, from_logits=False, mask_op=False):
    # p is always logits
    if not from_logits:
        # Compute softmax if not using logits
        p = tf.nn.softmax(p)
    if mask_op:
        # A dummy addition prevents Keras from detecting that
        # the value comes from a softmax operation
        p = p + tf.constant(0, p.dtype)
    return tf.keras.backend.categorical_crossentropy(y, p, from_logits=from_logits)

# Test
tf.random.set_seed(0)
y = tf.constant([1., 0., 0., 0.])

# Logits in [0, 1)
p = tf.random.uniform([4], minval=0, maxval=1)
tf.print(test_keras_xent(y, p, from_logits=True))
# 1.50469065
tf.print(test_keras_xent(y, p, from_logits=False, mask_op=False))
# 1.50469065
tf.print(test_keras_xent(y, p, from_logits=False, mask_op=True))
# 1.50469065

# Logits in [0, 10)
p = tf.random.uniform([4], minval=0, maxval=10)
tf.print(test_keras_xent(y, p, from_logits=True))
# 3.47569656
tf.print(test_keras_xent(y, p, from_logits=False, mask_op=False))
# 3.47569656
tf.print(test_keras_xent(y, p, from_logits=False, mask_op=True))
# 3.47569656

# Logits in [0, 100)
p = tf.random.uniform([4], minval=0, maxval=100)
tf.print(test_keras_xent(y, p, from_logits=True))
# 68.0106506
tf.print(test_keras_xent(y, p, from_logits=False, mask_op=False))
# 68.0106506
tf.print(test_keras_xent(y, p, from_logits=False, mask_op=True))
# 16.1180954

Taking your example:

import tensorflow as tf

tf.random.set_seed(42)
x = tf.random.normal((2, 1))
y = tf.constant(np.random.choice([0, 1], 2))
y1h = tf.one_hot(y, 2, dtype=x.dtype)
model = tf.keras.Sequential([
    # Linear activation because we want the logits for testing
    tf.keras.layers.Dense(2, use_bias=False, activation='linear', input_shape=[1,])
])
p = model(x)
tf.print(test_keras_xent(y1h, p, from_logits=True))
# [0.603375256 0.964639068]
tf.print(test_keras_xent(y1h, p, from_logits=False, mask_op=False))
# [0.603375256 0.964639068]
tf.print(test_keras_xent(y1h, p, from_logits=False, mask_op=True))
# [0.603375256 0.964638948]

The results here are almost identical, but you can see there is a small difference in the second value. This has in turn an effect (probably in amplified) in the computed gradients, which of course are as well "equivalent" mathematical expression but with different precision properties.

user_6057582 · Answer

It turns out this is a bug and I have raised it here.

GradientTape gives different gradients depending on loss function being decorated by tf.function or not

2 Answers