Chapter 183Optional

Reverse-Mode Automatic Differentiation and Cotangent Semantics

Prerequisites. Direct starred prerequisites: Chapter 182. No later core chapter depends on this route.

One evaluation of $D t$ delivers one directional derivative, by theorem 182.12. For a program $t$ of type $real$ in a context of $n$ variables of type $real$ , the gradient has $n$ components, so obtaining it by the tangent macro costs $n$ evaluations of $D t$ , hence about $n$ times the cost of evaluating $t$ . When $t$ is a loss function with $n$ in the millions, that factor decides whether the program is usable.

The gradient can be computed in a constant multiple of the cost of $t$ , and the reason is linear algebra rather than programming. Evaluating $t$ at a point composes Jacobians, $J = J_{k} \dots J_{1}$ ; the tangent macro multiplies this product by a column vector from the right, one column at a time, while the gradient is $J^{t}$ applied to a single output cotangent, and $J^{t} = J_{1}^{t} \dots J_{k}^{t}$ is a product of transposes in reverse order. So a transformation computing gradients must run the program forward to obtain the points at which the $J_{i}$ are taken, and then apply the transposes backward.

Nothing in chapter 182 does this, and no lemma of that chapter is reused here. What must be built is a translation producing, along with the value of $t$ , a backpropagator: a linear map from cotangents at the output to cotangents at the inputs. This chapter defines that translation, proves that its backpropagator denotes the transposed derivative, imports the corresponding statement for the language with higher-order subterms, and then treats the cost question separately, because the obvious representation of backpropagators is asymptotically wrong and the repair is not a corollary of correctness.

Signature delta

The forward development and the reverse development are two source systems, and the second is not an extension of the first. Before anything is reused, the differences are displayed.

	Forward chapter	This chapter
Source types	$real$ , finite products, functions	$real$ , finite products; functions only inside a term whose interface is first order
Target language	the same language	the source language extended by a type of linear maps between cotangent spaces
Translation	one macro $D$ on types and terms	one macro $\overset{\leftarrow}{D}$ producing a pair: a primal value and a backpropagator
Semantics	sets, with a logical relation on smooth curves	the same for the first-order fragment proved here; diffeological spaces for the imported theorem
Theorem proved here	theorem 182.12: tangent representation	theorem 183.5: transposed derivative
Theorem imported	theorem 182.14 (Huot–Staton–Vákár)	theorem 183.7 (Vákár), theorem 183.9 (Smeding–Vákár)

No lemma of chapter 182 is applied below. The denotational conventions of definition 182.5 — $real$ denotes $R$ , products denote products, each operation symbol denotes a fixed smooth function — are restated because they are the same conventions, not because a result is being transported.

Cotangents and backpropagators

Definition 183.1 — The first-order source fragment ¶

Types are $τ, σ ::= real ∣ (τ_{1} * \dots * τ_{n})$ and terms are $\begin{aligned} e ::= & x ∣ op (e_{1}, \dots, e_{k}) ∣ ⟨ e_{1}, \dots, e_{n} ⟩ \\ ∣ case e of ⟨ x_{1}, \dots, x_{n} ⟩ \to e^{'} ∣ let x = e in e^{'}, \end{aligned}$ typed as in definition 182.4 and denoted as in definition 182.5; every operation symbol denotes a smooth function and comes with terms $\partial_{i} op$ denoting its partial derivatives.

Referenced from 4 locations

Reverse-Mode Automatic Differentiation and Cotangent Semantics

Definition 183.2 — Cotangent spaces ¶

For each type put $C (real) : = R$ and $C (τ_{1} * \dots * τ_{n}) : = C (τ_{1}) \times \dots \times C (τ_{n})$ , and for a context $Γ = x_{1} : τ_{1}, \dots, x_{n} : τ_{n}$ put $C (Γ) : = \prod_{i} C (τ_{i})$ . Each $C (τ)$ is a real vector space with componentwise addition and a zero element $0_{τ}$ ; for $Γ$ and an index $i$ , write $ι_{i} : C (τ_{i}) \to C (Γ)$ for the map placing its argument in slot $i$ and $0$ elsewhere.

Referenced from 2 locations

Reverse-Mode Automatic Differentiation and Cotangent Semantics

Definition 183.3 — Backpropagator ¶

Let $Γ ⊢ e : τ$ and $γ \in [[Γ]]$ . A backpropagator for $e$ at $γ$ is a linear map $b : C (τ) \to C (Γ)$ such that for every $d \in C (τ)$ and every index $i$ , $(b (d))_{i} = \sum_{c} d_{c} \frac{\partial [[e {]]}_{c}}{\partial x_{i}} (γ),$ the sum running over the scalar components $c$ of $τ$ . Equivalently $b = (D [[e]] (γ))^{t}$ , the transpose of the Jacobian of $[[e]]$ at $γ$ .

