Jacobian-Vector and Vector-Jacobian Products in Autodiff

A function from $n$ real inputs to $m$ real outputs has a Jacobian of $mn$ entries, and for any network of interest that number is far larger than the machine can hold: a layer mapping 4096 units to 4096 units has roughly $1.7 \times 10^{7}$ entries, about 67 MB in 32-bit floats for one layer at one point. Differentiation machinery is therefore built around two products that never form the matrix. Both cost a constant multiple of one evaluation of the function, and the difference between them is the direction in which the chain rule is traversed.

The two products and the cost rule that picks between them

For $f:\mathbb{R}^{n} \to \mathbb{R}^{m}$ with Jacobian $J$ at a point, a Jacobian-vector product takes a direction $\dot{x} \in \mathbb{R}^{n}$ and returns $J\dot{x} \in \mathbb{R}^{m}$. A vector-Jacobian product takes a cotangent $\bar{v} \in \mathbb{R}^{m}$ and returns $\bar{v}^{\top} J \in \mathbb{R}^{n}$. The first is one step of forward-mode differentiation, propagating a directional derivative alongside the values. The second is the reverse sweep of the previous stage, described in reverse-mode differentiation, seeded with $\bar{v}$.

Recovering the full Jacobian requires $n$ Jacobian-vector products or $m$ vector-Jacobian products, so the cheaper mode is the one whose number of sweeps matches the smaller dimension. A loss function has $m = 1$ and $n$ equal to the parameter count, so a single reverse sweep gives the whole gradient while forward mode would need one sweep per parameter: for $10^{7}$ parameters, a factor of $10^{7}$ in work. The asymmetric case is a function with many outputs and few inputs, such as rendering a large image from a handful of parameters, where forward mode is the cheaper direction. When $m = n$ the arithmetic cost is comparable and the decision turns on memory: forward mode can be run as a single pass that carries the derivatives of the values it has computed so far and discards them immediately, while reverse mode must keep the intermediates of the forward pass addressable until the sweep reaches them, so its memory grows with the depth and the activation size.

The interfaces expose this directly. Libraries that support both directions provide a pair of entry points for one-off products and a way to transform a function into a differentiable one; the same primitive rules underlie both, and the survey literature calls the design principle that each rule be written once and used in both directions the central engineering requirement of an autodiff system.

Local rules and their transposes

Every primitive needs both products, and the two are transposes of each other in the sense that the identity $\langle \bar{v}, J\dot{x}\rangle = \langle J^{\top}\bar{v}, \dot{x}\rangle$ holds for all directions and cotangents. For the common operations the rules are short enough to enumerate, with $x, z$ vectors or matrices of compatible shape:

Operation	Jacobian-vector product	Vector-Jacobian product
$y = Wx$	$\dot{y} = W\dot{x}$	$\bar{x} = W^{\top}\bar{v}$
$y = \phi(x)$ elementwise	$\dot{y} = \phi’(x) \odot \dot{x}$	$\bar{x} = \phi’(x) \odot \bar{v}$
$y = x \odot z$	$\dot{y} = \dot{x} \odot z + x \odot \dot{z}$	$\bar{x} = \bar{v} \odot z$, $\bar{z} = \bar{v} \odot x$
$y = \sum_i x_i$	$\dot{y} = \sum_i \dot{x}_i$	$\bar{x}_i = \bar{v}$ broadcast
$Y = XZ$	$\dot{Y} = \dot{X}Z + X\dot{Z}$	$\bar{X} = \bar{V}Z^{\top}$, $\bar{Z} = X^{\top}\bar{V}$

The transpose relation is more than a curiosity: it is a check. Any pair of user-supplied rules for a primitive can be validated by drawing random $\dot{x}$ and $\bar{v}$ and comparing the two inner products, because the equality is exact in real arithmetic and therefore holds to rounding for consistent rules. Two evaluations per random pair make this the cheapest available structural test, and it catches the transposition errors that the finite-difference check of gradient checking detects only in the aggregate direction it happens to sample.

import numpy as np

def f(x, W1, W2):
    h = np.tanh(x @ W1.T)
    return h @ W2.T

def jvp(x, W1, W2, dx, dW1, dW2):
    h = np.tanh(x @ W1.T)
    dh = (1.0 - h ** 2) * (dx @ W1.T + x @ dW1.T)
    return dh @ W2.T + h @ dW2.T

