Reverse-mode automatic differentiation

Deep dive

How it works

In: a recorded computation graph ending in one scalar loss · Out: the gradient of the loss with respect to every input and parameter · State: the tape (parent indices plus saved tensors or partials) and one adjoint per node during the sweep · Cost: about 2-4x one forward pass, independent of the number of inputs; memory O(saved intermediates)

Forward (record). Evaluate the function. For each primitive y = op(x1, x2), append a node holding parent indices plus whatever the backward rule needs: local partials for scalars, saved inputs for tensors.
Seed. Set adj[loss] = 1. All other adjoints start at 0.
Reverse sweep. Visit nodes newest to oldest. For each node, adj[parent] += ∂node/∂parent · adj[node]. Tape order is already a reverse topological order, so no graph sort is needed.
Read out. adj[leaf] is ∂loss/∂leaf.

The += is essential. A value used k times (fan-out) collects k contributions.

x ─┬─► mul ─► add ─► z        forward: left to right, store partials
y ─┘    ▲      ▲
x ─► sq ─► sin ┘              backward: right to left, adj flows along the same edges

Tensor VJP rules (G = upstream gradient):

op	saved	backward
`Y = X W`	X, W	`dX = G Wᵀ`, `dW = Xᵀ G`
`Y = relu(X)`	mask	`dX = G ⊙ (X > 0)`
`y = softmax(x)`	y	`dx = y ⊙ (G − ⟨G, y⟩)`
`L = CE(logits, t)`	softmax p	`dlogits = (p − onehot(t)) / batch`
`Y = X + b` (broadcast)	none	`db = sum of G over broadcast dims`

Code

A scalar Wengert-list tape. Production engines keep the same structure, but each node is a tensor op with a VJP closure.

use std::cell::RefCell;
use std::ops::{Add, Mul, Sub};

#[derive(Clone, Copy)]
struct Node { parents: [usize; 2], partials: [f64; 2] } // d(node)/d(parent)

#[derive(Default)]
pub struct Tape { nodes: RefCell<Vec<Node>> }

#[derive(Clone, Copy)]
pub struct Var<'t> { tape: &'t Tape, idx: usize, pub val: f64 }

impl Tape {
    pub fn var(&self, val: f64) -> Var<'_> { self.push([0, 0], [0.0, 0.0], val) } // leaf: zero partials
    fn push(&self, parents: [usize; 2], partials: [f64; 2], val: f64) -> Var<'_> {
        let mut n = self.nodes.borrow_mut();
        n.push(Node { parents, partials });
        Var { tape: self, idx: n.len() - 1, val }
    }
}

impl<'t> Var<'t> {
    fn unary(self, val: f64, d: f64) -> Var<'t> { self.tape.push([self.idx, self.idx], [d, 0.0], val) }
    pub fn sin(self) -> Var<'t> { self.unary(self.val.sin(), self.val.cos()) }
    pub fn exp(self) -> Var<'t> { let e = self.val.exp(); self.unary(e, e) }
    pub fn ln(self) -> Var<'t> { self.unary(self.val.ln(), 1.0 / self.val) }
    pub fn relu(self) -> Var<'t> { self.unary(self.val.max(0.0), (self.val > 0.0) as u8 as f64) }

    /// One reverse sweep: adjoints of every node w.r.t. self.
    pub fn grad(self) -> Vec<f64> {
        let nodes = self.tape.nodes.borrow();
        let mut adj = vec![0.0; nodes.len()];
        adj[self.idx] = 1.0;
        for i in (0..=self.idx).rev() {
            let (n, a) = (nodes[i], adj[i]);
            if a == 0.0 { continue; }
            adj[n.parents[0]] += n.partials[0] * a; // += : fan-out accumulates
            adj[n.parents[1]] += n.partials[1] * a;
        }
        adj
    }
}

impl<'t> Add for Var<'t> { type Output = Self;
    fn add(self, o: Self) -> Self { self.tape.push([self.idx, o.idx], [1.0, 1.0], self.val + o.val) } }
impl<'t> Sub for Var<'t> { type Output = Self;
    fn sub(self, o: Self) -> Self { self.tape.push([self.idx, o.idx], [1.0, -1.0], self.val - o.val) } }
impl<'t> Mul for Var<'t> { type Output = Self;
    fn mul(self, o: Self) -> Self { self.tape.push([self.idx, o.idx], [o.val, self.val], self.val * o.val) } }

fn main() {
    let t = Tape::default();
    let (x, y) = (t.var(0.5), t.var(4.0));
    let z = x * y + (x * x).sin() - y.ln();           // dz/dx = y + 2x cos(x²), dz/dy = x - 1/y
    let g = z.grad();
    assert!((g[x.idx] - (4.0 + 1.0 * 0.25f64.cos())).abs() < 1e-12);
    assert!((g[y.idx] - (0.5 - 0.25)).abs() < 1e-12);

    // Finite-difference check: central difference in f64, h around 1e-6.
    let f = |x: f64, y: f64| x * y + (x * x).sin() - y.ln();
    let h = 1e-6;
    let fd = (f(0.5 + h, 4.0) - f(0.5 - h, 4.0)) / (2.0 * h);
    assert!((fd - g[x.idx]).abs() < 1e-8);
    println!("dz/dx = {:.6}, dz/dy = {:.6}", g[x.idx], g[y.idx]);
}

Complexity and constants

