jvp for Sensitivity Analysis

Estimate how a loss changes

A directional derivative describes how a loss responds to a small input change. For a scalar weight w and a direction dw, consider the path w + epsilon*dw. Its loss has rate of change loss_prime(w) * dw at epsilon=0. JAX calls this rate an output tangent.

import jax
import jax.numpy as jnp

x = jnp.array([1.0, 2.0])

def loss(w):
    return jnp.sum((w * x) ** 2)

value, tangent = jax.jvp(loss, (1.0,), (0.1,))
# value = 5.0
# tangent = 10.0 * 0.1 = 1.0

The linear estimate for loss(1 + epsilon*0.1) is 5 + epsilon*1. For epsilon=1, it predicts a change of 1, while the actual change is loss(1.1) - loss(1.0) = 1.05. A directional derivative is not an exact finite change. The approximation improves as the step gets smaller.

Keep the direction separate from the input

jax.jvp(loss, (w,), (dw,)) evaluates the derivative at w. It does not evaluate the loss at w + dw. Negative dw reverses the direction, and scaling dw scales the tangent. Do not normalize it.

The result is (value, tangent). The exercise needs only the second entry. Both w and dw are scalar here; zero-dimensional JAX arrays are valid too. Their floating-point dtypes must match. A length-one array is a different shape from a scalar.

JVP computes a derivative along one direction. Reverse-mode grad computes the full gradient of a scalar loss. Which is faster depends on the function and the derivatives you need; there is no guaranteed speed advantage for this scalar example.

Compute derivatives by following the operations

Automatic differentiation computes derivatives by applying the chain rule to a program's operations. It does not estimate them by evaluating nearby inputs, as finite differences do, or rearrange the whole expression using symbolic algebra.

For f(g(x)), the chain rule combines the derivative of g with the derivative of f at g(x). Repeating this process lets a framework differentiate a long sequence of operations.

Forward mode follows an input change

For a function with n inputs and m outputs, the Jacobian J is an (m, n) matrix of derivatives. Entry (i, j) describes how output i changes with input j.

Forward-mode differentiation computes Jv: the output's rate of change when the input moves in direction v. For example, if f(x) = sum(x_i^2), then at x = [1, 2] the gradient is [2, 4]. In direction v = [3, 0], the rate of change is 2 * 3 + 4 * 0 = 6.

This is a directional derivative. It gives one Jacobian column only when v selects one input coordinate, such as [1, 0].

Reverse mode works backward from an output

Reverse mode computes J^T u, where u weights the outputs. For a scalar loss, choosing u = 1 gives the gradient with respect to all inputs. This makes reverse mode useful for training: a model can have many parameters but just one loss.

Forward mode can be useful when you need sensitivity to only a few input directions. Neither mode needs to construct the full Jacobian for these products. That can save substantial memory. Higher-order differentiation can similarly compute a Hessian-vector product without storing the full matrix of second derivatives.

Changing gradient flow changes what is differentiated

A stop-gradient operation keeps a value in the forward calculation but treats it as constant during differentiation. A custom derivative supplies a derivative rule explicitly. It should match the intended function unless the exercise asks for a surrogate, such as an approximation used during training.