Automatic differentiation over frames

Automatic differentiation computes the derivatives of a program exactly, up to rounding, from the derivatives of the steps it is made of. This note states the three ways of composing them that a library needs (pushing tangents forward, pulling cotangents back, and differentiating the pull-back once more for second order), each over a frame of several directions at once.

Programs and their computation graphs#

A straight-line program has slots $s_0, s_1, \dots$, each holding a vector, and runs a sequence of steps

each writing one slot from slots written earlier. Each step’s map $\map_k$ is smooth, with partial Jacobians $\partial_j \map_k$, one per input, of shape (output dimension) by (dimension of input $j$). When an input is not differentiated, its partial Jacobian is simply never used.

A program is pictured by its computation graph. Its nodes are the slots, and it has an edge from each input of a step to that step’s output, so a node with incoming edges is computed by the step that writes it, from the nodes those edges come from. Because each step reads only slots written before it, the computation graph has no cycles, and the order of the steps is one of its topological orders. Its roots, the nodes with no incoming edge and so no ancestor, are the slots no step writes: the program’s inputs, the variables $\point$ we differentiate with respect to. Some nodes are targets, the values to differentiate, $f(\point)$. A node with several outgoing edges is a slot read by several steps.

$\point_1$ $\point_2$ $s_3$ $\map_1 = {\odot}$$s_4$ $\map_2 = \exp$$f$ $\map_3 = {+}$

• Root (an input)
• Slot written by a step
• input of a step

A frame is a matrix whose columns are directions: a tangent frame at a slot of dimension $\dimension$ is a $\dimension \times n$ matrix $\tangent$ of $n$ tangent vectors, and a cotangent frame a $\dimension \times n$ matrix $\alpha$ whose columns are linear functionals, written as columns. Below, $\tangent_s$ and $\alpha_s$ are the frames held at slot $s$.

Forward: pushing tangent frames#

Seed each root with a tangent frame $\tangent_{\point}$ (missing means zero), and run the steps in order, each pushing its inputs’ frames forward along the computation graph’s edges by its Jacobian–vector product:

By the chain rule, the frame reaching a target is $\sum_l \partial f / \partial \point_l \cdot \tangent_{\point_l}$: one pass with a frame of $n$ columns gives $n$ directional derivatives, and seeding the columns of an identity gives the whole Jacobian.

Reverse: pulling cotangent frames over the backpropagation graph#

The reverse pass runs on the backpropagation graph: the same nodes as the computation graph, with every edge reversed and labelled by the pull-back $\partial_j \map_k\transpose$ of the step it came from. Since reversing the edges swaps incoming for outgoing, a node is a root of the computation graph if and only if it is a leaf of the backpropagation graph, a node with no outgoing edge. Cotangents enter at the targets, flow along its edges, and end at those leaves, the program’s inputs. A node with several incoming edges, a slot read by several steps, receives a cotangent along each, and holds their sum.

$\partial_1 \map_1\transpose$$\partial_2 \map_1\transpose$$\partial \map_2\transpose$$\partial_1 \map_3\transpose$$\partial_2 \map_3\transpose$ $\point_1$ $\point_2$ $s_3$ $\map_1 = {\odot}$$s_4$ $\map_2 = \exp$$f$ $\map_3 = {+}$

• Root (an input)
• Slot written by a step
• pull-back

Seed each target with a cotangent frame $\alpha_f$, every other slot with zero, and run the steps backwards, each adding its vector–Jacobian product, the pull-back of its output’s cotangent, into each input:

The accumulation, which sums the cotangents arriving along every incoming edge of the backpropagation graph, is the chain rule for a slot with several consumers. When the steps are done, each root holds $\alpha_{\point_l} = \sum_f (\partial f / \partial \point_l)\transpose \alpha_f$. For a scalar target seeded with $\alpha_f = 1$, that is the gradient $\nabla_{\point_l} f$, whatever the number of roots.

Each column of a frame costs a small multiple of one evaluation of the program, in either direction. A Jacobian with few rows is cheaper in reverse, one with few columns forward.

Forward-over-reverse: Hessian products#

For second derivatives, differentiate the reverse pass itself along a path $\point + t\tangent$. Run the forward pass of (2) with velocities $\tangent$ seeded at the roots, so that every slot holds its value $s$ and its velocity $\tangent_s$, then run the reverse pass over the backpropagation graph carrying two cotangents per slot: $\alpha_s$, as in (3), and $\beta_s$, its derivative along the path. One column of $\tangent$ is one path; a frame of $n$ columns runs $n$ paths at once, and then $\beta_s$ is a frame of $n$ columns while $\alpha_s$ is the same for all of them. (A cotangent frame of several columns together with a velocity frame of several columns would make $\beta$ a three-way array; the passes here take one cotangent column.)

Differentiating (3) along the path by the product rule gives

The first term is the step’s own second-order map, which depends on the velocities of its inputs; the second is the step’s ordinary pull-back, applied to $\beta$. A library writes the first for each step and the program applies the rule. Two consequences are useful when writing $\hvp$ for a step. If a step is linear, its partial Jacobians are constant and its $\hvp$ is zero. If the pull-back into input $j$ does not involve input $l$, then the velocity of input $l$ contributes nothing to $\hvp_{k,j}$.

With all roots stacked into one vector, $\beta_\point = \nabla^2 f(\point)\,\tangent$: a frame of $n$ velocities gives $n$ columns of the Hessian for about $n$ times the cost of a gradient, and the identity frame gives all of it.