Referenced from 2 locations

Reverse-Mode Automatic Differentiation and Cotangent Semantics

Definition 183.4 — The reverse macro ¶

For $Γ ⊢ e : τ$ define a pair $\overset{\leftarrow}{D} (e) = (v_{e}, b_{e})$ , where $v_{e}$ is a term of type $τ$ in context $Γ$ and $b_{e}$ is a term denoting a map $C (τ) \to C (Γ)$ , by induction on the derivation: $\begin{aligned} \overset{\leftarrow}{D} (x_{i}) & : = (x_{i}, λ d . ι_{i} (d)), \\ \overset{\leftarrow}{D} (op (e_{1}, \dots, e_{k})) & : = (op (\vec{v}), λ d . \sum_{j = 1}^{k} b_{j} (\partial_{j} op (\vec{v}) * d)), \\ \overset{\leftarrow}{D} (⟨ e_{1}, \dots, e_{n} ⟩) & : = (⟨ \vec{v} ⟩, λ ⟨ \vec{d} ⟩ . \sum_{j = 1}^{n} b_{j} (d_{j})), \\ \overset{\leftarrow}{D} (let x = e_{1} in e_{2}) & : = (v_{2}, λ d . g + b_{1} (s)), where ⟨ g, s ⟩ = b_{2} (d), \end{aligned}$ writing $\vec{v}$ for $v_{1}, \dots$ and $\vec{d}$ for $d_{1}, \dots$ , and in the last clause $b_{2}$ is the backpropagator term produced for $e_{2}$ in the context $Γ, x : σ$ , so that it returns a pair consisting of a $C (Γ)$ component $g$ and a $C (σ)$ component $s$ . The clause for $case$ is the clause for $let$ with $n$ bound variables in place of one.

Referenced from 11 locations

Reverse-Mode Automatic Differentiation and Cotangent Semantics

Each clause is the transpose of the corresponding forward clause, and the reversal of order is visible: in $let$ , the backpropagator of the body runs first and its output cotangent for $x$ is then fed to the backpropagator of the bound expression.

Theorem 183.5 — Correctness of the reverse macro ¶

Let $Γ ⊢ e : τ$ in the fragment of definition 183.1 and let $γ \in [[Γ]]$ . Then $[[v_{e}]] (γ) = [[e]] (γ)$ and $[[b_{e}]] (γ)$ is a backpropagator for $e$ at $γ$ ; that is, $[[b_{e}]] (γ) = (D [[e]] (γ))^{t}$ .

Referenced from 13 locations

Reverse-Mode Automatic Differentiation and Cotangent Semantics

Proof of Theorem 183.5 — Correctness of the reverse macro

Proof. Induction on the derivation. In each case the primal claim is immediate from the shape of the clause, so only the backpropagator claim is written.

Variable. $[[e]] = π_{i}$ has Jacobian the $i$ -th projection, whose transpose is $ι_{i}$ , which is the clause.

Operation. Write $u : = ([[e_{1}]] (γ), \dots, [[e_{k}]] (γ))$ . The chain rule gives, for each input index $i$ , $\frac{\partial [[op (e_{1}, \dots, e_{k})]]}{\partial x_{i}} (γ) = \sum_{j = 1}^{k} \frac{\partial [[op]]}{\partial y_{j}} (u) \cdot \frac{\partial [[e_{j}]]}{\partial x_{i}} (γ) .$ Multiplying by $d$ and summing over the single output component, equation 183.1 for the composite equals $\sum_{j} (b_{j} (\partial_{j} [[op]] (u) \cdot d))_{i}$ by the induction hypothesis applied to each $e_{j}$ with the cotangent $\partial_{j} [[op]] (u) \cdot d$ . That is the displayed clause.

Tuple. The Jacobian of $⟨ e_{1}, \dots, e_{n} ⟩$ stacks the Jacobians of the $e_{j}$ , so its transpose is the sum of their transposes applied to the corresponding components of the cotangent, which is the clause.

