Skip to content

← JAX Conceptual Deep-Dives step 13 of 15

Hard Primitives

Selective stop_gradient

Jump to the task

Keep a value, block its derivative

lax.stop_gradient returns its input value unchanged, but tells automatic differentiation to treat that occurrence as constant. It blocks both forward- and reverse-mode derivatives through the wrapped expression, including nested differentiation.

import jax
import jax.numpy as jnp
from jax import lax

def loss(w_train, w_fixed):
    return (w_train + lax.stop_gradient(w_fixed)) ** 2

grads = jax.grad(loss, argnums=(0, 1))(1.0, 5.0)
# grads = (12.0, 0.0)
# loss(1.0, 5.0) is still 36.0

The fixed value still contributes 5 to the sum. The derivative with respect to the trainable value is therefore 2*(1+5)=12. The derivative with respect to the stopped value is zero. Without stop_gradient, both derivatives would be 12.

Stop a path, not every use of a variable

The operation affects only the expression it wraps. For example, z + lax.stop_gradient(z) has value 2*z, but autodiff reports a derivative of 1 through the unstopped occurrence of z.

For a parameter tree, jax.tree.map(lax.stop_gradient, params) applies the operation to its array leaves. Use the returned tree in the loss; this does not mutate the original tree or mark it permanently frozen. The diagram below illustrates a separate nested structure. Its numbered leaves are placeholders, not the scalar inputs in this exercise.

Loading visualization…

Blocking derivatives is only part of freezing parameters during training. An update rule can still change a parameter through existing momentum, weight decay or an explicit assignment. To keep a parameter unchanged, exclude it from those updates too.

Problem

Implement selective_stop_grad(w_train, w_freeze, x). The weights are real floating-point scalars, and x is a one-dimensional real floating-point JAX array.

Define the scalar loss sum((w_train*x + lax.stop_gradient(w_freeze))**2) and return its derivative with respect to w_train only using jax.grad. Return a scalar. Keep the stopped value in the forward calculation; do not replace it with zero.

For w_train=1.0, w_freeze=2.0, and x=[1.0, 3.0], the values before squaring are [3, 5]. The derivative is 2*3*1 + 2*5*3 = 36.

Because this function returns only the w_train derivative, omitting stop_gradient would give the same numerical answer. The joint-gradient example above shows the distinction that this output alone cannot test. Use stop_gradient here to practice controlling the derivative path.