Forward and reverse AD

All modes of automatic differentiation exploit the idea that derivatives satisfy the chain rule:

\[ D_x (f\circ g) = D_{g(x)}(f) \cdot D_x(g)\]

where we write the derivative of $h : \mathbb R^n \to \mathbb R^m$ at the point $x \in \mathbb R^n$ as $D_x (h)$ , a linear map $\underline{\mathbb{R}}^n\multimap \underline{\mathbb{R}}^m$ that computes directional derivatives.Footnote ² In linear algebra, this linear map corresponds to the Jacobian matrix of partial derivatives of h: if Jh(x) is the Jacobian of h at x, then $D_x(h)(v) = Jh(x) \cdot v$ .

Using the chain rule, one can mechanically compute derivatives of a composition of functions. The intermediate function values (such as x and g(x) in the example above) are called primals; the (forward) derivative values (e.g. the input and output of $D_x(g)$ and $D_{g(x)}(f)$ ) are called tangents.

Forward AD directly implements the chain rule above as a code transformation to compute derivatives. Reverse AD instead computes the transposed derivative, which satisfies this contravariant chain rule instead:

\[ D_x (f\circ g)^t = D_x(g)^t \cdot D_{g(x)}(f)^t\]

Here, we write $D_x (h)^t:\underline{\mathbb{R}}^m\multimap \underline{\mathbb{R}}^n$ for the transpose of the linear map $D_x(h):\underline{\mathbb{R}}^n\multimap \underline{\mathbb{R}}^m$ ; in relation to the Jacobian of h, we have $D_x(h)^t(w) = w \cdot Jh(x)$ . The values taken by and produced from transposed derivatives (or reverse derivatives), such as $D_x(g)^t$ and $D_{g(x)}(f)^t$ above, are called cotangents in this paper; other literature also uses the word adjoint for this purpose.Footnote ³

Computing the full Jacobian of a function $f : \mathbb R^n \to \mathbb R^m$ at $x : \mathbb R^n$ requires n evaluations of $D_x(f) : \underline{\mathbb{R}}^n \multimap \underline{\mathbb{R}}^m$ (at the n basis vectors of $\underline{\mathbb{R}}^n$ ; one such evaluation computes one column of the Jacobian), or alternatively m evaluations of $D_x(f)^t : \underline{\mathbb{R}}^m \multimap \underline{\mathbb{R}}^n$ (one such evaluation computes one row of the Jacobian). Because many applications have $n \gg m$ , reverse AD is typically of most interest, but forward AD is much simpler: for both modes (forward and reverse), we need the primals in the derivative computation, but in forward mode the derivatives are computed in the same order as the original computation (and hence as the computation of the primals). In reverse mode, the derivatives are computed in opposite order, requiring storage of allFootnote ⁴ primals during a forward pass of program execution, after which the reverse pass uses those stored primals to compute the derivatives.Footnote ⁵

This observation that, in forward AD, primals and tangents can be computed in tandem, leads to the idea of dual-numbers forward AD: pair primals and derivatives explicitly to interleave the primal and derivative computations, and run the program with overloaded arithmetic operators to propagate these tangents forward using the chain rule. The tangent of the output can be read from the tangents paired up with the output scalars of the program. For example, transforming the program in Figure 1(a) using dual-numbers forward AD yields Figure 1(b). As a result, we do not store the primals any longer than is necessary, and the resulting code transformation tends to be much simpler than a naive attempt that first computes and stores all primals before computing any derivatives.

Fig. 1. An example program together with its derivative, both using dual-numbers forward AD and using dual-numbers reverse AD. The original program is of type $(\mathbb R, \mathbb R) \rightarrow \mathbb R$ .

Naive dual-numbers reverse AD

For reverse AD, it is in general not possible to interleave the primal and derivative computations, as the reversal of derivatives generally requires us to have computed all primals before starting the derivative computation. However, by choosing a clever encoding, it is possible to make reverse AD look like a dual-numbers-style code transformation. Even if the primal and derivative computations will not be performed in an interleaved fashion, we can interleave the primal computation with a computation that builds a delayed reverse derivative function. Afterwards, we still need to call the derivative function that has been built. The advantage of such a “dual-numbers-style” approach to reverse AD is that the code transformation can be simple and widely applicable.

To make reverse AD in dual-numbers style possible, we have to encode the “reversal” in the tangent scalars that we called dx and dy in Figure 1(b). A solution is to replace those tangent scalars with linear functions that take the cotangent of the scalar it is paired with, and return the cotangent of the full input of the program. Transforming the same example program Figure 1(a) using this style of reverse AD yields Figure 1(c). The linearity indicated by the $\multimap $ -arrow here is that of a monoid homomorphism (a function preserving 0 and (+)Footnote ⁶ ); operationally, however, linear functions behave just like regular functions.

This naive dual-numbers reverse AD code transformation, shown in full in Figure 6, is simple and it is easy to see that it is correct via a logical relations argument (Lucatelli Nunes and Vákár, Reference Lucatelli Nunes and Vákár2024). The idea of this argument is to prove via induction that a backpropagator $\mathit{dx} : \underline{\mathbb{R}}\multimap c$ that is paired with an intermediate value $x : \mathbb R$ in the program, computes the reverse (i.e. transposed) derivative of the subcomputation that calculates $x : \mathbb R$ from the global input to the program.Footnote ⁷ c is a type parameter of the code transformation, and represents the type of cotangents to the program input; if this input (of type $\tau$ , say) is built from just scalars, discrete types and product and sum types, as we will assume in this paper, it suffices (for correctness, not efficiency) to set $c = \tau$ .Footnote ⁸

Dual-numbers forward AD has the very useful property that it generalises over many types (e.g. products, coproducts, and recursive types) and program constructs (e.g. recursion, higher-order functions), thereby being applicable to many functional languages; the same property is inherited by the style of dual-numbers reverse AD exemplified here. However, unlike dual-numbers forward AD (which can propagate tangents through a program with only a constant-factor overhead over the original runtime), naive dual-numbers reverse AD is wildly inefficient: calling $dx_n$ returned by the differentiated program in Figure 2 takes time exponential in n. Such overhead would make reverse AD completely useless in practice—particularly because other (less flexible) reverse AD algorithms exist that indeed do a lot better (see e.g. Griewank and Walther (Reference Griewank and Walther2008) and Baydin et al. (Reference Baydin, Pearlmutter, Radul and Siskind2017)).

Fig. 2. Left: an example showing how naive dual-numbers reverse AD can result in exponential blow-up when applied to a program with sharing. Right: the dependency graph of the backpropagators $dx_i$ .

Fortunately, it turns out that this naive form of dual-numbers reverse AD can be optimised to be as efficient (in terms of time complexity) as these other algorithms—and most of these optimisations are just applications of standard functional programming techniques. This paper presents a sequence of changes to the code transformation (see the overview in Figure 3) that fix all the complexity issues and, in the end, produce an algorithm with which the differentiated program has only a constant-factor overhead in runtime over the original program. This complexity is as desired from a reverse AD algorithm, and is best possible, while nevertheless being applicable to a wide range of programming language features. The last algorithm from Figure 3 can be enhanced to differentiate task-parallel source programs and can also be further optimised to something essentially equivalent to classical taping techniques.

Fig. 3. Overview of the optimisations to dual-numbers reverse AD as a code transformation that are described in this paper. († = inspired by Brunel et al. (Reference Brunel, Mazza and Pagani2020))

Optimisation steps

The first step in Figure 3 is to apply linear factoring: for a linear function f, such as a backpropagator, we have that $f\ x + f\ y = f\ (x + y)$ . Observing the form of the backpropagators in Figure 6, we see that in the end all we produce is a giant sum of applications of backpropagators to scalars; hence, in this giant sum, we should be able to contract applications of the same backpropagator using this linear factoring rule. The hope is that we can avoid executing f more often than is strictly necessary if we represent and reorganise these applications at runtime in a sufficiently clever way.

We achieve this linear factoring by not returning a plain c (presumably the type of (cotangents to) the program input) from our backpropagators, but instead a c wrapped in an object that can delay calls to linear functions producing a c. This object we call ${\textrm{Staged}}$ ; aside from changing the monoid that we are mapping into from $(c, \underline0, (+))$ to $({\textrm{Staged}} c, 0_{\text{Staged}}, (\mathbin{+_{\text{Staged}}}))$ , the only material change is that the calls to argument backpropagators in $\textbf{D}^{1}_{c}[\textit{op}]$ are now wrapped using a new function $\textrm{SCall}$ , which delays the calls to $d_i$ by storing the relevant metadata in the returned ${\textrm{Staged}}$ object.

However, it is not obvious how to implement this ${\textrm{Staged}}$ type: at the very least, we need decidable equality on linear functions to be able to implement the linear factoring rule; and if we want any hope of an efficient implementation, we even need a linear order on such linear functions so that we can use them as keys in a tree map in the implementation of ${\textrm{Staged}}$ . Furthermore, even if we can delay calls to backpropagators, we still need to call them at some point, and it is unclear in what order we should do so (while this order turns out to be very important for efficiency).

We thus need two orders on our backpropagators. It turns out that for sequential input programs, it suffices to generate a unique, monotonically increasing identifier (ID) for each backpropagator that we create and use that ID not only as a witness for backpropagator identity and as a key in the tree map, but also as a witness for the dependency order (see below) of the backpropagators. These IDs are generated by letting the differentiated program run in an ID generation monad (a special case of a state monad). The result is shown in Figure 8, which is very similar to the previous version in Figure 6 apart from threading through the next-ID-to-generate. (On first glance, the code looks very different, but this is only due to monadic bookkeeping.)

At this point, the code transformation already reaches a significant milestone: by staging (delaying) calls to backpropagators as long as possible, we can ensure that every backpropagator is called at most once. This milestone can be seen using the following observation: lambda functions in a pure functional program that do not take functions as arguments, can only call functions that appear in their closure. Because backpropagators are never mutually recursive (that could only happen if their corresponding scalars are defined mutually recursively, which, being scalars, would never terminate anyway), this observation means that calling backpropagators (really unfolding, as we delay evaluating subcalls) in their reverse dependency order achieves the goal of delaying the calls as long as possible. Thus, for monotonically increasing IDs, a backpropagator will only call other backpropagators with lower IDs. If we do not need parallelism in the differentiated programs, we can therefore suffice by simply resolving backpropagators from the highest to the lowest ID.

Effectively, one can think of the IDs as (labels of) the nodes in the computation graph of this run of the original program; from this perspective, the resolving process is nothing more than a choice of a reverse topological order of this graph, so that we can perform backpropagation (the reverse pass) in the correct order without redundant computation.