Let. Let $Γ ⊢ e_{1} : σ$ and $Γ, x : σ ⊢ e_{2} : τ$ , and abbreviate $γ^{+} : = (γ, [[e_{1}]] (γ))$ . The denotation of the $let$ is $[[e_{2}]] \circ ⟨ id, [[e_{1}]] ⟩$ , whose Jacobian at $γ$ is $J = J_{2}^{Γ} (γ^{+}) + J_{2}^{x} (γ^{+}) \cdot J_{1} (γ),$ where $J_{2}^{Γ}$ and $J_{2}^{x}$ are the blocks of the Jacobian of $[[e_{2}]]$ corresponding to the variables of $Γ$ and to $x$ . Transposing, $J^{t} (d) = (J_{2}^{Γ})^{t} (d) + J_{1}^{t} ((J_{2}^{x})^{t} (d))$ , and by the induction hypothesis for $e_{2}$ the pair $((J_{2}^{Γ})^{t} (d), (J_{2}^{x})^{t} (d))$ is exactly the pair $⟨ g, s ⟩$ computed by $b_{2} (d)$ ; the induction hypothesis for $e_{1}$ identifies $b_{1} (s)$ with $J_{1}^{t} (s)$ . The clause adds the two, which is $J^{t} (d)$ .

Linearity of each constructed map is immediate: $ι_{i}$ is linear, and the clauses build linear maps from linear maps by composition, addition, and scalar multiplication by values that do not depend on $d$ . ◻

Example 183.6 — One scalar output, counted ¶

Let $t_{n} : = let y_{2} = x_{1} * x_{2} in let y_{3} = y_{2} * x_{3} in \dots let y_{n} = y_{n - 1} * x_{n} in y_{n}$ , whose denotation is $x_{1} x_{2} \dots x_{n}$ . Evaluating $t_{n}$ costs $n - 1$ multiplications.

Forward mode: by theorem 182.12 one evaluation of $D t_{n}$ with seed $e_{i}$ returns $\partial / \partial x_{i}$ , and each evaluation costs $3 (n - 1)$ multiplications and $n - 1$ additions, by the operation clause of definition 182.6 at $*$ . The full gradient therefore costs $n$ such evaluations: about $3 n^{2}$ multiplications.

Reverse mode: one evaluation of $v_{t_{n}}$ costs $n - 1$ multiplications, and one run of $b_{t_{n}}$ at $d = 1$ performs, at each $let$ , two scalar multiplications and one addition of context cotangents. Counting only scalar arithmetic, the gradient costs about $3 n$ multiplications — one run, not $n$ . The asymptotic gain is a factor $n$ , and it is already visible in the count.

Referenced from 3 locations

Reverse-Mode Automatic Differentiation and Cotangent Semantics

Exercise 183.1 ¶

★☆☆ For $e : = ⟨ x_{1} * x_{2}, x_{1} + x_{2} ⟩$ compute $v_{e}$ and $b_{e}$ from definition 183.4, write the Jacobian of $[[e]]$ as a $2 \times 2$ matrix, and check equation 183.1 on the cotangents $(1, 0)$ and $(0, 1)$ .

Referenced from 2 locations

Reverse-Mode Automatic Differentiation and Cotangent Semantics

Exercise 183.2 ¶

★★☆ Exchange the two summands in the $let$ clause of definition 183.4, so that $b_{1}$ is applied to the $Γ$ component and $b_{2}$ ’s $Γ$ component is discarded. Exhibit a term on which the resulting map is not the transposed Jacobian, and identify the step of the proof of theorem 183.5 that fails.

Referenced from 2 locations

Reverse-Mode Automatic Differentiation and Cotangent Semantics

Higher-order subterms

Theorem 183.5 covers a fragment with no function type at all. The transformation extends to a source language with function types, at the price of a target language with a type of linear maps and a semantics in which higher types are interpreted; and the correctness statement then holds at exactly the same interface as in the forward chapter.

Theorem 183.7 — Correctness at first-order interfaces; imported ¶

Let the source language have ${real}^{n}$ types, finite products, and function types, and let the target language extend it with a type of linear maps. Write $\vec{D} (t) = (\vec{D} (t)_{1}, \vec{D} (t)_{2})$ and $\overset{\leftarrow}{D} (t) = (\overset{\leftarrow}{D} (t)_{1}, \overset{\leftarrow}{D} (t)_{2})$ for the forward and reverse macros. For every program $x : τ ⊢ t : σ$ between first-order types $τ$ and $σ$ , whose subterms may have higher types, $[[\vec{D} (t)_{1}]] = [[t]], [[\vec{D} (t)_{2}]] = D [[t]], [[\overset{\leftarrow}{D} (t)_{1}]] = [[t]], [[\overset{\leftarrow}{D} (t)_{2}]] = (D [[t]])^{t},$ where $D$ is the usual derivative and $(-)^{t}$ the transpose.

Referenced from 5 locations

Reverse-Mode Automatic Differentiation and Cotangent Semantics