Time: one forward plus one backward. For a GEMM, the backward does two GEMMs of the same size (dX and dW), so a training step costs about 3x forward: about 6·N FLOPs per token for an N-parameter dense model.
Memory: saved activations dominate. A transformer layer in 16-bit precision saves on the order of 34·s·b·h bytes, plus 5·a·s²·b for unfused attention (Korthikanti et al.). Weights are small by comparison at long sequence lengths.
A scalar tape like the one above costs about 32 bytes and one indirect jump per op. It runs 10-100x slower than tensor-level autodiff, which amortizes the bookkeeping over whole kernels.

Efficiency

Cost: a dense net's backward does two GEMMs per forward GEMM, so a step is about 3x forward, 6·N FLOPs per token for N parameters. Memory is the bound: a transformer layer in 16-bit saves order of 34·s·b·h bytes of activations (Korthikanti et al.). A scalar tape costs about 32 bytes and one indirect jump per op and runs 10-100x slower than tensor-level autodiff (order of, above).
Levers:
- recompute instead of store (checkpointing) — selective recomputation cuts activation memory about 5x and over 90% of the recompute time versus full recompute (Korthikanti et al. 2022) — memory-bound training
- save the cheapest thing (1-bit ReLU mask, softmax output, dropout seed) — fewer bytes held per op — every custom backward rule
- fused loss and gradient (softmax + cross-entropy gives p − onehot) — no N×V intermediate — large vocabularies
- trace and compile the backward (JAX/XLA, torch.compile) — fused kernels, planned memory — static shapes
- free saved tensors as the sweep passes them — peak memory does not double — define-by-run engines
Naive → standard → state of the art:
- one forward pass per input (finite differences) · cost grows with the number of inputs
- tape of tensor ops with VJP rules · about 3x forward
- compiled graph with fusion and selective recomputation · about 3x forward compute with a fraction of the activation memory

Implementation notes

Define-by-run (PyTorch, this tape): build the graph while executing. Python control flow works, and there is a per-op dispatch cost of a few µs. Trace-then-transform (JAX grad = linearize + transpose, XLA): differentiates a staged program, so the compiler can fuse and plan memory.
PyTorch reference-counts each grad_fn node and frees saved tensors as soon as the backward sweep passes them. Keep that property: release saved buffers in reverse order, or peak memory doubles.
Save the cheapest thing that suffices. For ReLU, keep a 1-bit mask, not the input. For softmax, keep the output, not the logits. For dropout, keep the RNG seed and offset, not the mask.
Fuse loss and its gradient (softmax + cross-entropy gives p − onehot). This is numerically stable and avoids the N×V intermediate.
Batch the gradients of parameters reused across time steps (RNNs, tied embeddings). The adjoint accumulation is the expensive memory write.

Pitfalls

Overwriting instead of accumulating adjoints gives wrong gradients for any reused value. Tied weights are the classic case.
In-place ops that clobber a saved tensor. PyTorch detects this with version counters ("one of the variables needed for gradient computation has been modified by an inplace operation").
Keeping the graph alive: total += loss (not loss.item()) retains every step's graph and leaks memory.
Non-smooth points: relu'(0), max ties and abs(0) use a subgradient convention. sqrt'(0) = ∞ produces NaN; add an epsilon inside sqrt for norms.
Gradient checks in f32 are noisy. Use f64, central differences, h ≈ 1e-6, and relative error < 1e-6 (torch.autograd.gradcheck does this).

Variants and alternatives

Forward mode / dual numbersone pass per input direction (JVP). Best for few inputs, or Hessian-vector products as forward-over-reverse.

Source transformationgenerate adjoint code ahead of time. Tapenade, and Enzyme at the LLVM IR level after optimization, so it differentiates C/C++/Rust/Julia.

Checkpointed reverse modestore some activations, recompute the rest (gradient checkpointing).

Implicit differentiationfor fixed points or optimizer layers, differentiate the optimality condition instead of unrolling.

In production

PyTorch autograd, JAX (jax.vjp, jax.grad), TensorFlow GradientTape, Enzyme (differentiates compiled C/C++/Rust/Julia numerical code), Stan's math library (reverse mode for Bayesian statistics), and the Rust frameworks Burn and candle.

Implementations

name	lang	repo	install	verdict	licence	health	as_of	notes
Burn	Rust	tracel-ai/burn	`cargo add burn`	—	—	—	—	autodiff decorator backend over any tensor backend
candle	Rust	huggingface/candle	`cargo add candle-core`	—	—	—	—	tensors with `backward()`; minimal ML framework
PyTorch autograd	Python / C++	pytorch/pytorch	`pip install torch`	—	—	—	—	define-by-run tape, `grad_fn` nodes
JAX	Python	jax-ml/jax	`pip install jax`	—	—	—	—	`jax.grad`/`jax.vjp` over traced programs, XLA
Enzyme	C++ / LLVM	EnzymeAD/Enzyme	`brew install enzyme`	—	—	—	—	reverse mode on LLVM IR; C, C++, Rust, Julia

Sources

Griewank and Walther, Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation, 2nd ed., SIAM, 2008. book
Baydin, Pearlmutter, Radul, Siskind, Automatic Differentiation in Machine Learning: a Survey, JMLR 2018. paper · arXiv
Paszke et al., Automatic differentiation in PyTorch, NIPS Autodiff Workshop, 2017. paper · code
Moses and Churavy, Instead of Rewriting Foreign Code for Machine Learning, Automatically Synthesize Fast Gradients (Enzyme), NeurIPS 2020. paper · code
Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models, 2022. paper
Carpenter et al., The Stan Math Library: Reverse-Mode Automatic Differentiation in C++, 2015. paper · code