But we are not done yet. The code transformation at this point ( $\textbf{D}^{2}_{c}$ in Figure 8) still has a glaring problem: orthogonal to the issue that backpropagators were called too many times (which we fixed), we are still creating one-hot input cotangents and adding those together. This problem is somewhat more subtle because it is not actually apparent in the program transformation itself; indeed, looking back at Figure 1(c), there are no one-hot values to be found. However, the only way to use the program in Figure 1(c) to do something useful, namely to compute the cotangent (gradient) of the input, is to pass $(\lambda z.\ (z, 0))$ to $\textit{dx}$ and $(\lambda z.\ (0, z))$ to $\textit{dy}$ ; it is easy to see that generalising this to larger input structures results in input values like $(0, \ldots, 0, z, 0, \ldots, 0)$ that get added together. (These are created in the wrapper in Figure 7.) Adding many zeros together can hardly be the most efficient way to go about things, and indeed this is a complexity issue in the algorithm.

The way we solve this problem of one-hots is less AD-specific: the most important optimisations that we perform are Cayley transformation (Section 7) and using a better sparse vector representation ( $\textrm{Map}\ \mathbb Z\ \underline{\mathbb{R}}$ instead of a plain c value; Section 8). Cayley transformation (also known in the Haskell community by a common use: difference lists (Hughes, Reference Hughes1986)) is a classic technique in functional programming that represents an element m of a monoid M (written additively in this paper) by the function $m + - : M \rightarrow M$ it induces through addition. Cayley transformation helps us because the monoid $M \rightarrow M$ has very cheap zero and plus operations: id and $(\circ)$ . Afterwards, using a better (sparse) representation for the value in which we collect the final gradient, we can ensure that adding a one-hot value to this gradient collector can be done in logarithmic time.

By now, the differentiated program can compute the gradient with a logarithmic overhead over the original program. If a logarithmic overhead is not acceptable, the log-factor in the complexity can be removed by using functional mutable arrays (Section 9). Such mutability can be safely accommodated in our code transformation by either swapping out the state monad for a resource-linear state monad, or by using mutable references in an ST-like monad (Section 9.1). (The latter can be generalised to the parallel case; see below.)

And then we are done for sequential programs, because we have now obtained a code transformation with the right complexity: the differentiated program computes the gradient of the source program at some input with runtime proportional to the runtime of the source program at that input.

Correctness

Correctness of the resulting AD technique follows in two steps: 1. the naive dual-numbers reverse AD algorithm we start with is correct by a logical relations argument detailed by Lucatelli Nunes and Vákár (Reference Lucatelli Nunes and Vákár2024); 2. we transform this into our final algorithm using a combination of (A) standard optimisations that are well-known to preserve semantics (and hence correctness)—notably sparse vectors via Cayley transformationFootnote ⁹ and using a mutable array to optimise a tree map—and (B) the custom optimisation of linear factoring, which is semantics-preserving because derivatives (backpropagators) are linear functions.

Parallelism

If the user is satisfied with a fully sequential (non-parallel) computation of the derivative, it is enough to generate monotonically increasing integers and use them, together with their linear order, as IDs for backpropagators during the forward pass. This is the approach described so far.

However, if the source program has (task) parallelism (we assume fork-join parallelism in this paper), we would prefer to preserve that parallelism when performing backpropagation on the (implied) computation graph. The linear order (chronological, by comparing IDs) that we were using to witness the dependency order so far does not suffice any more: we need to be more frugal with adding spurious edges in the dependency graph that only exist because one backpropagator happened to be created after another on the clock. However, we only need to make our order (i.e. dependency graph) precise enough that independent tasks are incomparable in the order (and hence independent in the dependency graph); recording more accurate (lack of) dependencies between individual scalar operations would even allow exploiting implicit parallelism within an a priori serial subcomputation, which is potentially interesting but beyond the scope of this paper.Footnote ¹⁰

Our solution is thus to switch from simple integer IDs to compound IDs, consisting of a job ID and a sequential ID within that job.Footnote ¹¹ We assume parallelism in the source program is expressed using a parallel pair constructor with the following typing rule:

\[ \frac{\Gamma \vdash t : \tau \qquad \Gamma \vdash s : \sigma} {\Gamma \vdash t \star s : (\tau, \sigma)}\]

(The method generalises readily to n-ary versions of this primitive.) To differentiate code using this construct, we take out the ID generation monad that the target program ran in so far, and replace it with a monad in which we can also record the dependency graph between parallel jobs. The derivative of $(\star)$ is the only place where we use the new methods of this monad: all other code transformation rules remain identical, save for writing the right-hand sides as a black-box monad instead of explicit state passing. We can then make use of this additional recorded information in the backpropagator resolution phase to do so in parallel; in this process, the net effect is that forks from the primal become joins in the derivative computation, and vice versa.

Comparison to other algorithms

We can relate the sequential version of our technique to that of Krawiec et al. (Reference Krawiec, Jones, Krishnaswami, Ellis, Eisenberg and Fitzgibbon2022) by noting that we can replace $\textbf{D}^{1}_{c}[\mathbb R] = (\mathbb R, \underline{\mathbb{R}}\multimap c)$ with the isomorphic definition $\textbf{D}^{1}_{c}[\mathbb R] = (\mathbb R, c)$ . This turns the linear factoring rule into a distributive law $v\cdot x+v\cdot y\leadsto v\cdot (x+y)$ that is effectively applied at runtime by using an intensional representation of the cotangent expressions of type c. While their development is presented very differently and the equivalence is not at all clear at first sight, we explain the correspondence in Section 10.

This perspective also makes clear the relationship between our technique and that of Shaikhha et al. (Reference Shaikhha, Fitzgibbon, Vytiniotis and Jones2019). Where they try to optimise vectorised forward AD to reverse AD at compile time by using a distributive law (which sometimes succeeds for sufficiently simple programs), our technique proposes a clever way of efficiently applying the distributive law in the required places at runtime, giving us the power to always achieve the desired reverse AD behaviour.

Finally, we are now in the position to note the similarity to (sequentialFootnote ¹² ) taping-based AD as used by Kmett and Contributors (2021), older versions of PyTorch (Paszke et al., Reference Paszke, Gross, Chintala, Chanan, Yang, DeVito, Lin, Desmaison, Antiga and Lerer2017), etc.: the incrementing IDs that we attached to backpropagators earlier give a mapping from $\{0,\ldots,n\}$ to our backpropagators. Furthermore, each backpropagator corresponds to either a primitive arithmetic operation performed in the source program, or to an input value; this already means that we have a tape, in a sense, of all performed primitive operations, albeit in the form of a chain of closures. The optimisation using mutable arrays (Section 9) which reifies this tape in a large array in the reverse pass, especially if one then proceeds to already use this array in the forward pass (Section 10.3), eliminates also this last difference.

Typing forward AD

For forward AD on first-order programs (or at least, programs for which the input and output does not contain function values), the desired type seems quite evident: $\mathcal F : (\sigma \rightarrow \tau) \rightsquigarrow ((\sigma, \underline{\sigma}) \rightarrow (\tau, \underline{\tau}))$ , where we write $T_1 \rightsquigarrow T_2$ for a (compiler) code transformation taking a program of type $T_1$ and returning a program of type $T_2$ , and where $\underline{\tau}$ is the type of tangent vectors (derivatives) of values of type $\tau$ . This distinction between the type of values $\tau$ and the type of their derivatives $\underline{\tau}$ is important in some versions of AD, but will be mostly cosmetic in this paper; in an implementation one can take $\underline{\tau} = \tau$ , but there is some freedom in this choice.Footnote ¹⁴ Given a program $f : \sigma \rightarrow \tau$ , $\mathcal F[f]$ is a program that takes, in addition to its regular argument, also a tangent at that argument; the output is then the regular result paired up with corresponding tangent at that result.

More specifically, for forward AD, we want the following in the case that $\sigma = \mathbb R^n$ and $\tau = \mathbb R^m$ (writing $\mathbf{x} = (x_1, \ldots, x_n)$ ):Footnote ¹⁵

\begin{align*}\mathcal F[f]\left(\mathbf{x}, \left(\frac{\partial x_1}{\partial\alpha}, \ldots, \frac{\partial x_n}{\partial \alpha}\right)\right)= \left(f(\mathbf{x}), \left(\frac{\partial f(\mathbf{x})_1}{\partial\alpha}, \ldots, \frac{\partial f(\mathbf{x})_m}{\partial\alpha}\right)\right)\end{align*}

Setting $\alpha = x_i$ means passing $(0, \ldots, 1, \ldots, 0)$ as the argument of type $\underline{\sigma}$ and computing the partial derivative with respect to $x_i$ of $f(\mathbf{x})$ . In other words, $\textrm{snd}(\mathcal F[f](x, \mathit{dx}))$ is the directional derivative of f at x in the direction $\mathit{dx}$ .

A first attempt at typing reverse AD

For reverse AD, the desired type is less evident. A first guess would be:

\begin{align*}\mathcal R_1 : (\sigma \rightarrow \tau) \rightsquigarrow ((\sigma, \underline{\tau}) \rightarrow (\tau, \underline{\sigma}))\end{align*}

with this intended meaning for $\sigma = \mathbb R^n$ and $\tau = \mathbb R^m$ (again writing $\mathbf{x} = (x_1, \ldots, x_n)$ ):

\begin{align*}\mathcal R_1[f]\left(\mathbf{x}, \left(\frac{\partial\omega}{\partial f(\mathbf{x})_1}, \ldots, \frac{\partial\omega}{\partial f(\mathbf{x})_m}\right)\right)= \left(f(\mathbf{x}), \left(\frac{\partial\omega}{\partial x_1}, \ldots, \frac{\partial\omega}{\partial x_n}\right)\right)\end{align*}

In particular, if $\tau = \mathbb R$ and we pass 1 as its cotangent (also called adjoint) of type $\underline{\tau} = \underline{\mathbb{R}}$ , the $\underline{\sigma}$ -typed output contains the gradient with respect to the input.

Dependent types

However, $\mathcal R_1$ is not readily implementable for even moderately interesting languages. One way to see this is to acknowledge the reality that the type $\underline{\tau}$ (of derivatives of values of type $\tau$ ) should really be dependent on the accompanying primal value of type $\tau$ . Let us write the type of derivatives at $\tau$ not as $\underline\tau$ but as $\mathcal D[\tau](x)$ , where $x : \tau$ is that primal value. With just scalars and product types this dependence does not yet occur (e.g. $\mathcal D[\mathbb R](x) = \mathbb R$ independent of the primal value x), but when adding sum types (coproducts), the dependence becomes non-trivial: the only sensible derivatives for a value $\textrm{inl}(x) : \sigma + \tau$ (for $x : \sigma$ ) are of type $\underline{\sigma}$ . Letting $\underline{\sigma + \tau} = \underline{\sigma} + \underline{\tau}$ would allow passing a derivative value of type $\underline{\tau}$ to $\textrm{inl}(x) : \sigma + \tau$ , which is nonsensical (and an implementation could do little else than return a bogus value like 0 or throw a runtime error). The derivative of $\textrm{inl}(2x) : \mathbb R + {\mathrm{Bool}}$ cannot be $\textrm{inr}({\mathrm{True}})$ ; it should at least somehow contain a real value.

Similarly, the derivative for a dynamically sized array, if the input language supports those, must really be of the same size as the input array. This, too, is a dependence of the type of the derivative on the value of the input.