This is Theorem 2 of Vákár, Reverse AD at Higher Types: Pure, Principled and Denotationally Correct, arXiv:2007.05283, proved there by a logical-relations argument over a subscone of diffeological spaces, the semantics being needed because Cartesian closure is unavailable in the categories of manifolds. The restriction to first-order $τ$ and $σ$ is part of the statement, not an artifact of the proof method: at higher types the transpose of a derivative is not determined by the source term, exactly as in remark 182.13. That paper conjectures, and does not prove, extensions to variant and inductive types and to general recursion; none of those is claimed here.

The cost of a backpropagator

Theorem 183.5 says what $b_{e}$ computes and says nothing about what it costs. The gap is real: the representation of backpropagators used in definition 183.4 is asymptotically wrong.

Example 183.8 — Where the naive representation loses ¶

In definition 183.4 every backpropagator returns a full element of $C (Γ)$ , a tuple with one slot per variable in scope, and the $let$ clause adds two such tuples. Consider the family $s_{m} : = let y_{1} = x_{1} * x_{1} in let y_{2} = y_{1} * y_{1} in \dots let y_{m} = y_{m - 1} * y_{m - 1} in y_{m}$ in a context of $n$ variables $x_{1}, \dots, x_{n}$ , of which only $x_{1}$ is used. Evaluating $s_{m}$ costs $m$ multiplications. Running $b_{s_{m}}$ performs $m$ additions of tuples of length $n$ , hence $Θ (m n)$ scalar operations, while the gradient has only one nonzero entry. The output is correct by theorem 183.5; the cost is $n$ times what the statement of the problem allows.

Referenced from 5 locations

Reverse-Mode Automatic Differentiation and Cotangent Semantics

The repair is a change of representation, not of the equations: cotangent tuples are accumulated destructively into one mutable environment, so that a variable’s contribution is a single update rather than an addition of two full tuples, and the zero tuple is created once. With that representation the transformed program admits a cost bound that is linear in the cost of the original.

Theorem 183.9 — Asymptotic efficiency; imported ¶

Let $Γ ⊢ t : τ$ be a program of the first-order source language, with $n$ variables in $Γ$ , the $i$ -th of type $σ_{i}$ . There are constants $c, c^{'}, c^{″} > 0$ , independent of $t$ , such that for all inputs and every output cotangent $d$ , $cost (run (snd D_{Γ} [t] d) ⟨ 0_{σ_{1}}, \dots, 0_{σ_{n}} ⟩) \leq c^{'} + c \cdot cost (t) + c^{″} \cdot n + c_{φ} \cdot size (d),$ where both costs are taken in one fixed call-by-value cost model, $run$ executes the cotangent-accumulation monad starting from the zero environment, and $size (d)$ is the size of the incoming cotangent.

Referenced from 8 locations

Reverse-Mode Automatic Differentiation and Cotangent Semantics

This is Equation (7) of Smeding and Vákár [SV24], derived there from the amortised criterion of their Equation (6) by initialising the accumulator with zeros; the cost model is the paper’s explicitly displayed call-by-value model, and the statement is machine-checked in the accompanying Agda development [SV23]. Three things are not imported with it. The bound is stated for first-order programs, and the paper does not claim it for arbitrary higher-order closure conversion. It says nothing about floating-point stability, which is a numerical property of the primitives, not of the transformation. And it says nothing about checkpointing, which trades memory for recomputation and changes the constant $c$ as well as the space usage; that is a separate design with its own analysis.

Remark 183.10 — Three distinct statements ¶

Derivative correctness (theorem 183.5, theorem 183.7), asymptotic efficiency (theorem 183.9), and numerical accuracy are three claims with three different proofs and three different failure modes. A transformation can be correct and quadratic, as in example 183.8; it can be fast and wrong, as any mutation of the $let$ clause shows; and it can be correct and fast while accumulating catastrophic rounding error, which neither theorem constrains, since both are stated over $R$ .

Referenced from 2 locations

Reverse-Mode Automatic Differentiation and Cotangent Semantics

Linearize, unzip, transpose. A different arrangement of the same material first linearizes a program to a tangent program, then unzips the result into a primal pass and a linear pass, and finally transposes the linear pass. The transposition step is where the reversal happens, and the design keeps it syntactically separate, at the price of a substructurally linear intermediate language in which a linear variable is used exactly once. This is a comparison of designs only: no theorem of that development is imported here, and theorem 183.5 is proved above without any linear intermediate language.

Array implementations. Reverse-mode implementations for parallel array languages report speedups on benchmarks. Benchmarks are execution evidence about particular programs and machines. They do not establish theorem 183.9, whose content is a bound in a displayed cost model, and they are not evidence about the transformation defined in definition 183.4, which is a different program transformation.

