Blocking Gradients with `jax.lax.stop_gradient`

Hello! In our last lesson, we dove deep into the internals of automatic differentiation by implementing custom backward passes with jax.custom_vjp. This gave you ultimate control over the gradient computation. Now, we'll turn our attention to a related but much simpler and equally common requirement: how to completely prevent gradients from flowing through certain parts of your computational graph.

This lesson focuses on jax.lax.stop_gradient, a fundamental tool for manipulating the backpropagation process. While custom_vjp allows you to redefine the backward pass, stop_gradient lets you prune it. You'll learn how this function works and apply it to several important patterns in modern machine learning, such as creating fixed targets in self-supervised learning and implementing non-standard update rules.

This skill is crucial for implementing a wide range of advanced models where the default gradient isn't what the algorithm requires.

1. The `stop_gradient` Mechanism

At its core, jax.lax.stop_gradient is a straightforward function with a dual identity:

In the forward pass, it's the identity function. It returns its input completely unchanged: jax.lax.stop_gradient(x) simply evaluates to x.
In the backward pass, it acts as a wall. It tells JAX to treat its input as a constant, meaning no gradients will be propagated back through it. The gradient of jax.lax.stop_gradient(x) with respect to x is always zero.

This behavior allows you to selectively detach parts of your computation from the gradient tape, without altering the forward pass values. Let's examine a clear example that illustrates this.

Handling Non-Differentiable Functions JAX

This article provides a very clear, practical demonstration of stop_gradient. It compares a normal gradient calculation with one where the gradient flow is partially blocked, making the effect unambiguous.

Please read the section "Explicitly Stopping Gradients with jax.lax.stop_gradient". Pay close attention to the code block and the explanation of the resulting gradients for f_normal and f_stopped. Note how the forward pass output is identical for both functions, but the derivative is fundamentally different.

As the example shows, in the f_stopped function, the term z = jax.lax.stop_gradient(jnp.cos(x)) ensures that while the value of cos(x) is used in the forward computation, the differentiation process treats z as a pre-computed constant. The chain rule is therefore applied as if the function were y * K, where K is the constant value of z.

2. Use Case: Creating Fixed Targets

A primary application of stop_gradient is to prevent a network from "chasing its own tail" in algorithms where a model's output is used as a target for itself or another network. This is common in self-supervised learning, generative adversarial networks, and some reinforcement learning methods.

Consider the following computational graph, which represents a common pattern in self-supervised learning (like a simplified BYOL or SimSiam):

A computational graph where an input `x_t1` is processed by networks `N1` and `N2` to produce a prediction `s_tilde_t2`. Another input, `x_t2`, is processed by `N1` to produce a target `s_t2`. The loss compares the prediction to the target.

In this setup, we have two branches:

Prediction Branch (Top): An input $x_{t_{1}}$ goes through network $N_{1}$ and then network $N_{2}$ , producing $s ~_{t_{2}}$ . This could be an "online" network.
Target Branch (Bottom): A different input $x_{t_{2}}$ goes through network $N_{1}$ to produce $s_{t_{2}}$ . This could be a "target" network, which might be an older version of the online network.

The goal is to train the networks so that the prediction $s ~_{t_{2}}$ matches the target $s_{t_{2}}$ , measured by a loss function, $Loss = f (s_{t_{2}}, s ~_{t_{2}})$ .

If we compute the gradient of this loss naively, backpropagation will flow through both arguments of the loss function.

Gradient from $s ~_{t_{2}}$ flows back through $N_{2}$ and $N_{1}$ . This is what we want; it updates the prediction network to be better.
Gradient from $s_{t_{2}}$ flows back through $N_{1}$ . This is problematic. It encourages the network to change the target to make the loss smaller (a phenomenon known as model collapse), rather than improving the prediction.

The solution is to treat the target as a fixed value for the purpose of the backward pass. We use stop_gradient to cut off the gradient flow into the target branch:

\text{Loss} = f(\text{jax.lax.stop_gradient}(s_{t_2}), \tilde{s}_{t_2})

Now, gradients only flow from the loss into the prediction branch, training N1 and N2 to produce better predictions of the fixed target s_t2.

3. Use Case: Custom Learning Rules & Non-Differentiable Functions

stop_gradient is also a key component in implementing learning rules that are not pure gradients of a loss function, or for defining "gradients" for non-differentiable operations.

Temporal Difference (TD) Learning

In reinforcement learning, the TD(0) update for a value function $v_{θ} (s)$ is not a true gradient. However, it can be expressed as the gradient of a "pseudo-loss" if the target value is treated as a constant. The JAX documentation provides a canonical example of this.

Higher-order derivatives — JAX documentation

The official JAX documentation demonstrates how to use stop_gradient to implement a classic reinforcement learning update rule. This is a perfect example of applying stop_gradient for algorithmic correctness.

Please read the section "Stopping gradients". Focus on how the TD pseudo-loss is constructed and why jax.lax.stop_gradient is necessary to make the gradient of this pseudo-loss equal to the desired TD(0) update.

Straight-Through Estimator (STE)

Another clever application is the straight-through estimator, a technique to train networks that include non-differentiable operations like rounding or sign functions. The idea is to use the non-differentiable function in the forward pass but pretend its derivative is that of the identity function ( $f (x) = x$ ) in the backward pass.

This can be implemented elegantly with stop_gradient using the formula:
y_forward = f(x)
y_backward_friendly = x + jax.lax.stop_gradient(y_forward - x)

Let's analyze this:

Forward pass: The result is x + (f(x) - x) = f(x). It correctly computes the output of the non-differentiable function.
Backward pass: The gradient is $\nabla_x(x + \text{stop_gradient}(\dots)) = \nabla_x(x) + 0 = 1$ . The gradient is simply 1, as if the operation were the identity.

This technique allows gradient-based optimization to proceed through the non-differentiable step.

Higher-order derivatives — JAX documentation

This section of the JAX docs shows the code for the Straight-Through Estimator trick we just discussed.

Now, please read the short section "Straight-through estimator using stop_gradient". This example, using a step function, solidifies how this trick is implemented in just a single line of code.

Conclusion

In this lesson, we've explored jax.lax.stop_gradient, a simple but powerful primitive for controlling the flow of information during backpropagation. While it may seem less complex than the custom_vjp we saw previously, its correct application is fundamental to implementing many sophisticated machine learning algorithms.

Key Takeaways:

jax.lax.stop_gradient acts as the identity function during the forward pass but blocks gradients in the backward pass, effectively treating its argument as a constant.
It is essential for creating fixed targets in loss functions, preventing undesirable gradient flows in algorithms like self-supervised learning, GANs, and certain RL methods.
It enables the implementation of non-standard update rules (e.g., TD learning) by framing them as gradients of a carefully constructed pseudo-loss.
It is a key ingredient for techniques like the Straight-Through Estimator, which provide a "pseudo-gradient" for mathematically non-differentiable functions.

In our next lesson, we will continue exploring mechanisms that control the backward pass by examining gradient checkpointing (also known as remat). Unlike stop_gradient, which alters the computational graph for modeling purposes, checkpointing modifies the backward pass purely for performance, trading extra computation to significantly reduce memory usage.

Recommended resources