Forward-Mode and Reverse-Mode Autodiff

Same Jacobian product, opposite evaluation order. Forward mode multiplies right to left and gives \(Jv\), one column of the Jacobian per pass, at forward-pass memory. Reverse mode multiplies left to right and gives \(u^\top J\), one row per pass, but must store intermediates until the backward sweep. Filling the whole Jacobian costs \(n\) forward passes or \(m\) reverse passes, where \(n\) is inputs and \(m\) outputs.

Each JVP yields the derivative along one input direction, so recovering the full gradient needs \(n\) passes, roughly \(10^{10}\) of them per step. Reverse mode needs one pass per output, and the loss is a single scalar, so \(m = 1\). The asymmetry between \(n \approx 10^{10}\) and \(m = 1\) is the whole reason backpropagation is the algorithm rather than a choice.

Reverse mode cannot begin until the forward pass ends, and each backward step needs intermediates from its forward counterpart, so activations stay alive across the whole forward pass. That is why activation memory dominates training memory rather than weights. Gradient checkpointing exists to trade it back: discard intermediates on the way forward and recompute them during the backward sweep, buying memory with roughly one extra forward pass.

Compose the two modes: apply forward mode to the gradient function, giving \(Hv\) as a JVP of a VJP. Cost is a small multiple of one gradient evaluation and memory never exceeds the \(n\)-vector, versus \(n^2\) for the explicit Hessian. Pearlmutter's 1994 trick underpins every matrix-free second-order method, including conjugate gradient on Gauss-Newton systems, Lanczos spectrum estimation, and loss-sharpness probes.

Reverse mode costs one backward pass per output row, so a 2,000-dimensional output means 2,000 passes. If the input dimension is smaller than 2,000, forward mode fills the same Jacobian in fewer passes. The habit of always using reverse mode comes from the scalar-loss case; once \(m\) is comparable to or larger than \(n\), the cost model inverts and the wall clock shows it.

Because autodiff differentiates the program you wrote, not the mathematics you meant. A naive exp then sum then log overflows and cancels catastrophically, and the derived gradient inherits every one of those problems. Writing the op as a single primitive lets the library apply the numerically stable formulation forward and a closed-form, stable gradient backward. Autodiff is exact with respect to your code, which is not the same as accurate.

Same Jacobian product, opposite evaluation order. Forward mode multiplies right to left and gives \(Jv\), one column of the Jacobian per pass, at forward-pass memory. Reverse mode multiplies left to right and gives \(u^\top J\), one row per pass, but must store intermediates until the backward sweep. Filling the whole Jacobian costs \(n\) forward passes or \(m\) reverse passes, where \(n\) is inputs and \(m\) outputs.

Each JVP yields the derivative along one input direction, so recovering the full gradient needs \(n\) passes, roughly \(10^{10}\) of them per step. Reverse mode needs one pass per output, and the loss is a single scalar, so \(m = 1\). The asymmetry between \(n \approx 10^{10}\) and \(m = 1\) is the whole reason backpropagation is the algorithm rather than a choice.

Reverse mode cannot begin until the forward pass ends, and each backward step needs intermediates from its forward counterpart, so activations stay alive across the whole forward pass. That is why activation memory dominates training memory rather than weights. Gradient checkpointing exists to trade it back: discard intermediates on the way forward and recompute them during the backward sweep, buying memory with roughly one extra forward pass.

Compose the two modes: apply forward mode to the gradient function, giving \(Hv\) as a JVP of a VJP. Cost is a small multiple of one gradient evaluation and memory never exceeds the \(n\)-vector, versus \(n^2\) for the explicit Hessian. Pearlmutter's 1994 trick underpins every matrix-free second-order method, including conjugate gradient on Gauss-Newton systems, Lanczos spectrum estimation, and loss-sharpness probes.

Reverse mode costs one backward pass per output row, so a 2,000-dimensional output means 2,000 passes. If the input dimension is smaller than 2,000, forward mode fills the same Jacobian in fewer passes. The habit of always using reverse mode comes from the scalar-loss case; once \(m\) is comparable to or larger than \(n\), the cost model inverts and the wall clock shows it.

Because autodiff differentiates the program you wrote, not the mathematics you meant. A naive exp then sum then log overflows and cancels catastrophically, and the derived gradient inherits every one of those problems. Writing the op as a single primitive lets the library apply the numerically stable formulation forward and a closed-form, stable gradient backward. Autodiff is exact with respect to your code, which is not the same as accurate.