Suggested first pass.

None of these problems is a prerequisite for a later chapter. Begin with exercise 183.3, then exercise 183.4; the implementation project exercise 183.6 may be attempted at any time.

Exercise 183.3 ¶

★★☆ Prove that for composable first-order programs $e_{1}$ and $e_{2}$ the backpropagator of $let x = e_{1} in e_{2}$ constructed by definition 183.4 equals the composite $J_{1}^{t} \circ J_{2}^{t}$ when $e_{2}$ does not mention the variables of $Γ$ other than $x$ , and identify the extra summand that appears when it does. Relate the two cases to the block decomposition in the proof of theorem 183.5.

Referenced from 3 locations

Reverse-Mode Automatic Differentiation and Cotangent Semantics

Exercise 183.4 ¶

★★☆ Count the scalar operations performed by $v_{t}$ and by one run of $b_{t}$ for $t : = let y = ς (x_{1} * x_{2}) in y * y$ , using definition 183.4 literally. Then count them again for the version in which the backpropagator accumulates into a mutable environment, and exhibit the difference as a function of the number $n$ of variables in scope.

Referenced from 3 locations

Reverse-Mode Automatic Differentiation and Cotangent Semantics

Exercise 183.5 ¶

★★☆ Let $t$ have $n$ real inputs and $m$ real outputs. Using theorem 182.12, theorem 183.5, state the number of runs of each transformed program needed to obtain the full Jacobian, and determine for which pairs $(n, m)$ forward mode is preferable. Give a term for which the two counts are equal.

Referenced from 2 locations

Reverse-Mode Automatic Differentiation and Cotangent Semantics

Exercise 183.6 ¶

★★★ Practical project.reverse-cotangent-transposer Implement in Agda or Kappa the fragment of definition 183.1, its evaluator, the macro $\overset{\leftarrow}{D}$ of definition 183.4, and an operation counter for both the source program and the transformed program. Represent cotangent environments in two interchangeable ways: as tuples added pointwise, and as a sparse accumulator with in-place update.

The invariant to maintain is equation 183.1: for every term, every input, and every output cotangent, the value returned by the backpropagator equals the transposed Jacobian applied to that cotangent. The concrete result is a function taking a term, an input environment, and an output cotangent, and returning the primal value, the input cotangent, and the two operation counts.

Acceptance test. For $e = ⟨ x_{1} * x_{2}, x_{1} + x_{2} ⟩$ at $(3, 5)$ , the cotangent $(1, 0)$ returns $(5, 3)$ and $(0, 1)$ returns $(1, 1)$ ; for $t_{n}$ of example 183.6 with $n = 6$ at all inputs equal to $2$ , one run returns the gradient $(32, 32, 32, 32, 32, 32)$ and the reported operation count of the transformed program is at most six times that of the source program; for $s_{m}$ of example 183.8 with $m = 8$ and $n = 64$ variables in scope, the tuple representation reports a count of order $m n$ while the accumulator representation reports a count of order $m$ , and both return the same gradient. A mutation that drops the $ι_{i}$ in the variable clause, adding the cotangent to every slot, must fail the first test.

The program checks finitely many terms at finitely many points. It illustrates theorem 183.5 and the contrast of example 183.8 with theorem 183.9; it proves neither, and its operation counts are not the cost model of theorem 183.9 unless the two are compared explicitly.

Referenced from 3 locations

Reverse-Mode Automatic Differentiation and Cotangent Semantics

Bibliographic notes

The reverse macro of definition 183.4 is the first-order fragment of the CHAD transformation, presented so that the transposition structure is visible in each clause; theorem 183.5 is proved here rather than imported because at that fragment the induction is short and exhibits the mechanism. Theorem 183.7 is Theorem 2 of Vákár, Reverse AD at Higher Types: Pure, Principled and Denotationally Correct, arXiv:2007.05283, whose categorical derivation of both macros from a single indexed structure explains why the forward and reverse clauses are transposes of each other. Logical relations for typed AD in the presence of partial language features are treated by Nunes and Vákár, whose source-language restrictions are explicit and are not weakened here.

Theorem 183.9 is Equation (7) of [SV24], with the machine-checked development [SV23]; the cotangent-accumulation representation described before it is theirs. The functional reverse-mode transformation of Pearlmutter and Siskind [PS08] is the classic implementation-oriented predecessor and is a comparison only. The linearize–unzip–transpose arrangement discussed in the comparison is the design of You Only Linearize Once; it owns no theorem here.