Stopping and stacking

Multiply a vector by itself and differentiate: the answer is 2x, one term per factor. Now wrap the second factor in lax.stop_gradient. The value is unchanged at 14.0, and the gradient is x rather than 2x.

The gradient jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → shows what happened with no interpretation required. stop_gradient a runs in the forward half, the product uses its output, and the backward half has exactly one multiply, against that stopped value. The term that would have differentiated the second factor was never emitted, because there was no rule to emit it through.

Reading it as deletion from a program, rather than as zeroing something at runtime, is what makes the effect predictable. Nothing is multiplied by zero and nothing is masked. One path through the derived program does not exist.

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

def cut(x):
    return jnp.sum(x * lax.stop_gradient(x))

x = jnp.array([1.0, 2.0, 3.0])
print(cut(x), jax.grad(cut)(x))
print(jax.grad(lambda v: jnp.sum(v * v))(x))
print(jax.make_jaxpr(jax.grad(cut))(x))

# 14.0 [1. 2. 3.]
# [2. 4. 6.]
# { lambda ; a:f32[3]. let
#     b:f32[3] = stop_gradient a
#     c:f32[3] = mul a b
#     _:f32[] = reduce_sum[axes=(0,)] c
#     d:f32[3] = broadcast_in_dim[
#       broadcast_dimensions=()
#       shape=(3,)
#       sharding=None
#     ] 1.0
#     e:f32[3] = mul d b
#   in (e,) }

Rounding has a derivative of zero wherever it has one at all, so a model with a rounding step in it learns nothing through that step. jnp.round differentiates to zeros, as the run below confirms.

The construction that gets around it is x + stop_gradient(round(x) - x). The value is exactly round(x), since the added term is the difference. The derivative is the identity, because the only part of the expression with a live gradient is the leading x. Quantizers and discrete samplers are built from this pattern.

Be precise about the price. The gradient you get is not an approximation of this function's gradient; it is the exact gradient of a different function, chosen because it is useful for optimization. That is a modelling decision made in the code, and the only place it is written down is the stop_gradient call itself.

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

def ste(x):
    return x + lax.stop_gradient(jnp.round(x) - x)

y = jnp.array([0.2, 1.7, -0.5])
print(ste(y))
print(jax.grad(lambda v: jnp.sum(ste(v)))(y))
print(jax.grad(lambda v: jnp.sum(jnp.round(v)))(y))

# [0. 2. 0.]
# [1. 1. 1.]
# [0. 0. 0.]

jax.grad returns a function, and that function is a program of the same kind as the one it came from, so grad applies to it again. Stack four of them on sin and what prints is sin, cos, -sin, -cos at the same point, matching the library's own values digit for digit.

Nothing about the second application is special-cased. The gradient program was assembled out of the same per-primitive rules as the program it came from, so those rules apply to it in turn, which is also what gives an order-2 check_grads something to check.

Where the stacking gets interesting is the order you stack in, and what the composition costs. The next two sections are those two questions.

import jax
import jax.numpy as jnp

d0 = jnp.sin
d1 = jax.grad(d0)
d2 = jax.grad(d1)
d3 = jax.grad(d2)
print(d0(0.7), d1(0.7), d2(0.7), d3(0.7))
print(jnp.sin(0.7), jnp.cos(0.7), -jnp.sin(0.7), -jnp.cos(0.7))

# 0.64421767 0.7648422 -0.64421767 -0.7648422
# 0.64421767 0.7648422 -0.64421767 -0.7648422

jax.hessian is not separate machinery. Its whole body in the wheel is one expression, jacfwd(jacrev(fun, ...), ...), and reading the order gives you the shape. jacrev produces the gradient as a row, then jacfwd pushes one tangent per input through that gradient program and fills the matrix column by column.

Whether the other order would cost more is a question you can put to the compiler instead of arguing about. Priced with the counter from lesson two, jacfwd(jacrev(loss)) and jacrev(jacrev(loss)) come out the same on a 32-input loss: 147,552 flops each, 12,288 bytes of scratch each, and the same matrix. On this function XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → does not separate them, so take the order in the body as the one JAX fixes for you and not as a measured win.

The order also explains the refusal from the previous lesson. A Hessian needs a forward pass over a gradient, so a function whose only custom rule is a backward one cannot supply one at all.

  return jacfwd(jacrev(fun, argnums, has_aux=has_aux, holomorphic=holomorphic),
                argnums, has_aux=has_aux, holomorphic=holomorphic)

Stacking transforms is cheap to write and expensive to run, and the same flop counter from lesson two prices it. On a 32-input loss the forward pass is 2,111 flops, the gradient 4,256, and the full Hessian 147,552: about 35 gradients to fill a 32 by 32 matrix.

Most methods that want curvature never want the matrix. Newton-style solvers, conjugate gradient, and Gauss-Newton approximations all consume the Hessian one product at a time, and one product is jax.jvp(jax.grad(loss), (x,), (u,))[1]: push a tangent forward through the gradient program. That costs 6,624 flops, a little over one and a half gradients, and 22 times less than building the matrix.

The two agree to 7.5e-09 in the largest component, which is float32 rounding rather than a different answer. When the input is a model rather than a 32-vector, the matrix is not merely expensive but unrepresentable, and the product is the only form that exists.

A Hessian-vector product is a jvp of a grad. Both transforms you already have, composed in the order that skips the matrix.

import jax
import jax.numpy as jnp

def flops(fn, *args):
    return jax.jit(fn).lower(*args).compile().cost_analysis()[0]['flops']

W = jnp.ones((32, 32)) * 0.01

def loss(x):
    return jnp.sum(jnp.tanh(x @ W) ** 2)

x = jnp.ones(32)
u = jnp.arange(32.) / 32.0

print(flops(loss, x), flops(jax.grad(loss), x))
print(flops(jax.hessian(loss), x))
print(flops(lambda x, u: jax.jvp(jax.grad(loss), (x,), (u,))[1], x, u))

hvp = jax.jvp(jax.grad(loss), (x,), (u,))[1]
print(float(jnp.max(jnp.abs(hvp - jax.hessian(loss)(x) @ u))))

# 2111.0 4256.0
# 147552.0
# 6624.0
# 7.450580596923828e-09