TensorFlow
tf.gradients
machine learning
gradient computation
deep learning

Separate gradients in tf.gradients

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

When people ask for "separate gradients" in tf.gradients, they usually mean one of two things: gradients with respect to multiple variables, or separate gradient contributions from different loss terms. TensorFlow can do both, but it is important to understand that tf.gradients aggregates contributions by default unless you ask for them separately.

What tf.gradients Returns

In TensorFlow 1 style graph mode, tf.gradients(ys, xs) computes the derivative of ys with respect to each tensor in xs.

If you pass multiple target tensors in ys, TensorFlow sums their contributions. Conceptually, this:

python
tf.compat.v1.gradients([loss_a, loss_b], [w])

is treated like the gradient of loss_a + loss_b with respect to w.

That default behavior is convenient for training, but it is exactly why people sometimes think the gradients are not "separate."

Separate Gradients for Different Variables

If you only need one gradient per variable, pass all variables in xs.

python
1import tensorflow as tf
2
3tf.compat.v1.disable_eager_execution()
4
5w1 = tf.Variable(2.0, name="w1")
6w2 = tf.Variable(3.0, name="w2")
7loss = w1 * w1 + 2.0 * w2
8
9grads = tf.compat.v1.gradients(loss, [w1, w2])
10
11with tf.compat.v1.Session() as sess:
12    sess.run(tf.compat.v1.global_variables_initializer())
13    g1, g2 = sess.run(grads)
14    print("dw1:", g1)
15    print("dw2:", g2)

This is already "separate" in the variable dimension. You get one gradient result for w1 and one for w2.

Separate Contributions from Different Loss Terms

If you want to know how much each loss term contributes, compute gradients for each term independently.

python
1import tensorflow as tf
2
3tf.compat.v1.disable_eager_execution()
4
5w = tf.Variable(2.0, name="w")
6loss_a = w * w
7loss_b = 3.0 * w
8
9grad_a = tf.compat.v1.gradients(loss_a, w)[0]
10grad_b = tf.compat.v1.gradients(loss_b, w)[0]
11grad_total = tf.compat.v1.gradients(loss_a + loss_b, w)[0]
12
13with tf.compat.v1.Session() as sess:
14    sess.run(tf.compat.v1.global_variables_initializer())
15    ga, gb, gt = sess.run([grad_a, grad_b, grad_total])
16    print("grad_a:", ga)
17    print("grad_b:", gb)
18    print("grad_total:", gt)

This gives you the cleanest view of each term. In this example, grad_total equals grad_a + grad_b.

When to Use stop_gradient

Sometimes you want one branch to contribute to the forward pass but not to the backward pass. That is what tf.stop_gradient is for.

python
1import tensorflow as tf
2
3tf.compat.v1.disable_eager_execution()
4
5w = tf.Variable(2.0)
6term_a = w * w
7term_b = tf.stop_gradient(3.0 * w)
8loss = term_a + term_b
9
10grad = tf.compat.v1.gradients(loss, w)[0]
11
12with tf.compat.v1.Session() as sess:
13    sess.run(tf.compat.v1.global_variables_initializer())
14    print(sess.run(grad))

Here, term_b affects the loss value but contributes nothing to the gradient.

Use GradientTape in TensorFlow 2

If you are writing new TensorFlow code, use tf.GradientTape instead of tf.gradients. The same idea applies: ask for each contribution explicitly if you need them separately.

python
1import tensorflow as tf
2
3w = tf.Variable(2.0)
4
5with tf.GradientTape(persistent=True) as tape:
6    loss_a = w * w
7    loss_b = 3.0 * w
8    total = loss_a + loss_b
9
10print("grad_a:", tape.gradient(loss_a, w).numpy())
11print("grad_b:", tape.gradient(loss_b, w).numpy())
12print("grad_total:", tape.gradient(total, w).numpy())

This is often easier to debug because the gradient requests read more directly.

Why the Default Aggregation Exists

Optimizers normally need the total derivative of the total loss, not each part separately. That is why TensorFlow defaults to aggregation. Separate gradients are mostly useful for:

  • debugging training behavior
  • logging loss-term influence
  • implementing custom optimization rules
  • freezing or blocking parts of the graph

For ordinary training, the default is usually what you want.

Common Pitfalls

The biggest mistake is passing multiple losses to tf.gradients and expecting a separate tensor back for each loss term. TensorFlow sums them unless you compute them in separate calls.

Another issue is forgetting that tf.gradients is a TensorFlow 1 style graph API. In TensorFlow 2, GradientTape is the normal choice.

A third problem is using stop_gradient without realizing it changes only the backward pass, not the forward value of the expression.

Summary

  • 'tf.gradients returns one gradient per tensor in xs.'
  • If ys contains multiple losses, their gradient contributions are summed by default.
  • Compute each loss gradient in separate calls if you need them separately.
  • Use tf.stop_gradient to block selected branches from contributing to backpropagation.
  • Prefer tf.GradientTape for new TensorFlow 2 code.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.