Therefore, the output type of forward AD which we wrote above as $(\sigma, \underline{\sigma}) \rightarrow (\tau, \underline{\tau})$ should reallyFootnote ¹⁶ be $(\Sigma_{x : \sigma}\, \mathcal D[\sigma](x)) \rightarrow (\Sigma_{y : \tau}\, \mathcal D[\tau](y))$ , rendering what were originally pairs of value and tangent now as dependent pairs of value and tangent. This is a perfectly sensible type, and indeed correct for forward AD, but it does not translate at all well to reverse AD in the form of $\mathcal R_1$ : the output type would be something like $(\Sigma_{x : \sigma}\, \mathcal D[\tau](y)) \rightarrow (\Sigma_{y : \tau}\, \mathcal D[\sigma](x))$ , which is nonsense because both x and y are out of scope.

Let-bindings

A different way to see that the type of $\mathcal R_1$ is unusable, is to note that one cannot even differentiate let-bindings using $\mathcal R_1$ . In order to apply to (an extension of) the lambda calculus, let us rewrite the types somewhat: where we previously put a function $f : \sigma \rightarrow \tau$ , we now put a term $x : \sigma \vdash t : \tau$ with its input in a free variable and producing its output as the returned value. Making the modest generalisation to support any full environment as input (instead of just a single variable), we get $\mathcal R_1 : (\Gamma \vdash t : \tau) \rightsquigarrow (\Gamma, d : \underline{\tau} \vdash \mathcal R_1[t] : (\tau, \underline{\Gamma}))$ , where $\underline{\Gamma}$ is a tuple containing the derivatives of all elements in the environment $\Gamma$ . (To be precise, we define $\underline{\varepsilon} = ()$ for the empty environment and $\underline{\Gamma, x : \tau} = (\underline{\Gamma}, \underline{\tau})$ inductively.)

Now, consider differentiating the following program using $\mathcal R_1$ :

\begin{align*}\Gamma \vdash (\mathbf{let}\ x = e_1\ \mathbf{in}\ e_2) : \tau\end{align*}

where $\Gamma \vdash e_1 : \sigma$ and $\Gamma, x : \sigma \vdash e_2 : \tau$ . Substituting, we see that $\mathcal R_1$ needs to somehow build a program of this type:

(4.1) \begin{align}\Gamma, d : \underline{\tau} \vdash \mathcal R_1[\mathbf{let}\ x = e_1\ \mathbf{in}\ e_2] : (\tau, \underline{\Gamma})\end{align}

However, recursively applying $\mathcal R_1$ on $e_1$ and $e_2$ yields terms:

\begin{align*}\Gamma, d : \underline{\sigma} &\vdash \mathcal R_1[e_1] : (\sigma, \underline{\Gamma}) \\\Gamma, x : \sigma, d : \underline{\tau} &\vdash \mathcal R_1[e_2] : (\tau, (\underline{\Gamma}, \underline{\sigma}))\end{align*}

To produce the program in Equation (4.1), we cannot use $\mathcal R_1[e_2]$ because we do not yet have an $x : \sigma$ (which needs to come from $\mathcal R_1[e_1]$ ), and we cannot use $\mathcal R_1[e_1]$ because the $\underline{\sigma}$ needs to come from $\mathcal R_1[e_2]$ ! The type of $\mathcal R_1$ demands the cotangent of the result too early.

Of course, one might argue that we can just use $e_1$ to compute the $\sigma$ , $\mathcal R_1[e_2]$ to get the $\underline{\sigma}$ and $e_2$ ’s contribution to $\underline{\Gamma}$ , and finally $\mathcal R_1[e_1]$ to get $e_1$ ’s contribution to $\underline{\Gamma}$ based on its own cotangent of type $\underline{\sigma}$ . However, this would essentially compute $e_1$ twice (once directly and once as part of $\mathcal R_1[e_1]$ ), meaning that the time complexity becomes super-linear in the depth of let-bindings, which is quite disastrous for typical functional programs.

So in addition to not being precisely typeable, $\mathcal R_1$ is also not implementable in a compositional way.

Fixing the type of reverse AD

Both when looking at the dependent type of $\mathcal R_1$ and when looking at its implementation, we found that the cotangent $\mathit{dy} : \mathcal D[\tau](y)$ was required before the result $y : \tau$ was itself computed. One way to solve this issue is to just postpone requiring the cotangent of y, i.e. to instead look at $\mathcal R_2$ :Footnote ¹⁷

\begin{align*}\begin{array}{l@{\qquad}c} \mathcal R_2 : (\sigma \rightarrow \tau) \rightsquigarrow (\sigma \rightarrow (\tau, \underline{\tau} \rightarrow \underline{\sigma})) & \text{(non-dependent version)} \\[0.2em] \mathcal R_2 : (\sigma \rightarrow \tau) \rightsquigarrow (\Pi_{x : \sigma}\, \Sigma_{y : \tau}\, (\mathcal D[\tau](y) \rightarrow \mathcal D[\sigma](x))) & \text{(dependent version)}\end{array}\end{align*}

Note that this type is well-scoped. Furthermore, this “derivative function” mapping the cotangent of the result to the cotangent of the argument is actually a linear function, in the sense of a vector space homomorphism: indeed, it is multiplication by the Jacobian matrix of f, the input function. Thus we can write:

\begin{align*}\mathcal R_2 : (\sigma \rightarrow \tau) \rightsquigarrow (\Pi_{x : \sigma}\, \Sigma_{y : \tau}\, (\mathcal D[\tau](y) \multimap \mathcal D[\sigma](x)))\end{align*}

which is the type of the reverseFootnote ¹⁸ AD code transformation derived by Elliott (Reference Elliott2018) and in CHAD (e.g. Vákár and Smeding (Reference Vákár and Smeding2022); see also Section 15.3).Footnote ¹⁹

While this formulation of reverse AD admits a rich mathematical foundation (Nunes and Vákár, Reference Nunes and Vákár2023) and has the correct complexity (Smeding and Vákár, Reference Smeding and Vákár2024), the required program transformation is more complex than the formulation $\mathcal F$ that we have for forward AD. In particular, we need to compute for each programming language construct what its CHAD transformation is, which may be non-trivial (for example, for the case of function types). This difficulty motivates us to pursue a reverse AD analogue of $\mathcal F$ .

Applying Yoneda/CPS