def vjp(x, W1, W2, v):
    h = np.tanh(x @ W1.T)
    gh = v @ W2                       # cotangent at the hidden layer
    gz = gh * (1.0 - h ** 2)          # through the elementwise tanh
    return gz @ W1, gz.T @ x, v.T @ h  # dL/dx, dL/dW1, dL/dW2

rng = np.random.default_rng(0)
x, W1, W2 = rng.normal(size=(1, 4)), rng.normal(size=(5, 4)), rng.normal(size=(3, 5))
dx, dW1, dW2, v = rng.normal(size=x.shape), rng.normal(size=W1.shape), rng.normal(size=W2.shape), rng.normal(size=(1, 3))

y_dot = jvp(x, W1, W2, dx, dW1, dW2)
gx, gW1, gW2 = vjp(x, W1, W2, v)
lhs = np.sum(v * y_dot)                                  # <v, J dx>
rhs = np.sum(gx * dx) + np.sum(gW1 * dW1) + np.sum(gW2 * dW2)   # <J^T v, dx>
assert np.allclose(lhs, rhs, atol=1e-12)                 # adjoint identity

Second-order information without a Hessian

The Hessian of a scalar loss has $p^{2}$ entries, so it is never formed for a real model; the same dimension argument that motivated the gradient products applies one order higher. A Hessian-vector product takes a direction $v$ and returns $Hv$, and the trick that produces it is to differentiate the gradient map itself. Applying a forward-mode product to the reverse sweep with direction $v$ yields exactly $Hv$: forward-over-reverse, with the reverse sweep at the outer level because the gradient is a vector-valued function of a large number of parameters and reverse mode is what makes the sweep cheap at that level. The reverse-over-reverse arrangement returns the same product but requires differentiating the backward pass, and therefore requires the intermediate quantities of the backward pass to be retained, which multiplies memory by a factor related to the number of layers. This is the result that Hessian-vector products are available at a small constant multiple of the cost of the gradient itself.

The practical consequence is that truncated second-order methods, conjugate-gradient inner loops, and curvature diagnostics that need only products are technically within reach of any system that implements both modes, without a Hessian and without an approximation of one. The same construction gives the product of the Fisher matrix by a vector, since that matrix is a covariance of gradients and each of its factors is a vector-Jacobian product of the log-probability. The cost is bounded by a small multiple of the gradient’s cost, not by the parameter count, which is the property that makes the whole family of iterative second-order methods applicable at all: an inner conjugate-gradient loop with $r$ iterations costs $r$ products, and $r$ is chosen for the problem rather than forced by the size of the matrix.

Many directions at once

A single direction is not the only option. The rules behind one product are maps over tensors, so a direction can be replaced by a batch of directions of the same shape, and the rule evaluates them together with an extra leading axis. In the table above this changes nothing structurally: the elementwise rule becomes a broadcast multiply and the matrix product gains a dimension. The arithmetic grows linearly in the number of directions, but the traffic over the parameter tensors does not, because the values being differentiated are read once for the whole batch of directions rather than once per direction. Past the point where the working set no longer fits in the fastest levels of memory, that amortization is worth more than the arithmetic, so a batched evaluation of 32 directions costs appreciably less than 32 separate evaluations and returns the same numbers.

This is also the practical way to obtain a small Jacobian in full. Directions chosen as the standard basis vectors yield the columns of $J$, so $n$ batched products with a one-hot direction each reproduce the matrix; for a layer with a hundred outputs and a hundred inputs that is $10^{4}$ entries and a few megabytes, cheap enough to inspect directly. For a quantity such as the derivative of a class score with respect to an input image, the reverse map is more convenient: one product per class with a one-hot cotangent gives the gradient for that class, and $K$ products cover all $K$ classes, which is the standard construction behind gradient-based attribution methods. The number of products needed is set by the dimension that is being enumerated, and the choice of mode decides which dimension that is. Batching directions costs memory in proportion to their number, since each one needs its own intermediate value wherever the rule that consumes it needs to be evaluated, so the number of directions computed at once is bounded by the same budget that bounds the batch of examples.

What is never formed, and what that costs

The two products make the correct direction cheap and hide the structure that the full Jacobian would reveal. A diagonal, a spectrum, or a rank estimate requires either $n$ or $m$ sweeps, or an iterative method that uses the products as an oracle. This is not an inefficiency in the machinery but a statement about what is being asked for: the products answer questions about a single direction of response, and everything that requires the shape of the response as a whole costs proportionally more. The same limitation applies to the memory of the reverse sweep, which grows with the number of recorded intermediates rather than with the parameter count, and is therefore bounded by batch size and depth rather than by model size.