Automatic Differentiation

The Computation Graph

Every differentiable program can be unrolled into a Wengert list (also called a tape): a directed acyclic graph (DAG) where nodes are intermediate values and edges represent elementary operations. Given a function $f: \mathbb{R}^n \to \mathbb{R}^m$, AD exploits the chain rule in one of two traversal orders.

Forward Mode (Tangent Linear) Forward mode propagates tangent vectors (directional derivatives) alongside primal values. For a single input perturbation $\dot{x} = \frac{\partial x}{\partial t}$, each intermediate variable $v_i$ carries a "dot" value $\dot{v}_i = \frac{\partial v_i}{\partial t}$.

Explanation: The tangent of each node is the sum of its parents' tangents weighted by local Jacobian factors — a single forward sweep computes one directional derivative of the output with respect to a chosen input direction. Forward mode requires $n$ sweeps to compute the full Jacobian $J \in \mathbb{R}^{m \times n}$, making it efficient when $n \ll m$ (few inputs, many outputs).

Reverse Mode (Adjoint / Backpropagation) Reverse mode propagates adjoint values (also called cotangents or bars) backward through the tape. Define the adjoint of node $v_i$ as the sensitivity of the scalar output $L$ to $v_i$:

Explanation: The adjoint $\bar{v}_i$ measures how much a small perturbation in intermediate value $v_i$ changes the final loss $L$ — it is the fundamental quantity backpropagation accumulates. The reverse sweep applies the chain rule to accumulate adjoints from outputs back to inputs:

Explanation: Each parent node $v_j$ receives a contribution to its adjoint from every child $v_i$, scaled by the local partial derivative — this accumulation over all children implements the multivariate chain rule exactly. A single reverse sweep computes $\nabla_x L \in \mathbb{R}^n$ — all $n$ parameter gradients simultaneously — at cost $O(\text{forward pass})$. This is the fundamental asymmetry that makes deep learning feasible.

Jacobian Cost Comparison For $f: \mathbb{R}^n \to \mathbb{R}^m$ with $W$ operations in the forward pass:

Explanation: When $m=1$ (scalar loss, as in training), reverse mode computes the entire gradient in $O(W)$ — independent of $n$ — while forward mode would require $n$ passes, one per parameter.

Operator Overloading in PyTorch

PyTorch implements reverse-mode AD via operator overloading: every `torch.Tensor` with `requires_grad=True` records operations into a dynamic tape. ```python import torch

# Define inputs x = torch.tensor([2.0, 3.0], requires_grad=True) W = torch.tensor([[1.0, 2.0], [3.0, 4.0]], requires_grad=True)

# Forward pass — tape is recorded automatically h = torch.tanh(W @ x) # shape: [2] L = h.sum() # scalar loss

# Reverse pass — adjoints flow backward through the tape L.backward() print(x.grad) # dL/dx — exact, not approximate print(W.grad) # dL/dW — computed in a single backward sweep ``` PyTorch's autograd engine stores a `grad_fn` on every non-leaf tensor, forming a linked list of `Function` objects. Each `Function` implements both `forward()` (primal computation) and `backward()` (local Jacobian-vector product). During `loss.backward()`, the engine performs a topological sort of this graph and calls `backward()` in reverse order, accumulating `.grad` on leaf tensors.

Custom Backward Passes You can register custom primitives with explicitly specified derivatives: ```python class MySigmoid(torch.autograd.Function): @staticmethod def forward(ctx, x): s = torch.sigmoid(x) ctx.save_for_backward(s) # save primal for use in backward return s @staticmethod def backward(ctx, grad_output): (s,) = ctx.saved_tensors

# dL/dx = dL/ds ds/dx = grad_output s(1-s) return grad_output s * (1 - s) sigmoid = MySigmoid.apply ```

Explanation: The custom backward multiplies the incoming adjoint $\bar{s}$ by the local derivative of sigmoid — a Jacobian-vector product — demonstrating the modular, chain-rule composition that all AD systems exploit.

Backpropagation · Computational Graphs · JAX and Functional Transformations · Higher-Order Derivatives · Differentiable Programming