An instance of the Yoneda lemma (or in this case: continuation-passing style; see also Boisseau and Gibbons (Reference Boisseau and Gibbons2018)) is that $\sigma \multimap \tau$ is equivalent to $\forall r.\, (\sigma \multimap r) \to (\tau \multimap r)$ . We can apply this to the $\multimap $ -arrow in $\mathcal R_2$ to obtain a type for reverse AD that is somewhat reminiscent of our formulation $\mathcal F$ of forward AD. With just Yoneda, we get $\mathcal R_3'''$ below; we then weaken this type somewhat by enlarging the scope of the $\forall c$ quantifier, weaken some more by taking the $(\mathcal D[\tau](x) \multimap c)$ argument before returning y, and finally we uncurry to arrive at $\mathcal R_3$ :

We give both a dependently typed (black) and a simply typed (grey) signature for $\mathcal R_3$ .

The $\multimap $ -arrows in these types, as well as the c bound by the $\forall$ -quantifier, live in the category of commutative monoids. Indeed, c will always have a commutative monoid structure in this paper; that is: it has a zero $\underline 0$ as well as a commutative, associative addition operation $(+) : (c, c) \rightarrow c$ for which $\underline 0$ is the unit. (The $\multimap $ -arrows in these types are really vector space homomorphisms, but since we will only use the substructure of commutative monoids in this paper, and forget about scalar multiplication, we will always consider $\multimap $ -functions (commutative) monoid homomorphisms.)

Returning to the types in question, we see that we can convert $\mathcal R_2[t]$ to $\mathcal R_3[t]$ :

\[ \begin{array}{@{}l@{}} (\lambda (x : \sigma, \mathit{dx} : \mathcal D[\sigma](x) \multimap c). \\ \hspace{3em} \mathbf{let}\ (y : \tau, \mathit{dy} : \mathcal D[\tau](y) \multimap \mathcal D[\sigma](x)) = \mathcal R_2[t]\ x\ \mathbf{in}\ (y, \mathit{dx} \circ \mathit{dy})) \\ \hspace{1em} : \forall c.\, (\Sigma_{x : \sigma}\, (\mathcal D[\sigma](x) \multimap c)) \rightarrow \Sigma_{y : \tau}\, (\mathcal D[\tau](y) \multimap c) \end{array}\]

where we write $\circ$ for the composition of linear functions. We can also convert $\mathcal R_3[t]$ back to $\mathcal R_2[t]$ , but due to how we weakened the types above, only in the non-dependent world:

\begin{align*}&(\lambda (x : \sigma).\ \mathcal R_3[t]\ (x, \underline{\lambda} (z: \underline \sigma).\ z)) : \sigma \rightarrow (\tau, \underline \tau \multimap \underline \sigma)\end{align*}

So, in some sense, $\mathcal R_2$ and $\mathcal R_3$ compute the same thing, albeit with types that differ in how precisely they portray the dependencies.

In fact, $\mathcal R_3$ admits a very elegant implementation as a program transformation that is structurally recursive over all language elements except for the primitive operations in the leaves. However, there are some issues with the computational complexity of this straightforward implementation of $\mathcal R_3$ , one of which we will fix here immediately, and the other of which are the topic of the rest of this paper.

Moving the pair to the leaves

Let us return to forward AD for a moment. Recall the type we gave for forward AD:Footnote ²⁰

\begin{align*}\mathcal F : (\sigma \rightarrow \tau) \rightsquigarrow ((\sigma, \underline{\sigma}) \rightarrow (\tau, \underline{\tau}))\end{align*}

Supposing we have a program $f : ((\mathbb R_1, \mathbb R_2), \mathbb R_3) \rightarrow \mathbb R_4$ , we get: (the subscripts are semantically meaningless and are just for tracking arguments)

\begin{align*}\mathcal F[f] : (((\mathbb R_1, \mathbb R_2), \mathbb R_3), ((\underline{\mathbb R_1}, \underline{\mathbb R_2}), \underline{\mathbb R_3})) \rightarrow (\mathbb R_4, \underline{\mathbb R_4})\end{align*}

While this is perfectly implementable and correct and efficient, it is not the type that corresponds to what is by far the most popular implementation of forward AD, namely dual-numbers forward AD, which has the following type:

\begin{gather*}\mathcal F_{\text{dual}} : (\sigma \rightarrow \tau) \rightsquigarrow (\textrm{Dual}[\sigma] \rightarrow \textrm{Dual}[\tau]) \\\textrm{Dual}[\mathbb R] = (\mathbb R, \underline{\mathbb{R}}) \qquad\textrm{Dual}[()] = () \qquad\textrm{Dual}[(\sigma, \tau)] = (\textrm{Dual}[\sigma], \textrm{Dual}[\tau])\end{gather*}

Intuitively, instead of putting the pair at the root like $\mathcal F$ does, $\mathcal F_{\text{dual}}$ puts the pair at the leaves—more specifically, at the scalars in the leaves, leaving non- $\mathbb R$ types like () or $\mathbb Z$ alone. For the given example program f, dual-numbers forward AD would yield the following derivative program type:

\begin{align*}\mathcal F_{\text{dual}}[f] : (((\mathbb R_1, \underline{\mathbb R_1}), (\mathbb R_2, \underline{\mathbb R_2})), (\mathbb R_3, \underline{\mathbb R_3})) \rightarrow (\mathbb R_4, \underline{\mathbb R_4})\end{align*}

Of course, for any given types $\sigma, \tau$ the two versions are trivially inter-converted, and as stated, for forward AD both versions can be defined inductively equally well, resulting in efficient programs in terms of time complexity.

However, for reverse AD in the style of $\mathcal R_3$ , the difference between $\mathcal R_3$ and its pair-at-the-leaves dual-numbers variant ( $\mathcal R_{3\text{dual}}$ below) is more pronounced. First note that indeed both styles (with the pair at the root and with the pair at the leaves) produce a sensible type for reverse AD: (again for $f : ((\mathbb R_1, \mathbb R_2), \mathbb R_3) \rightarrow \mathbb R_4$ )

\begin{align*}\mathcal R_3[f] &: \forall c.\, (((\mathbb R_1, \mathbb R_2), \mathbb R_3), ((\underline{\mathbb R_1}, \underline{\mathbb R_2}), \underline{\mathbb R_3}) \multimap c) \rightarrow (\mathbb R_4, \underline{\mathbb R_4} \multimap c) \\\mathcal R_{3\text{dual}}[f] &: \forall c.\, (((\mathbb R_1, \underline{\mathbb R_1} \multimap c), (\mathbb R_2, \underline{\mathbb R_2} \multimap c)), (\mathbb R_3, \underline{\mathbb R_3} \multimap c)) \rightarrow (\mathbb R_4, \underline{\mathbb R_4} \multimap c)\end{align*}

The individual functions of type $\underline{\mathbb R} \multimap c$ are usually called backpropagators in literature, and we will adopt this terminology.

Indeed, these two programs are again easily inter-convertible, if one realises that:

1. 1. c is a commutative monoid and thus possesses an addition operation, which can be used to combine the three c results into one for producing the input of $\mathcal R_3$ from the input of $\mathcal R_{3\text{dual}}$ ;

2. 2. The function $g : ((\underline{\mathbb R_1}, \underline{\mathbb R_2}), \underline{\mathbb R_3}) \multimap c$ is linear, and hence, e.g. $\lambda (x : \underline{\mathbb R_2}).\ g\ ((0, x), 0)$ suffices as value for $\underline{\mathbb R_2} \multimap c$ .

However, the problem arises when defining $\mathcal R_3$ inductively as a program transformation. To observe this difference between $\mathcal R_3$ and $\mathcal R_{3\text{dual}}$ , consider the term $t = \lambda(x : (\sigma, \tau)).\ \textrm{fst}(x)$ of type $(\sigma, \tau) \rightarrow \sigma$ and the types of its derivative using both methods:

\begin{gather*}\begin{array}{r@{\ }l}\mathcal R_3[t] &: \forall c.\ ((\sigma, \tau), (\underline{\sigma}, \underline{\tau}) \multimap c) \rightarrow (\sigma, \underline{\sigma} \multimap c) \\[2pt]\mathcal R_{3\text{dual}}[t] &: \forall c.\ (\textrm{Dual}_c[\sigma], \textrm{Dual}_c[\tau]) \rightarrow \textrm{Dual}_c[\sigma]\end{array} \\\textrm{Dual}_c[\mathbb R] = (\mathbb R, \underline{\mathbb{R}}\multimap c) \qquad\textrm{Dual}_c[()] = () \qquad\textrm{Dual}_c[(\sigma, \tau)] = (\textrm{Dual}_c[\sigma], \textrm{Dual}_c[\tau])\end{gather*}

Their implementations look as follows:

\begin{align*}\mathcal R_3[t] &= \lambda (x : (\sigma, \tau), \mathit{dx} : (\underline{\sigma}, \underline{\tau}) \multimap c).\ (\textrm{fst}(x), \lambda (d : \underline{\sigma}).\ \mathit{dx}\ (d, \underline{0}_{\underline{\tau}})) \\\mathcal R_{3\text{dual}}[t] &= \lambda (x : (\textrm{Dual}_c[\sigma], \textrm{Dual}_c[\tau])).\ \textrm{fst}(x)\end{align*}

where $\underline{0}_{\underline{\tau}}$ is the zero value of the cotangent type of $\tau$ . The issue with the first variant is that $\tau$ may be an arbitrarily complex type, perhaps even containing large arrays of scalars, and hence this zero value $\underline{0}_{\underline{\tau}}$ may also be large. Having to construct this large zero value is not, in general, possible in constant time, whereas the primal operation ( $\textrm{fst}$ ) was a constant-time operation; this is anathema to getting a reverse AD code transformation with the correct time complexity. Further, on our example program, we see that the variant $\mathcal R_{3}$ results in a more complex code transformation than $\mathcal R_{3\text{dual}}$ , and this observation turns out to hold more generally. $\mathcal R_{3}$ shares both these challenges with the CHAD formulation $\mathcal R_{2}$ of reverse AD.

As evidenced by the complexity analysis and optimisation of the CHAD reverse AD algorithm (Smeding and Vákár, Reference Smeding and Vákár2024), there are ways to avoid having to construct a non-constant-size zero value here. In fact, we use one of those ways, in a different guise, later in this paper in Section 7. However, in this paper, we choose the $\mathcal R_{3\text{dual}}$ approach. We pursue the dual numbers approach not to avoid having to deal with the issue of large zeros—indeed, skipping the problem here just moves it somewhere else, namely to the implementation of the backpropagators ( $\underline{\mathbb R} \multimap c$ ). Rather, we pursue this approach because $\mathcal R_{3\text{dual}}$ extends more easily to a variety of language features (see Sections 11 and 12).

Evaluation order

Ensuring this complete merging of linear function calls is really a question of choosing an order of evaluation for the tree of function calls created by the backpropagators. Consider for example the (representative) situation where a program generates the following backpropagators:

Suppose $f_4$ is the (only) backpropagator contained in the result. Normal call-by-value evaluation of $f_4$ would yield two invocations of $f_2$ and five invocations of $f_1$ , following the call graph on the right.

However, taking inspiration from symbolic evaluation and moving away from standard call-by-value for a moment, we could also first invoke $f_3$ to expand the body of $f_4$ to $f_2\ z + f_2\ (4 \cdot (2 \cdot z)) + f_1\ (5 \cdot (2 \cdot z))$ . Now we can take the two invocations of $f_2$ together using linear factoring to produce $f_2\ (z + 4 \cdot (2 \cdot z)) + f_1\ (5 \cdot (2 \cdot z))$ ; then invoking $f_2$ first, producing two more calls to $f_1$ , we are left with three calls to $f_1$ which we can take together to a single call using linear factoring, which we can then evaluate. With this alternate evaluation order, we have indeed ensured that every linear function is invoked at most (in this case, exactly) once.

If we want to obtain something like this evaluation order, the first thing that we must achieve is to postpone invocation of linear functions until we conclude that we have merged all calls to that function and that its time for evaluation has arrived. To achieve this goal, we would like to change the representation of c to a dictionary mapping linear functions to the argument at which we intend to later call them.Footnote ²⁴ Note that this uniform representation in a dictionary works because all backpropagators in the core transformed program (outside of the wrapper) have the same domain ( $\underline{\mathbb{R}}$ ) and codomain (c). The idea is that we replace what are now applications of linear functions with the creation of a dictionary containing one key-value (function-argument) pair and to replace addition of values in c with taking the union of dictionaries, where arguments for common keys are added together.

Initial Staged object

More concretely, we want to change $\textbf{D}^{1}_{c}[\mathbb R] = (\mathbb R, \underline{\mathbb{R}}\multimap c)$ to instead read $\textbf{D}^{1}_{c}[\mathbb R] = (\mathbb R, \underline{\mathbb{R}}\multimap {\textrm{Staged}} c)$ , where ‘ ${\textrm{Staged}} c$ ’ is our “dictionary”.Footnote ²⁵ We define, or rather, would like to define ${\textrm{Staged}} c$ as follows: (‘ $\textrm{Map} \ k \ v$ ’ is the usual type of persistent tree-maps with keys of type k and values of type v)

\begin{align*} {\textrm{Staged}} c = (c, \textrm{Map}\ (\underline{\mathbb{R}}\multimap {\textrm{Staged}} c)\ \underline{\mathbb{R}})\end{align*}

Suspending disbelief about implementability, this type can represent both literal c values (necessary for the one-hot vectors returned by the input backpropagators created in $\textrm{Interleave}^{1}$ ) and staged (delayed) calls to linear functions. We use $\textrm{Map}$ to denote a standard (persistent) tree-map as found in every functional language. The intuitive semantics of a value $(x, \{f_1 \mapsto a_1, f_2 \mapsto a_2\}) $ of type $ {\textrm{Staged}} c$ is its resolution $x + f_1\ a_1 + f_2\ a_2:c$ .

To be able to replace c with ${\textrm{Staged}} c$ in $\textbf{D}^{1}_{c}$ , we must support all operations that we perform on c also on ${\textrm{Staged}} c$ . We implement them as follows:

• • $\underline0 : c$ becomes simply $0_{\text{Staged}} := (\underline0, \{\}) : {\textrm{Staged}} c$ .

• • $(+) : c \rightarrow c \rightarrow c$ becomes $(\mathbin{+_{\text{Staged}}})$ , adding c values using $(+)$ and taking the union of the two $\textrm{Map}$ s. Here we apply linear factoring: if the two $\textrm{Map}$ s both have a value for the same key (i.e. we have two staged invocations of the same linear function f), the resulting map will have one value for that same key f: the sum of the arguments stored in the two separate $\textrm{Map}$ s. For example:

$ \begin{array}{l} (c_1, \{f_1 \mapsto a_1, f_2 \mapsto a_2\}) \mathbin{+_{\text{Staged}}} (c_2, \{f_2 \mapsto a_3\}) \\ \qquad = (c_1 + c_2, \{f_1 \mapsto a_1, f_2 \mapsto a_2 + a_3\}) \end{array} $

• • The one-hot c values created in the backpropagators from $\textrm{Interleave}^{1}$ are stored in the c component of ${\textrm{Staged}} c$ .

• • An application $f\ x$ of a backpropagator $f : \underline{\mathbb{R}}\multimap c$ to an argument $x : \underline{\mathbb{R}}$ now gets replaced with $\textrm{SCall}\ f\ x := (\underline0, \{f \mapsto x\}) : {\textrm{Staged}} c$ . This occurs in $\textbf{D}^{1}_{c}[\textit{op}(...)]$ and in $\textrm{Deinterleave}^{1}$ .

Essentially, this step of replacing c with ${\textrm{Staged}} c$ can be seen as a clever partial defunctionalisation of our backpropagators.

What is missing from this list is how to “resolve” the final ${\textrm{Staged}} c$ value produced by the derivative computation down to a plain c value—we need this at the end of the wrapper. This resolve algorithm:

\[\textrm{SResolve} : ({\textrm{Staged}} c) \to c\]

will need to call functions stored in the ${\textrm{Staged}} c$ object in the correct order, ensuring that we only invoke a backpropagator when we are sure that we have collected all calls to it in the $\textrm{Map}$ . For example, in the example at the beginning of this section, $f_4\ 1$ returns $(\underline0, \{f_2 \mapsto 1, f_3 \mapsto 2\})$ . At this point, “resolving $f_3$ ” means calling $f_3$ at 2, observing the return value $(\underline0, \{f_2 \mapsto 8, f_1 \mapsto 10\})$ , and adding it to the remainder (i.e. without the $f_3$ entry) of the previous ${\textrm{Staged}} c$ object to get $(\underline0, \{f_2 \mapsto 9, f_1 \mapsto 10\})$ .

But as we observed above, the choice of which function to invoke first is vital to the complexity of the reverse AD algorithm: if we chose $f_2$ first instead of $f_3$ , the later call to $f_3$ would produce another call to $f_2$ , forcing us to evaluate $f_2$ twice—something that we must avoid. There is currently no information in a ${\textrm{Staged}} c$ object from which we can deduce the correct order of invocation, so we need something extra.

There is another problem with the current definition of ${\textrm{Staged}} c$ : it contains a $\textrm{Map}$ keyed by functions, meaning that we need equality—actually, even an ordering—on functions! This is nonsense in general. Fortunately, both problems can be tackled with the same solution.

Resolve order

The backpropagators that occur in the derivative program (as produced by $\textbf{D}^{1}_{c}$ from Figure 6) are not just arbitrary functions. Indeed, taking the target type c of the input backpropagators to be equal to the input type $\sigma$ of the original program (of type $\sigma \rightarrow \tau$ ), as we do in $\textrm{Wrap}^{1}$ in Figure 7, all backpropagators in the derivative program have one of the following three forms:

1. 1. $(\underline{\lambda}(z : \underline{\mathbb{R}}).\ t)$ where t is a tuple (of type $\sigma$ ) filled with zero scalars except for one position, where it places z; we call such tuples one-hot tuples. These backpropagators result, after trivial beta-reduction of the intermediate linear functions, from the way that $\textrm{Interleave}^{1}_\sigma$ (Figure 7) handles references to the global inputs of the program.

2. 2. $(\underline{\lambda}(z : \underline{\mathbb{R}}).\ \underline0)$ occurs as the backpropagator of a scalar constant r. Note that since this $\underline0$ is of type $\sigma$ , operationally it is equivalent to a tuple filled completely with zero scalars.

3. 3. $(\underline{\lambda}(z : \underline{\mathbb{R}}).\ d_1\ (\partial_1\textit{op}(x_1,\ldots,x_n)(z)) + \cdots + d_n\ (\partial_n\textit{op}(x_1,\ldots,x_n)(z)))$ for an $\textit{op} \in \mathrm{Op}_n$ where $d_1,\ldots,d_n$ are other linear backpropagators: these occur as the backpropagators generated for primitive operations.

Insight: Hence, we observe that a backpropagator $f_r$ paired with a scalar r will only ever call backpropagators $f_s$ that are paired with scalars s, such that r already has a dependency on s in the source program. In particular, $f_s$ must have been created (at runtime of the derivative program) before $f_r$ itself was created. Furthermore, $f_r$ is not the same function as $f_s$ because that would mean that r depends on itself in the source program. Therefore, if, at runtime, we define a partial order on backpropagators with the property that $f_r \geq f_s$ if r depends on s (and $f_r \gt f_s$ if they are not syntactically equal), we obtain that a called backpropagator is always strictly lower in the order than the backpropagator it was called from.

In practice, we achieve this by giving unique IDs, of some form, to backpropagators and defining a partial order on those IDs at runtime, effectively building a computation graph. This partial order tells us in which order to resolve backpropagators: we walk the order from top to bottom, starting from the maximal IDs and repeatedly resolving the predecessors in the order after we finish resolving a particular backpropagator. After all, any calls to other backpropagators that it produces in the returned ${\textrm{Staged}} c$ value will have lower IDs, and so cannot be functions that we have already resolved (i.e. called) before. And as promised, giving backpropagators IDs also solves the issue of using functions as keys in a $\textrm{Map}$ : we can use the ID as the $\textrm{Map}$ key, which is perfectly valid and efficient as long as the IDs are chosen to be of some type that can be linearly ordered to perform binary search (such as tuples of integers).

We have still been rather vague about how precisely to assign the IDs and define their partial order. In fact, there is some freedom in how to do that. For the time being, we will simply work with sequentially incrementing integer IDs with their linear order, which suffices for sequential programs. Concretely, we number backpropagators with incrementing integer IDs at runtime, at the time of their creation by a $\underline{\lambda}{}$ . We then resolve them from top to bottom, starting from the unique maximal ID. To support parallelism in Section 12, we will revisit this choice and work instead with pairs of integers (a combination of a job ID and a sequentially increasing ID within that job) with a partial order that encodes the fork-join parallelism structure of the source program. That choice of non-linear partial order allows us to reflect the parallelism present in the source program in a parallel reverse pass to compute derivatives. But because we can mostly separate the concerns of ID representation and differentiation, we will focus on simple, sequential integer IDs for now.

When we give backpropagators integer IDs, we can rewrite ${\textrm{Staged}} c$ and $\textrm{SCall}$ :

\begin{align*} &{\textrm{Staged}} c = (c, \textrm{Map}\ \mathbb Z\ (\underline{\mathbb{R}}\multimap {\textrm{Staged}} c, \underline{\mathbb{R}})) \\ &\begin{array}{@{\,}l@{\ }c@{\ }l@{}} \textrm{SCall} &:& (\mathbb Z, \underline{\mathbb{R}}\multimap {\textrm{Staged}} c) \rightarrow \underline{\mathbb{R}}\multimap {\textrm{Staged}} c \\ \textrm{SCall}\ (i, f)\ x &=& (\underline0, \{i \mapsto (f, x)\}) \end{array}\end{align*}

We call the second component of a ${\textrm{Staged}} c$ value, which has type $\textrm{Map}\ \mathbb Z\ (\underline{\mathbb{R}}\multimap {\textrm{Staged}} c, \underline{\mathbb{R}}))$ , the staging map, after its function to stage (linear) function calls.

The only thing that remains is to actually generate the IDs for the backpropagators at runtime. This we do using an ID generation monad (a state monad with a state of type $\mathbb Z$ to keep track of our integer IDs). The resulting new program transformation, modified from Figures 6 and 7, is shown in Figure 8.

Fig. 8. The monadically transformed code transformation (from Figures 4 to 5 plus ${\textrm{Staged}}$ operations), based on Figure 6. Grey parts are unchanged or simply monadically lifted.

New program transformation

In Figure 8, the term transformation now produces a term in the ID generation monad ( $\mathbb Z \rightarrow (-, \mathbb Z)$ ); therefore, all functions in the original program will also need to run in the same monad. This gives the second change in the type transformation (aside from $\textbf{D}^{2}_{c}[\mathbb R]$ , which now tags backpropagators with an ID): $\textbf{D}^{2}_{c}[\sigma \rightarrow \tau]$ now produces a monadic function type instead of a plain function type.

On the term level, notice that the backpropagator for primitive operations (in $\textbf{D}^{2}_{c}[\textit{op}(...)]$ ) now no longer calls $d_1,\ldots,d_n$ (the backpropagators of the arguments to the operation) directly, but instead registers the calls as pairs of function and argument in the ${\textrm{Staged}} c$ returned by the backpropagator. The $\cup$ in the definition of $(\mathbin{+_{\text{Staged}}})$ refers to map union including linear factoring; for example:

\[ \{i_1 \mapsto (f_1, a_1), i_2 \mapsto (f_2, a_2)\} \cup \{i_2 \mapsto (f_2, a_3)\} = \{i_1 \mapsto (f_1, a_1), i_2 \mapsto (f_2, a_2 + a_3)\} \]

Note that the transformation assigns consistent IDs to backpropagators: it will never occur that two staging maps have an entry with the same key (ID) but with a different function in the value. This invariant is quite essential in the algorithms in this paper.

In the wrapper, $\textrm{Interleave}^{2}$ is lifted into the monad and generates IDs for scalar backpropagators; $\textrm{Deinterleave}^{2}$ is essentially unchanged. The initial backpropagator provided to $\textrm{Interleave}^{2}$ in $\textrm{Wrap}^{2}$ , which was $(\underline{\lambda}{(z:\underline\sigma)}.\ z) : \underline\sigma \multimap \underline\sigma$ in Figure 7, has now become $\textrm{SCotan} : \underline\sigma \multimap {\textrm{Staged}} \underline\sigma$ , which injects a cotangent into a ${\textrm{Staged}} c$ object. $\textrm{Interleave}^{2}$ “splits” this function up into individual $\underline{\mathbb{R}}\multimap {\textrm{Staged}} \underline\sigma$ backpropagators for each of the individual scalars in $\sigma$ .

At the end of the wrapper, we apply the insight that we had earlier: by resolving (calling and eliminating) the backpropagators in the final ${\textrm{Staged}} c$ returned by the differentiated program in order from the highest ID to the lowest ID, we ensure that every backpropagator is called at most once.Footnote ²⁶ This is done in an additional function, $\textrm{SResolve}$ , which can be written as follows:

\begin{align*} \begin{array}{l} \textrm{SResolve} : {\textrm{Staged}} c \to c\\ \textrm{SResolve}\ (\textit{grad} : c, m : \textrm{Map}\ \mathbb Z\ (\underline{\mathbb{R}}\multimap {\textrm{Staged}} c, \underline{\mathbb{R}})) := \\ \quad\textbf{if}\ m\ \text{is empty}\ \begin{array}[t]{@{}l@{}} \textbf{then}\ \textit{grad} \\ \textbf{else}\ \begin{array}[t]{@{}l@{}} \textbf{let}\ i = \text{highest key in}\ m \\ \textbf{in}\ \textbf{let}\ (f, a) = \text{lookup}\ i\ \text{in}\ m \\ \textbf{in}\ \textbf{let}\ m' = \text{delete}\ i\ \text{from}\ m \\ \textbf{in}\ \textrm{SResolve}\ (f\ a \mathbin{+_{\text{Staged}}} (c', m')) \end{array} \end{array} \end{array} \hspace{0.5cm} \end{align*}

The three operations on m are standard logarithmic-complexity tree-map operations.

Complexity

Now that we have fixed (in Section 7) the first complexity problem identified in Section 6.2 (expensive monoid operations) and reduced the second (expensive input backpropagators) to a logarithmic overhead over the original program, we have reached the point where we satisfy the complexity requirement stated in Sections 3 and 6.2 apart from log-factors. More precisely, if we set $T = \textrm{cost}(P\ x)$ and $I = \textrm{size}(x)$ , then $\textrm{Wrap}^{}[P]$ computes the gradient of P at x not in the desired time $O(T + I)$ but instead in time $O\bigl(T \max(\log(T), \log(I)) + I \log(I)\bigr)$ .Footnote ³¹ If we accept logarithmic overhead, we could choose to stop here. However, if we wish to strictly conform to the required complexity, or if we desire to lose the non-negligible constant-factor overhead of dealing with a persistent tree-map, we need to make the input backpropagators and $\textrm{Map}$ operations in $\textrm{SResolve}$ constant-time; we do this using mutable arrays in Section 9.

Time complexity

We now satisfy all the requirements of the analysis in Section 6.2, and hence have the correct time complexity for reverse AD. In particular, let I denote the size of the input and T the runtime of the original program. Let $\textbf{D}^{4}_{}$ , $\textrm{Interleave}^{4}$ , $\textrm{Deinterleave}^{4}$ , etc. be the definitions that use arrays as described above (see Section 9.1 for details). Then we can observe the following:

• • The number of operations performed by $\textbf{D}^{4}_{c}[t]$ (with the improvements from Sections 8 and 9) is only a constant factor times the number of operations performed by t, and hence in O(T). This was already observed for $\textbf{D}^{2}_{c}[t]$ in Section 6.2, and still holds.

• • The number of backpropagators created while executing $\textbf{D}^{4}_{c}[t]$ is clearly also in O(T).

• • The number of operations performed in any one backpropagator is constant. This is new, and only true because id (replacing $0_{\text{Staged}}$ ), $(\circ)$ (replacing $(\mathbin{+_{\text{Staged}}})$ ), $\textrm{SCotan}$ (with a constant-time mutable array updater as argument) and $\textrm{SCall}$ are now all constant-time.

• • Hence, because every backpropagator is invoked at most once thanks to our staging, and because the overhead of $\textrm{SResolve}$ is constant per invoked backpropagator, the amount of work performed by calling the top-level input backpropagator is again in O(T).

• • Finally, the (non-constant-time) extra work performed in $\textrm{Wrap}^{4}$ is interleaving (O(I)), deinterleaving ( $O(\text{size of output})$ and hence $O(T + I)$ ), resolving (O(T)) and reconstructing the gradient from the scalars in the ${\textrm{Array}} \underline{\mathbb{R}}$ in ${\textrm{Staged}} c$ (O(I)); all this work is in $O(T + I)$ .

Hence, calling $\textrm{Wrap}^{4}[t]$ with an argument and calling its returned top-level derivative once takes time $O(T + I)$ , i.e. at most proportional to the runtime of calling t with the same argument, plus the size of the argument itself. This is indeed the correct time complexity for an efficient reverse AD algorithm, as discussed in Section 3.

Mutable arrays interface

We assume an interface to mutable arrays that is similar to that for IOVector in the Haskell vector library, cited earlier. In summary, we assume the following functions:

\[ \begin{array}{@{}l@{\ }c@{\ }l@{}} \texttt{alloc} &:& \mathbb Z \rightarrow \alpha \rightarrow \texttt{IO}\,\ ({\textrm{Array}} \alpha) \\ \texttt{get} &:& \mathbb Z \rightarrow {\textrm{Array}} \alpha \rightarrow \texttt{IO}\,\ \alpha \\ \texttt{modify} &:& \mathbb Z \rightarrow (\alpha \rightarrow \alpha) \rightarrow {\textrm{Array}} \alpha \rightarrow \texttt{IO}\,\ () \\ \texttt{freeze} &:& {\textrm{Array}} \alpha \rightarrow \texttt{IO}\,\ ({\textrm{IArray}} \alpha) \\ (@) &:& {\textrm{IArray}} \alpha \rightarrow \mathbb Z \rightarrow \alpha \end{array}\]

where ${\textrm{IArray}}$ is an immutable array type. $\texttt{alloc}$ and $\texttt{freeze}$ are linear in the length of the array; $\texttt{get}$ , $\texttt{modify}$ and $(@)$ are constant-time (in addition to calling the function once, of course, for $\texttt{modify}$ ).

Let us see how the code transformation changes with mutable arrays.

Code transformation

Using the new definition of ${\textrm{Staged}}$ , we change the code transformation once more, this time from Figure 9 to the version given in Figure 10. The transformation on types and terms simply sees the type of scalar backpropagators change to use effectful updating instead of functional updating, so they do not materially change: we just give yet another interpretation of $0_{\text{Staged}}$ and $(\mathbin{+_{\text{Staged}}})$ , using the monoid structure on ${\textrm{Staged}} c \rightarrow \texttt{IO}\,\ ()$ defined above in terms of $(\mathbin{>\!\!>})$ . However, some important changes occur in the ${\textrm{Staged}} c$ interface and the wrapper. Let us first look at the algorithm from the top, by starting with $\textrm{Wrap}^{4}$ ; after understanding the high-level idea, we explain how the other components work.

Fig. 9. The Cayley-transformed code transformation, based on Figure 8. Grey parts are unchanged.

Fig. 10. Code transformation plus wrapper using mutable arrays, modified from Figure 9. Grey parts are unchanged.

In basis, $\textrm{Wrap}^{4}$ does the same as $\textrm{Wrap}^{3}$ from Figure 9: interleave injector backpropagators with the input of type $\sigma$ , execute the transformed function body using the interleaved input, and then deinterleave the result. However, because we (since Section 8) represent the final cotangent not directly as a value of type $\sigma$ in a ${\textrm{Staged}} \sigma$ but instead as an array of only the embedded scalars ( ${\textrm{Array}} \underline{\mathbb{R}}$ ), some more work needs to be done.

First, $\textrm{Interleave}^{4}_\sigma$ (monadically, in the ID generation monad that we are writing out explicitly) produces, in addition to the interleaved input, also a reconstruction functionFootnote ³⁷ of type ${\textrm{IArray}} \underline{\mathbb{R}} \rightarrow \sigma$ . This rebuilder takes an array with precisely as many scalars as were in the input, and produces a value of type $\sigma$ with the structure (and discrete-typed values) of the input, but the scalars from the array. The mapping between locations in $\sigma$ and indices in the array is the same as the numbering performed by $\textrm{Interleave}^{4}$ .

Having x, $\textit{rebuild}$ and i (the next available ID), we execute the transformed term $\textbf{D}^{4}_{\sigma}[t]$ monadically (with x in scope), resulting in an output $y' : \textbf{D}^{4}_{\sigma}[\tau]$ . This output we deinterleave to $y : \tau$ and $d : \tau \multimap ({\textrm{Staged}} \sigma \rightarrow \texttt{IO}\,\ ())$ .

The final result then consists of the regular function result (y) as well as the top-level derivative function of type $\tau \multimap \sigma$ . In the derivative function, we allocate two arrays to initialise an empty ${\textrm{Staged}} \sigma$ (note that the given sizes are indeed precisely large enough), and apply d to the incoming $\tau$ cotangent. This gives us an updater function that (because of how $\textrm{Deinterleave}^{4}$ works) calls the top-level backpropagators contained in y’ in the ${\textrm{Staged}}$ arrays, and we apply this function to the just-allocated ${\textrm{Staged}}$ object. Then we use the new $\textrm{SResolve}$ to propagate the cotangent contributions backwards, by invoking each backpropagator in turn in descending order of IDs. Like before in the Cayley-transformed version of our AD technique, those backpropagators update the state (now mutably) to record their own contributions to (i.e. invocations of) other backpropagators. At the end of $\textrm{SResolve}$ , the backpropagator staging array is dropped and the cotangent collection array is frozen and returned as an ${\textrm{IArray}} \underline{\mathbb{R}}$ (corresponding to the c value in a ${\textrm{Staged}} c$ for the Cayley-transformed version in Figure 9).

This whole derivative computation is, in the end, pure (in that it is deterministic and has no side-effects), so we can safely evaluate the IO using unsafePerformIO to get a pure ${\textrm{IArray}} \underline{\mathbb{R}}$ , which contains the scalar cotangents that $\textit{rebuild}$ from $\textrm{Interleave}^{4}$ needs to put in the correct locations in the input, thus constructing the final gradient.

Implementation of the components

Having discussed the high-level sequence of operations, let us briefly discuss the implementation of the ${\textrm{Staged}} c$ interface and (de)interleaving. In $\textrm{Interleave}^{4}$ , instead of passing structure information down in the form of a setter ( $(\tau \rightarrow \tau) \multimap ({\textrm{Staged}} c \rightarrow {\textrm{Staged}} c)$ ) like we did in $\textrm{Interleave}^{3}$ in Figure 9, we build structure information up in the form of a getter ( ${\textrm{IArray}} \underline{\mathbb{R}} \rightarrow \tau$ ). This results in a somewhat more compact presentation, but in some sense the same information is still communicated.

The program text of $\textrm{Deinterleave}^{4}$ is again unchanged, because it is agnostic about the codomain of the backpropagators, as long as it is a monoid, which it remains.

On the ${\textrm{Staged}}$ interface, the transition to mutable arrays had a significant effect. The $0_{\text{Staged}}$ created by $\textrm{Wrap}^{3}$ in Section 7 is now essentially in $\textrm{SAlloc}$ , which uses $\texttt{alloc}$ to allocate a zero-filled cotangent collection array of size $i_{\text{inp}}$ , and the backpropagator staging array of size $i_{\text{out}}$ filled with zero-backpropagators with an accumulated argument of zeros.

$\textrm{SCall}$ has essentially the same type, but its implementation differs because it now performs a constant-time mutable update on the backpropagator staging array instead of a logarithmic-complexity immutable $\textrm{Map}$ update. Note that, unlike in Section 7, there is no special case if i is not yet in the array, because unused positions are already filled with zeros.

$\textrm{SOneHot}$ takes the place of $\textrm{SCotan}$ , with the difference that we have specialised it using the knowledge that all relevant $c \rightarrow c$ functions add a particular scalar to a particular index in the input, and that these functions can hence be defunctionalised to a pair $(\mathbb Z, \underline{\mathbb{R}})$ . The monoid-linearity here is in the real scalar, as it was before, hence the placement of the $\multimap $ -arrow.

Finally, $\textrm{SResolve}$ takes an additional $\mathbb Z$ argument that should contain the output ID of $\textbf{D}^{4}_{c}[t]$ , i.e. one more than the largest ID generated. $\textit{loop}$ then does what the original $\textrm{SResolve}$ did directly, iterating over all IDs in descending order and applying the state updaters in the backpropagator staging array one-by-one to the state. After the loop is complete, we freeze and return just the cotangent collection array, because we have no need for the staging array any more. This frozen collection array will then be used to build the final gradient in $\textrm{Wrap}^{4}$ .

(Mutual) recursion

For example, we can allow recursive functions in our source language by adding recursive let-bindings with the following typing rule:

\[ \frac{\Gamma, f : \sigma \rightarrow \tau, x : \sigma \vdash s : \tau \qquad \Gamma, f : \sigma \rightarrow \tau \vdash t : \rho} {\Gamma \vdash \textbf{letrec}\ f\ x = s\ \textbf{in}\ f\ t : \rho}\]

The code transformation $\textbf{D}^{i}_{}$ for all i then treats $\textbf{letrec}$ exactly the same as $\textbf{let}$ —note that the only syntactic difference between $\textbf{letrec}$ and $\textbf{let}$ is the scoping of f—and the algorithm remains both correct and efficient. Note that due to the assumed call-by-value semantics, recursive non-function definitions would make little sense.

Recursion introduces the possibility of non-termination because the reverse pass is nothing more than a loop over the primitive scalar operations performed in the forward execution, the derivative program terminates exactly if the original program terminates (on a machine with sufficient memory).

Coproducts

To support dynamic control flow (necessary to make recursion useful), we can easily add coproducts to the source language. First add coproducts to the syntax for types () both in the source language and in the target language and add constructors and eliminators to all term languages (both linear and non-linear):

where x is in scope in $t_1$ and y is in scope in $t_2$ . Then the type and code transformations extend in the unique structure-preserving manner:

To create an interesting interaction between control flow and differentiation, we can add a construct ‘ $\textrm{sign}$ ’ with the unsurprising typing rule

\[ \frac{\Gamma \vdash t : \mathbb R} {\Gamma \vdash \textrm{sign}(t) : {\mathrm{Bool}}}\]

where ${\mathrm{Bool}} = () \sqcup ()$ , which allows us to perform a case distinction on the sign of a real number. For differentiation of this construct, it suffices to define $\textbf{D}^{1}_{c}[\textrm{sign}(t)] = \textrm{sign}(\textrm{fst}(\textbf{D}^{1}_{c}[t]))$ .

If one accepts losing some of the structure-preserving nature of the transformation, it is possible to prevent redundant differentiation of t in $\textbf{D}^{1}_{c}[\textrm{sign}(t)]$ by making clever substitutions in t’s free variables, converting back from dual-number form to the plain data representation. The idea is to define functions $\varphi_\tau : \textbf{D}^{1}_{c}[\tau] \rightarrow \tau$ and $\psi_\tau : \tau \rightarrow \textbf{D}^{1}_{c}[\tau]$ by induction on $\tau$ , where $\varphi_\tau$ projects out the value from a dual number and $\psi_\tau$ injects scalars into a dual number as constants (i.e. with a zero backpropagator). $\varphi_\tau$ and $\psi_\tau$ are mutually recursive at function types: e.g. $\varphi_{\sigma \rightarrow \tau}\ f = \varphi_\tau \circ f \circ \psi_\sigma$ . Then one can define a non-differentiating transformation $\textbf{D}^{\text{plain}}_{c}$ on terms that uses $\varphi$ to convert free variables to plain values, and otherwise keeps all code plain. We then get $\textbf{D}^{1}_{c}[\textrm{sign}(t)] = \textrm{sign}(\textbf{D}^{\text{plain}}_{c}[t])$ .

When moving to $\textbf{D}^{4}_{c}$ , the type transformation for coproducts stays unchanged, and the term definitions change only by transitioning to monadic code in $\textbf{D}^{2}_{c}$ . Lifting a computation to monadic code is a well-understood process. The corresponding cases in $\textrm{Interleave}^{}$ and $\textrm{Deinterleave}^{}$ are the only reasonable definitions that type-check.

The introduction of dynamic control flow complicates the correctness story for any AD algorithm. The approach presented here has the usual behaviour: derivatives are correct in the interior of domains leading to a particular execution path (if ‘ $\textrm{sign}$ ’ is the only “continuous conditional”, this is for inputs where none of the ‘ $\textrm{sign}$ ’ operations receive zero as their arguments), but may be unexpected at the points of branching. For discussion see e.g. (Hückelheim et al. Reference Hückelheim, Menon, Moses, Christianson, Hovland and Hascoët2023, § 3.3); for proofs see e.g. (Lucatelli Nunes and Vákár Reference Lucatelli Nunes and Vákár2024, § 11) or Mazza and Pagani (Reference Mazza and Pagani2021).

Polymorphic and (mutually) recursive types

In Haskell one can define (mutually) recursive data types, e.g. as follows:

\begin{align*} &\textbf{data}\ T_1\ \alpha = C_1\ \alpha\ (T_2\ \alpha) \mid C_2\ \mathbb R \\ &\textbf{data}\ T_2\ \alpha = C_3\ \mathbb Z\ (T_1\ \alpha)\ (T_2\ \alpha)\end{align*}

If the user has defined some data types, then we can allow these data types in the code transformation. We generate new data type declarations that simply apply $\textbf{D}^{1}_{c}[-]$ to all parameter types of all constructors:

\begin{align*} &\textbf{data}\ DT_1\ \alpha = DC_1\ \alpha\ (DT_2\ \alpha) \mid DC_2\ (\mathbb R,\underline{\mathbb{R}}\multimap c) \\ &\textbf{data}\ DT_2\ \alpha = DC_3\ \mathbb Z\ (DT_1\ \alpha)\ (DT_2\ \alpha)\end{align*}

and we add one rule for each data type that simply maps:

\begin{gather*} \textbf{D}^{1}_{c}[T_1\ \tau] = DT_1\ \textbf{D}^{1}_{c}[\tau] \qquad \textbf{D}^{1}_{c}[T_2\ \tau] = DT_2\ \textbf{D}^{1}_{c}[\tau]\end{gather*}

Furthermore, for plain type variables, we set $\textbf{D}^{1}_{c}[\alpha] = \alpha$ .Footnote ³⁹

The code transformation on terms is completely analogous to a combination of coproducts (given above in this section, where we take care to match up constructors as one would expect: $C_i$ gets sent to $DC_i$ ) and products (given already in Figure 6). The wrapper also changes analogously: $\textrm{Interleave}^{}$ and $\textrm{Deinterleave}^{}$ get clauses for $\textrm{Interleave}^{}_{(T_i\,\tau)}$ and $\textrm{Deinterleave}^{}_{(T_i\,\tau)}$ .

Finally, we note that with the mentioned additional rule that $\textbf{D}^{1}_{c}[\alpha] = \alpha$ , polymorphic functions can also be differentiated transparently, similarly to how the above handles polymorphic data types.

Differentiation

We keep the differentiation rules for all existing language constructs the same, except for changing the type of IDs to $\textrm{CID}$ and using the monad $\operatorname*{\mathcal{M}}$ instead of doing manual state passing of the next ID to generate. A representative rule showing how this looks is: (compare Figure 8)

\[ \textbf{D}^{5}_{c}[(s, t)] = \textbf{D}^{5}_{c}[s] \mathbin{>\!\!>\!\!=} \lambda x.\ \textbf{D}^{5}_{c}[t] \mathbin{>\!\!>\!\!=} \lambda y.\ \mathbf{return}\ (x, y)\]

Because we now generate fresh compound IDs rather than plain integer IDs when executing primitive operations, we change $\mathbb Z$ to $\textrm{CID}$ in $\textbf{D}^{4}_{c}[\mathbb R]$ :Footnote ⁴¹

The new rule for the parallel pairing construct is simply:

\[ \textbf{D}^{5}_{c}[t\star s] = \textbf{D}^{5}_{c}[t] \star' \textbf{D}^{5}_{c}[s]\]

This is the only place where we use the operation $(\star')$ , and thus the only place in the forward pass where we directly use the extended functionality of our monad $\operatorname*{\mathcal{M}}$ .

Duality

The usual mantra in reverse-mode AD is that “sharing in the primal becomes addition in the dual”. When parallelism comes into play, not only do we have this duality in the data flow graph, we also get an interesting duality in the control-flow graph: $\textrm{SResolve}$ forks where the primal joined and joins where the primal forked. Perhaps we can add a mantra for task-parallel reverse AD: “forks in the primal become joins in the dual”.

Benchmarks

To check that our implementation has reasonable performance in practice, we benchmark (in bench/Main.hs) against Kmett’s Haskell ad library (Kmett and Contributors, 2021) (version 4.5.6) on a few basic functions.Footnote ⁴³ These functions are the following (abbreviating Double as “D”):

• • A single scalar multiplication of type (D, D) -> D;

• • Dot product of type ([D], [D]) -> D;

• • Matrix–vector multiplication, then sum: of type ([[D]], [D]) -> D;

• • The rotate_vec_by_quat example from Krawiec et al. (Reference Krawiec, Jones, Krishnaswami, Ellis, Eisenberg and Fitzgibbon2022) of type (Vec3 D, Quaternion D) -> Vec3 D, with data Vec3 s = Vec3 s s s and data Quaternion s = Quaternion s s s s;

• • A simple, dense neural network with 2 hidden layers, ReLU activations and (safe) softmax output processing. The result vector is summed to make the output a single scalar. The actual Haskell function is generic in the number of layers and has type ([([[D]], [D])], [D]) -> D: the first list contains a matrix and a bias vector for each hidden layer, and the second tuple component is the input vector. In the benchmark, the input has length 50, and the two hidden layers have size 100 and 50.

• • Simulation of 4 particles in a simple force field with friction for 1000 time steps; this example has type [((D, D), (D, D))] -> D. The input is a list (of length 4) of initial positions and velocities; the output is $\sum_{(x,y)} x \cdot y$ ranging over the 4 result positions p, to ensure that the computation has a single scalar as output. The four particles are simulated in parallel using the $(\star)$ combinator from Section 12.

The fourth test case, rotate_vec_by_quat, has a non-trivial return type; the benchmark executes its reverse pass three times (“3” being the number of scalars in the function output) to get the full Jacobian. The fifth test case, “particles”, is run on 1, 2 and 4 threads, where the ideal result would be perfect scaling due to the four particles being independent.

Table 1. Benchmark results of Section 12 + Sections 10.1 and 10.2 versus ad-4.5.6. The “TH” and “ad” columns indicate runtimes on one machine for our implementation and the ad library, respectively. The last column shows the ratio between the previous two columns. We give the size of the largest side of criterion’s 95% bootstrapping confidence interval, rounded to 2 decimal digits. Setup: GHC 9.6.6 on Linux, Intel i9-10900K CPU, with Intel Turbo Boost disabled (i.e. running at a consistent 3.7 GHz).

The benchmark results are summarised in Figure 1 and shown in detail (including evidence of the desired linear complexity scaling) in Appendix A. The benchmarks are measured using the criterion Footnote ⁴⁴ library. To get statistically significant results, we measure how the timings scale with increasing n:

• • Scalar multiplication, rotate_vec_by_quat, the neural network and the particle simulation are simply differentiated n times;

• • Dot product is performed on lists of length n;

• • Matrix multiplication is done for a matrix and vector of size $\sqrt n$ , to get linear scaling in n.

The reported time is the deduced actual runtime for $n = 1$ .

By the results in Figure 1, we see that on these simple benchmark programs, our implementation is faster than the existing ad library. While this is encouraging, it is not overly surprising: because our algorithm is implemented as a compile-time code transformation, the compiler is able to optimise the derivative code somewhat before it gets executed.

Of course, performance results may well be different for other programs, and AD implementations that have native support for array operations can handle some of these programs much more efficiently. Furthermore, there are various implementation choices for the code transformation that may seem relatively innocuous but have a large impact on performance (we give some more details below in Section 13.1).

For these reasons, our goal here is merely to substantiate the claim that the implementation exhibits constant-factor performance in the right ballpark (in addition to it having the right asymptotic complexity, as we have argued). Nevertheless, our benchmarks include key components of many AD applications, and seeing that we have not at all special-cased their implementation (the implementation knows only primitive arithmetic operations such as scalar addition, multiplication, etc.), we believe that they suffice to demonstrate our limited claim.

Garbage collection

The graph of backpropagators (a normal data structure, “ $\textrm{Contrib}$ ”, since Section 10.2) is a big data structure of size proportional to the number of scalar operations performed in the source program. While this data structure grows during the forward pass, the nursery (zeroth generation) of GHC’s generational garbage collector (GC) repeatedly fills up, triggering garbage collection of the heap. Because a GC pass takes time on the order of the amount of live data on the heap, these passes end up very expensive: a naive taping reverse AD algorithm becomes quadratic in the source program runtime. Using a GHC runtime system flag (e.g. -A500m) to set the nursery size of GHC’s GC sufficiently large to prevent the GC from running during a benchmark (criterion runs the GC explicitly between each benchmark invocation), timings on some benchmarks above decrease significantly: on the largest benchmark (particles), this can save 10% off ad’s runtime and 25% off our runtime (though precise timings vary). The times reported in Figure 1 are with GHC’s default GC settings.

While this is technically a complexity problem in our implementation, we gloss over this because it is not fundamental to the algorithm: the backpropagator graph does not contain cycles, so it could be tracked with reference-counting GC to immediately eliminate the quadratic blowup. Using a custom, manual allocator, one could also eliminate all tracking of the liveness of the graph because we know from the structure of the algorithm exactly when we can free a node in the graph (namely when we have visited it in the reverse pass). Our reference implementation does not do these things to be able to focus on the workings of the algorithm.

Laziness

Because data types are lazy by default in Haskell, a naive encoding of the $\textrm{Contrib}$ data type from Section 10.2 would make the whole graph lazily evaluated (because it is never demanded during the forward pass). This results in a significant constant-factor overhead (more than $2\times$ ), and furthermore means that part of the work of the forward pass happens when the reverse pass first touches the top-level backpropagator; this work then happens sequentially, even if the forward pass was meant to be parallel. To achieve proper parallel scaling, it was necessary to make the $\textrm{Contrib}$ graph strict, and furthermore to make the $\textbf{D}^{}_{c}[\mathbb R] = (\mathbb R, (\textrm{CID}, \textrm{Contrib}))$ triples strict, to ensure that the graph is fully evaluated as it is being constructed, not when it is demanded in the reverse pass.

Using some well-chosen -# UNPACK #- pragmas on some of these fields also had a significant positive effect on performance.

Thread pool

In Section 12, we used an abstract inParallel operation for running two jobs in parallel, assuming some underlying thread pool for efficient evaluation of the resulting parallel job graph. In a standard thread pool implementation, spawning tasks from within tasks can result in deadlocks. Because nested tasks are essential to our model of task parallelism, the implementer has to take care to further augment $\operatorname*{\mathcal{M}}$ from Section 12 to be a continuation monad as well: this allows the continuation of the inParallel operation to be captured and scheduled separately in the thread pool, so that thread pool jobs are indeed individual jobs as defined in Section 12.

The GHC runtime system has a thread scheduler that should handle this completely transparently, but in our tests it was not eager enough in assigning virtual Haskell threads to actual separate kernel threads, resulting in a sequential benchmark. This motivated a (small) custom thread pool implementation that sufficed for our benchmarks, but has a significant amount of overhead that can be optimised with more engineering effort.

Imperfect scaling

Despite efforts to the contrary, it is evident from Figure 1 that even on an embarrassingly parallel task like “particles”, the implementation does not scale perfectly. From inspection of more granular timing of our implementation, we suspect that this is a combination of thread pool overhead and the fact that the forward pass simply allocates too quickly, exhausting the memory bandwidth of our system when run sufficiently parallel. However, more research is needed here to uncover the true bottlenecks.

Distributive law

Under the isomorphism $\underline{\mathbb{R}}\multimap c\cong c$ , the type ${\textrm{Staged}} c$ can be thought of as a type ${\textrm{Expr}} c$ of ASTs of expressions (with sharing) of type c.Footnote ⁴⁵ The linear factoring rule $f\ x+ f\ y\leadsto f\ (x+y)$ for a linear function $f:\underline{\mathbb{R}}\multimap c$ that rescales a vector $v:c$ with a scalar then corresponds to the distributive law $v\cdot x+ v\cdot y\leadsto v\cdot (x+y)$ . This highlights the relationship between our work and that of Shaikhha et al. (Reference Shaikhha, Fitzgibbon, Vytiniotis and Jones2019), who try to (statically) optimise vectorised forward AD to reverse AD using precisely this distributive law. A key distinction is that we apply this law (in the form of the linear factoring rule) at runtime rather than compile time, allowing us to always achieve the complexity of reverse AD, rather than merely on some specific programs with straightforward control and data flow. The price we pay for this generality is a runtime overhead, similar to the usual overhead of taping.

CHAD and category theory inspired AD

Rather than interleaving backpropagators by pairing them with scalars in a type, we can also try to directly implement reverse AD as a structurally recursive code transformation that does not need a (de)interleaving wrapper: $\mathcal{R}_2$ from Section 4. This is the approach taken by Elliott (Reference Elliott2018). It pairs vectors (and values of other composite types) with a single composite backpropagator, rather than decomposing to the point where each scalar is paired with a mini-backpropagator like in our dual-numbers approach. The resulting algorithm is extended to source languages with function types in Vákár (Reference Vákár2021), Vákár and Smeding (Reference Vákár and Smeding2022) and Vytiniotis et al. (Reference Vytiniotis, Belov, Wei, Plotkin and Abadi2019) and to sum and (co)inductive types in Nunes and Vákár (Reference Nunes and Vákár2023). Smeding and Vákár (Reference Smeding and Vákár2024) give a mechanised proof that the resulting algorithm attains the correct computational complexity, after a series of optimisations that are similar to the ones considered in this paper. Like our dual-numbers reverse AD approach, the algorithm arises as a canonical structure-preserving functor on the syntax of a programming language. However, due to a different choice in target category (a Grothendieck construction of a linear $\lambda$ -calculus for CHAD rather than the syntax of a plain $\lambda$ -calculus for dual-numbers AD), the resulting algorithm looks very different.

Taping-like methods and non-local control flow

Another family of approaches to AD recently taken by the PL community contains those that rely on forms of non-local control flow such as delimited continuations (Wang et al., Reference Wang, Zheng, Decker, Wu, Essertel and Rompf2019) or effect handlers (Sigal, Reference Sigal2021; de Vilhena and Pottier, Reference de Vilhena and Pottier2023). These techniques are different in the sense that they generate code that is not purely functional. This use of non-local control flow makes it possible to achieve an efficient implementation of reverse AD that looks strikingly simple compared to alternative approaches. Where the CHAD approaches and our dual-numbers reverse AD approach both have to manually invert control flow at compile time by passing around continuations, possibly combined with smart staging of those continuations like in this paper, this inversion of control can be deferred to run time by clever use of delimited control operators or effect handlers. de Vilhena and Pottier (Reference de Vilhena and Pottier2023) give a mechanised proof that the resulting (rather subtle) code transformation is correct.

Operationally, however, these effect-handler-based techniques are essentially equivalent to taping: the mutable cells for cotangent accumulation scope over the full remainder of the program. In this sense, they are operationally similar to the algorithm of Section 10.3, which is also essentially taping. The AD algorithm in Dex (Paszke et al., Reference Paszke, Johnson, Duvenaud, Vytiniotis, Radul, Johnson, Ragan-Kelley and Maclaurin2021a) also achieves something like taping in a functional style, by conversion to an A-normal form.

AD of parallel code

Bischof et al. (Reference Bischof, Griewank and Juedes1991); Bischof (Reference Bischof1991) present some of the first work in parallel AD. Rather than starting with a source program (with potential dynamic control flow) that has (fork-join) parallelism like we do and mirroring that in the reverse pass of the algorithm, they focus on code without dynamic control flow and analyse the dependency structure of the reverse pass at compile time to automatically parallelise it. This approach is developed further by Bucker et al. (Reference Bucker, Lang, Rasch, Bischof and Mey2002).

Building on the classic imperative AD tool Tapenade (Hascoët and Pascual, Reference Hascoët and Pascual2013), Hückelheim and Hascoët (Reference Hückelheim and Hascoët2022) discuss a method for differentiating parallel for-loops (with shared memory) in a parallelism-preserving way. Industrial machine learning frameworks such as TensorFlow (Abadi et al., Reference Abadi, Barham, Chen, Chen, Davis, Dean, Devin, Ghemawat, Irving, Isard, Kudlur, Levenberg, Monga, Moore, Murray, Steiner, Tucker, Vasudevan, Warden, Wicke, Yu and Zheng2016), JAX (Bradbury et al., Reference Bradbury, Frostig, Hawkins, Johnson, Leary, Maclaurin, Necula, Paszke, VanderPlas, Wanderman-Milne and Zhang2018) and PyTorch (Paszke et al., Reference Paszke, Gross, Chintala, Chanan, Yang, DeVito, Lin, Desmaison, Antiga and Lerer2017) focus on data parallelism through parallel array operations—typically first-order ones.

Kaler et al. (Reference Kaler, Schardl, Xie, Leiserson, Chen, Pareja and Kollias2021) focus on AD of programs with similar forms of fork-join parallelism as we consider in this paper. Their implementation builds on the Adept C++ library, which implements an AD algorithm that is very different from the one discussed in this paper. A unique feature of their work is that they give a formal analysis of the complexity of the technique and give bounds for both the work and the span of the resulting derivative code; it is possible that ideas from the cited work can be used to improve and bound the span of our parallel algorithm.

Other recent work has focussed on developing data-parallel derivatives for data-parallel array programs (Paszke et al., Reference Paszke, Johnson, Duvenaud, Vytiniotis, Radul, Johnson, Ragan-Kelley and Maclaurin2021a; b; Schenck et al., Reference Schenck, Rønning, Henriksen and Oancea2022). These methods are orthogonal to the ideas focussed on task parallelism that we present in the present paper.

In recent work, we have shown that an optimised version of CHAD seems to preserve task and data parallelism (Smeding and Vákár, Reference Smeding and Vákár2024). We are working on giving formal guarantees about the extent (work and span) to which different forms of parallelism is preserved. Compared to the dual numbers reverse AD method of the present paper, CHAD seems to accommodate higher order array combinators more easily, including their data parallel versions. This makes us hopeful for the use of CHAD as a candidate algorithm for parallel AD.