Chapter 8Automatic differentiation

In Chapter 2, we reviewed the fundamentals of differentiation and stressed the importance of two linear maps: the Jacobian-vector product (JVP) and its adjoint, the vector-Jacobian product (VJP). In this chapter, we review forward-mode and reverse-mode autodiff using these two linear maps. We start with computation chains and then generalize to feedforward networks and general computation graphs. We also review checkpointing, reversible layers and randomized estimators.

8.1 Computation chains

To begin with, consider a computation chain (Section 4.1.1) representing a function $f \colon \cS_0 \to \cS_K$ expressed as a sequence of compositions $f \coloneqq f_K \circ \dots \circ f_1$, where $f_k \colon \cS_{k-1} \to \cS_k$. The computation of $f$ can be unrolled into a sequence of operations $$\begin{align*} \s_0 &\in \cS_0 \\ \s_1 &\coloneqq f_1(\s_0) \in \cS_1 \\ & \hspace{6pt} \vdots \\ \s_K &\coloneqq f_K(\s_{K-1}) \in \cS_K \\ f(\s_0) &\coloneqq \s_K. \tag{8.1}\end{align*}$$ Our goal is to compute the variations of $f$ around a given input $\s_0$. In a feedforward network, this amounts to estimating the influence of a given input $\s_0$ for fixed parameters (we will see how to estimate the variations w.r.t. parameters $\w$ in the sequel).

Jacobian matrix.

We first consider the computation of the full Jacobian $\jac f(\s_0)$, seen as a matrix, as the notation $\bm{\partial}$ indicates. Following Proposition 2.2, we have $$\jac f(\s_0) = \jac f_K(\s_{K-1}) \ldots \jac f_1(\s_0), \tag{8.2}$$ where $\jac f_k(\s_{k-1})$ are the Jacobians of the intermediate functions computed at $\s_0, \ldots, \s_K$, as defined in Eq. (8.1). We may also want to compute the transpose of the Jacobian $$\jac f(\s_0)^\top = \jac f_1(\s_0)^\top \ldots \jac f_K(\s_{K-1})^\top.$$ In both cases, the main drawback of this approach is computational: computing the full $\jac f(\s_0)$ requires materializing the intermediate Jacobians in memory and performing matrix-matrix multiplications. However, in practice, computing the full Jacobian is rarely needed. Indeed, oftentimes, we only need to right-multiply or left-multiply with $\jac f(\s_0)$. This gives rise to forward-mode and reverse-mode autodiff, respectively.

8.1.1 Forward-mode

We now interpret the Jacobian $\partial f(\s_0)$ as a linear map, as the non-bold $\partial$ indicates. Following Proposition 2.6, $\partial f(\s_0)$ is the composition of the intermediate linear maps, $$\partial f(\s_0) = \partial f_K(\s_{K-1}) \circ \ldots \circ \partial f_1(\s_0).$$ Evaluating $\partial f(\s_0)$ on an input direction $\v \in \cS_0$ can be decomposed, like the function Eq. (8.1) itself, into intermediate computations $$\begin{align*} \t_0 & \coloneqq \v \\ \t_1 & \coloneqq \partial f_1(\s_0)[\t_0] \\ & \hspace{6pt} \vdots\\ \t_K & \coloneqq \partial f_K(\s_{K-1})[\t_{K-1}]\\ \partial f(\s_0)[\v] &\coloneqq \t_K. \end{align*}$$ Each intermediate $\partial f_k(\s_{k-1})[\t_{k-1}]$ amounts to a Jacobian-vector product (JVP) and can be performed in a forward manner, alongside the computation of the intermediate states $\s_k$. This can also be seen as multiplying the matrix defined in Eq. (8.2) with a vector, from right to left. This is illustrated in Fig. 8.1 and the procedure is summarized in Algorithm 8.1.

Computational complexity.

The JVP follows exactly the computations of $f$, with an additional variable $\t_k$ being propagated. If we consider that computing $\partial f_k$ is roughly as costly as computing $f_k$, then computing a JVP has roughly twice the computational cost of $f$. See Section 8.3.3 for a more general and more formal statement.

Memory usage.

The memory usage of a program at a given evaluation step is the number of variables that need to be stored in memory to ensure the execution of all remaining steps. The memory cost of a program is then the maximal memory usage over all evaluation steps. For our purposes, we analyze the memory usage and memory cost by examining the given program. Formal definitions of operations on memory such as read, write, delete and associated memory costs are presented by Griewank and Walther (2008, 4).

For example, to execute the chain $f = f_K \circ \dots \circ f_1$, at each step $k$, we only need to have access to $\s_{k-1}$ to execute the rest of the program. As we compute $\s_k$, we can delete $\s_{k-1}$ from memory and replace it by $\s_k$. Therefore, the memory cost associated with the evaluation of $f$ is just the maximal dimension of the $\s_k$ variables.

For forward mode autodiff, as we follow the computations of $f$, at each step $k$, we only need to have access to $\s_{k-1}$ and $\t_{k-1}$ to execute the rest of the program. The memory used by $\s_{k-1}$ and $\t_{k-1}$ can directly be used for $\s_k, \t_k$ once they are computed. The memory usage associated with the JVP is summarized in Fig. 8.2. Overall the memory cost of the JVP is then exactly twice the memory cost of the function itself.

8.1.2 Reverse-mode

In machine learning, most functions whose gradient we need to compute take the form $\ell \circ f$, where $\ell$ is a scalar-valued loss function and $f$ is a network. As seen in Proposition 2.3, the gradient takes the form $$\nabla (\ell \circ f)(\s_0) = \partial f(\s_0)^*[\nabla \ell(f(\s_0))].$$ This motivates the need for applying the adjoint $\partial f(\s_0)^*$ to $\nabla \ell(f(\s_0)) \in \cS_K$ and more generally to any output direction $\u \in \cS_K$. From Proposition 2.7, we have $$\partial f(\s_0)^* = \partial f_1(\s_0)^* \circ \ldots \circ \partial f_K(\s_{K-1})^*.$$ Evaluating $\partial f(\s_0)^*$ on an output direction $\u \in \cS_K$ is decomposed as $$\begin{align*} \r_K & \coloneqq \u \\ \r_{K-1} &\coloneqq \partial f_K(\s_{K-1})^*[\r_K] \\ & \hspace{6pt} \vdots \\ \r_0 &\coloneqq \partial f_1(\s_0)^*[\r_1] \\ \partial f(\s_0)^*[\u] &\coloneqq \r_0. \end{align*}$$ Each intermediate adjoint $\partial f_k(\s_{k-1})^*$ amounts to a vector-Jacobian product (VJP). The key difference with the forward mode is that the procedure runs backward through the chain, hence the name reverse mode autodiff. This can also be seen as multiplying Eq. (8.2) from left to right. The procedure is illustrated in Fig. 8.3 and summarized in Algorithm 8.2.

Computational complexity.

In terms of number of operations, the VJP simply passes two times through the chain, once forward, then backward. If we consider the intermediate VJPs to be roughly as costly as the intermediate functions themselves, the VJP amounts to just twice the cost of the original function, just as the JVP. See Section 8.3.3 for a more generic and formal statement.

Memory usage.

Recall that the memory usage of a program at a given evaluation step is the number of variables that need to be stored in memory to ensure the execution of the remaining steps. If we inspect Algorithm 8.3, to execute all backward steps, that is the loop in line 4, we need to have access to all the intermediate inputs $\s_0, \ldots, \s_{K-1}$. Therefore, the memory cost of reverse-mode autodiff is proportional to the length of the chain $K$. Fig. 8.4 illustrates the memory usage during reverse mode autodiff. It grows linearly until the end of the forward pass and then progressively decreases until it outputs the value of the function and the VJP. The memory cost can be mitigated by means of checkpointing techniques presented in Section 8.5.

Decoupled function and VJP evaluations.

The additional memory cost of reverse mode autodiff comes with some advantages. If we need to compute $\partial f(\s_0)^*[\u_i]$ for $n$ different output directions $\u_i$, we only need to compute and store once the intermediate computations $\s_k$ and then make $n$ calls to the backward pass. In other words, by storing in memory the intermediate computations $\s_k$, we may instantiate a VJP operator, which we may apply to any $\u$ through the backward pass. Formally, the forward and backward passes can be decoupled as $$\mathrm{forward}(f, \s_0) \coloneqq (f(\s_0), \partial f(\s_0)^*)$$ where $$\partial f(\s_0)^*[\u] \coloneqq \mathrm{backward}(\u; \s_0, \ldots, \s_{K-1}).$$ In functional programming terminology, the VJP $\partial f(\s_0)^*$ is a closure, as it contains the intermediate computations $\s_0, \dots, \s_{K-1}$. The same can be done for the JVP $\partial f(\s_0)$ if we want to apply it to multiple directions $\v_i$.

8.1.3 Complexity of computing entire Jacobians

In this section, we analyze the time and space complexities of forward-mode and reverse-mode autodiff for computing the entire Jacobian matrix $\jac f(\s_0)$ of a computation chain $f = f_K \circ \dots \circ f_1$, where $f_k \colon \cS_{k-1} \to \cS_k$. We assume $\cS_k \subseteq \RR^{D_k}$, $D_K = M$ and $D_0 = D$. Therefore, we have $f \colon \RR^D \to \RR^M$ and $\jac f(\s_0) \in \RR^{M \times D}$.

Complexity of forward-mode autodiff

Using Definition 2.9, we find that we can extract each column $[\jac f(\s_0)]_{:,j} \in \RR^M$ of the Jacobian matrix, for $j \in [D]$, by multiplying with the standard basis vector $\e_j \in \RR^D$: $$\begin{align*} [\jac f(\s_0)]_{:, 1} &= \partial f(\s_0)[\e_1] \\ &\vdots \\ [\jac f(\s_0)]_{:, D} &= \partial f(\s_0)[\e_D]. \end{align*}$$ Computing the full Jacobian matrix therefore requires $D$ JVPs with vectors in $\RR^D$. Assuming each $f_k$ in the chain composition has the form $f_k: \R^{D_{k-1}} \rightarrow \R^{D_k}$, seen as a matrix, $\jac f_k(\s_{k-1})$ has size $D_{k} \times D_{k-1}$. Therefore, the computational cost of $D$ JVPs is $O\left(D \sum_{k=1}^K D_k D_{k-1}\right)$. The memory cost is $O(\max_{k \in [K]} D_k)$, since we can release intermediate computations after each layer is processed. Setting $D_1 = \dots = D_{K-1} = D$ for simplicity and using $D_K=M$, we obtain that the computational cost of computing $D$ JVPs and therefore of computing the full Jacobian matrix by forward-mode autodiff is $O(MD^2 + KD^3)$. The memory cost is $O(\max\{D, M\})$. If a function has a single input ($D=1$), then the forward mode computes the entire Jacobian at once, which reduces to a single directional derivative.

Complexity of reverse-mode autodiff

Using Definition 2.9, we find that we can extract each row of the Jacobian matrix $[\jac f(\s_0)]_{i} \in \RR^D$, for $i \in [M]$, by multiplying with the standard basis vector $\e_i \in \RR^M$: $$\begin{align*} [\jac f(\s_0)]_{1} &= \partial f(\s_0)^*[\e_1] \\ &\vdots \\ [\jac f(\s_0)]_{M} &= \partial f(\s_0)^*[\e_M]. \end{align*}$$ Computing the full Jacobian matrix therefore requires $M$ VJPs with vectors in $\RR^M$. Assuming as before that each $f_k$ in the chain composition has the form $f_k: \R^{D_{k-1}} \rightarrow \R^{D_k}$, the computational cost of $M$ VJPs is $O\left(M \sum_{k=1}^K D_k D_{k-1}\right)$. However, the memory cost is $O(\sum_{k=1}^{K} D_k)$, as we need to store the intermediate computations for each of the $K$ layers. Setting $D_0 = \dots = D_{K-1} = D$ for simplicity and using $D_K=M$, we obtain that the computational cost of computing $M$ VJPs and therefore of computing the full Jacobian matrix by reverse-mode autodiff is $O(M^2D + KMD^2)$. The memory cost is $O(KD + M)$. If the function has a single output ($M=1$), reverse-mode autodiff computes the entire Jacobian at once, which reduces to the gradient.

When to use forward-mode vs. reverse-mode autodiff?

We summarize the time and space complexities in Table 8.1. Generally, if $M<D$, reverse-mode is more advantageous at the price of some memory cost. If $M\geq D$, forward mode is more advantageous.

8.2 Feedforward networks

In the previous section, we derived forward-mode autodiff and reverse-mode autodiff for computation chains with an input $\s_0 \in \cS_0$. In this section, we now derive reverse-mode autodiff for feedforward networks, in which each layer $f_k$ is now allowed to depend explicitly on some additional parameters $\w_k \in \cW_k$. The recursion is $$\begin{align*} \s_0 & \coloneqq \x \in \cS_0 \\ \s_1 & \coloneqq f_1(\s_0, \w_1) \in \cS_1 \\ & \hspace{6pt} \vdots \\ \s_K & \coloneqq f_K(\s_{K-1}, \w_K) \in \cS_K \\ f(\x, \w) & \coloneqq \s_K, \end{align*}$$ where $\cS_0 = \cX$ and $\w = (\w_1, \dots, \w_K) \in \cW_1 \times \dots \times \cW_K$. Each $f_k$ is now a function of two arguments. The first argument depends on the previous layer, but the second argument does not. This is illustrated in Fig. 8.5. We now explain how to differentiate a feedforward network.

8.2.1 Computing the adjoint

The function has the form $f \colon \cE \to \cF$, where $\cE \coloneqq \cX \times (\cW_1 \times \dots \times \cW_K)$ and $\cF \coloneqq \cS_K$. From Section 2.3, we know that the VJP has the form $\partial f(\x, \w)^* \colon \cF \to \cE$. Therefore, we want to be able to compute $\partial f(\x, \w)^*[\u] \in \cE$ for any $\u \in \cF$.

Fortunately, the backward recursion is only a slight modification of the computation chain case. Indeed, since $f_k \colon \cE_k \to \cF_k$, where $\cE_k \coloneqq \cS_{k-1} \times \cW_k$ and $\cF_k \coloneqq \cS_k$, the intermediate VJPs have the form $\partial f_k(\s_{k-1}, \w_k)^* \colon \cF_k \to \cE_k$. We therefore arrive at the recursion $$\begin{align*} \r_K &= \u \in \cS_K \\ (\r_{K-1}, \g_K) &= \partial f_K(\s_{K-1}, \w_K)^* [\r_K] \in \cS_{K-1} \times \cW_K \\ &\hspace{6pt} \vdots \\ (\r_0, \g_1) &= \partial f_1(\s_0, \w_1)^* [\r_1] \in \cS_0 \times \cW_1. \end{align*}$$ The final output is $$\partial f(\x, \w)^*[\u] = (\r_0, (\g_1, \dots, \g_K)).$$

8.2.2 Computing the gradient

We often compose a network with a loss function $$L(\w; \x, \y) \coloneqq \ell(f(\x, \w); \y) \in \RR.$$ From Proposition 2.7, the gradient is given by $$\nabla L(\w; \x, \y) = (\g_1, \dots, \g_K) \in \cW_1 \times \dots \times \cW_K$$ where $$\partial f(\x, \w)^*[\u] = (\r_0, (\g_1, \dots, \g_K)),$$ with $\u = \nabla \ell(f(\x, \w); \y) \in \cS_K$. The output $\r_0 \in \cS_0$, where $\cS_0 = \cX$, corresponds to the gradient w.r.t. $\x \in \cX$ and is typically not needed, except in generative modeling settings. The full procedure is summarized in Algorithm 8.3.

8.3 Computation graphs

In the previous sections, we reviewed autodiff for computation chains and its extension to feedforward networks. In this section, we review its generalization to computation graphs, introduced in Section 4.1.3. Our formalism assumes, without loss of generality, that functions can take multiple inputs but only produce a single output, as we can always group multiple outputs as a tuple. This enables a one-to-one correspondence between output variables $\s_k$ and functions $f_k$.

8.3.1 Forward mode

The forward mode corresponds to computing the JVP of the program in an input direction $\v \in \cS_0$. In the case of a computation chain $f \coloneqq f_K \circ \dots \circ f_1$, each $f_k$ takes a single input $\s_{k-1} \in \cS_{k-1}$ and each $\partial f_k(\s_{k-1})$ takes a single input direction $\t_{k-1} \in \cS_{k-1}$. The forward pass on iteration $k \in [K]$ is then, starting from $\t_0 \coloneqq \v$, $$\begin{align*} \s_k &\coloneqq f_k(\s_{k-1}) \in \cS_k \\ \t_k &\coloneqq \partial f_k(\s_{k-1})[\t_{k-1}] \in \cS_k. \end{align*}$$ In the case of a computation graph, specified by functions $f_1, \dots, f_K$ in topological order, each $f_k$ may now take multiple inputs $(\s_{i_1}, \dots, \s_{i_{p_k}}) \in \cS_{i_1} \times \dots \times \cS_{i_{p_k}}$, and each $\partial f_k(\s_{i_1}, \dots, \s_{i_{p_k}})$ takes as many input directions $(\t_{i_1}, \dots, \t_{i_{p_k}}) \in \cS_{i_1} \times \dots \times \cS_{i_{p_k}}$, where $(i_1, \ldots, i_{p_k}) \coloneqq \parent(k)$ are the parents of $f_k$. The forward pass on step $k \in [K]$ then computes both intermediate inputs and directions as $$\begin{align*} \s_k &\coloneqq f_k(\s_{i_1}, \dots, \s_{i_{p_k}}) \in \cS_k \\ \t_k &\coloneqq \partial f_k(\s_{i_1}, \dots, \s_{i_{p_k}})[\t_{i_1}, \dots, \t_{i_{p_k}}] \in \cS_k. \end{align*}$$ Let us recall that $\partial_i f_k$ means that we differentiate $f_k$ w.r.t.  its $i$^th argument. Using the fan-in rule in Proposition 2.8, we obtain that the derivatives are propagated as $$\t_k = \sum_{j=1}^{p_k} \partial_{j} f_k(\s_{i_1}, \dots, \s_{i_{p_k}})[\t_{i_j}] \in \cS_k.$$ The final output is $\partial f(\s_0)[\v] = \t_K$. The resulting generic forward-mode procedure is summarized in Fig. 8.7 and in Algorithm 8.4. Intuitively, the algorithm consists in computing and summing intermediate JVPs along the forward pass. Although not explicitly mentioned, we can release $\s_k$ and $\t_k$ from memory when no child node depends on node $k$.

8.3.2 Reverse mode

The reverse mode corresponds to computing the VJP of the program in an output direction $\u \in \cS_K$. In the case of a computation chain $f \coloneqq f_K \circ \dots \circ f_1$, each intermediate variable $\s_k$ is used only once to compute the next variable $\s_{k+1} \coloneqq f_{k+1}(\s_k)$. To compute the VJP of the chain, it then suffices to reverse the order of the operations and to use the corresponding VJPs. Since each function takes a single input, each VJP $\partial f_k(\s_k)^*[\r_k]$ produces a single output direction $\r_{k-1} \in \cS_k$. A backward pass therefore computes from $k \coloneqq K$ to $k \coloneqq 1$, starting from $\r_K \coloneqq \u$, $$\r_{k-1} \coloneqq \partial f_k(\s_{k-1})^*[\r_k] \in \cS_{k-1}.$$ The final output is $\r_0 \coloneqq \partial f(\s_0)^*[\u]$.

In computation graphs, intermediate variables may be used more than once, since a node $k$ may have several children nodes. This complicates reversing the computation graph. To circumvent this issue, following the formalism of Roy Frostig (Murphy, 2023, Section 6.2.2.2), we may formally distinguish between the output $\s_k$ of $f_k$ and the inputs $\s_{k \rightarrow j}$ of $f_j$ for $j \in \child(k)$, the children of $f_k$. We do so by introducing an operation that simply duplicates $\s_k$, $$\begin{align*} \dup(\s_k) & \coloneqq (\s_k, \ldots, \s_k). \end{align*}$$ The tuple output by $\mathrm{dup}$ is of length $|\child(k)|$. The forward pass can then be formally rewritten as, for $k \in [K]$, $$\s_k \coloneqq f_k(\s_{i_1 \rightarrow k}, \ldots, \s_{i_{p_k} \rightarrow k})$$ where $(i_1, \dots, i_{p_k}) \coloneqq \parent(k)$ followed by $$(\s_{k \rightarrow j_1}, \ldots, \s_{k \rightarrow j_{c_k}}) \coloneqq \dup(\s_k),$$ where $(j_1, \ldots, j_{c_k}) \coloneqq \child(k)$. The benefit of this approach is that each duplicated output is input only once to the subsequent child functions. Thanks to this, we can now associate to each variable, $\s_{i \rightarrow j}$ or $\s_k$, a single corresponding intermediate variable, $\r_{i \rightarrow j}$ or $\r_k$, in the backward pass. The reverse mode can then simply be written by going through the VJPs of the functions $f_k$ and the VJP of $\dup$ in reverse order.

For the VJP of the $\mathrm{dup}$ operation, let us denote the intermediate variables in the backward pass associated with $\s_{k \rightarrow j_1}, \ldots, \s_{k \rightarrow j_{c_k}}$ by $\r_{k \rightarrow j_1}, \dots, \r_{k \rightarrow j_{c_k}}$. Following the fan-out rule in Proposition 2.9, the VJP of the $\dup$ operation on iteration $k$ is $$\begin{align*} \r_k & = \dup(\s_k)^*[\r_{k \rightarrow j_1}, \dots, \r_{k \rightarrow j_{c_k}}] \\ & = \sum_{i=1}^{c_k} \dup_i(\s_k)^*[\r_{k \rightarrow {j_i}}]\\ & = \sum_{i=1}^{c_k} \r_{k \rightarrow {j_i}}. \end{align*}$$ In the last line, we used that $\dup_i(\s) = \s$ by definition of the duplication operation so $\partial\dup_i(\s) = \idm$ and $\partial\dup_i(\s)^* = \idm$. The VJP of $\dup$ justifies why, if an intermediate value $\s_k$ is used by later functions $f_{j_1},\dots, f_{j_{c_k}}$, for $(j_1,\dots,j_{c_k}) \coloneqq \child(k)$, the derivatives with respect to $\s_k$ need to sum all the variations through the $f_j$ functions into the variable $\r_k$.

For the VJP of the $f_k$ functions, following the fan-in rule in Proposition 2.8, the VJP returns multiple output variations, $$\r_{i_1 \rightarrow k}, \dots, \r_{i_{p_k} \rightarrow k} = \partial f_k(\s_{i_1 \rightarrow k}, \ldots, \s_{i_{p_k} \rightarrow k})^*[\r_k],$$ where $(i_1, \dots, i_{p_k}) \coloneqq \parent(k)$, and where, for $j \in \{1, \dots, p_k\}$, $$\r_{i_j \rightarrow k} \coloneqq \partial_j f_k(\s_{i_1 \rightarrow k}, \ldots, \s_{i_{p_k} \rightarrow k})^*[\r_k] \in \cS_{i_j}.$$ The overall formalism is summarized in Fig. 8.8.

Implementation

The duplication operation $\dup$ is just a formalism to mathematically derive the reverse mode. In practice, the variables $\s_k$ are not duplicated but accessed several times during the backward pass. The variables $\r_k$ are computed by accumulating $\r_{k \rightarrow j}$ into $\r_k$ each time a variable $\r_{k \rightarrow j}$ is computed. Therefore, for each $k \in [K]$, we can compute the VJP and perform the in-place updates, $$\begin{align*} \r_{i_1 \rightarrow k}, \dots, \r_{i_{p_k} \rightarrow k} & = \partial f_k(\s_{i_1}, \ldots, \s_{i_{p_k}})^*[\r_k] \\ \r_{i_j} & \leftarrow \r_{i_j} + \r_{i_j \rightarrow k} \quad \forall j \in \{1, \dots, p_k\}. \end{align*}$$ The topological ordering ensures that $\r_k$ has been fully computed when we reach $f_k$. The resulting generic reverse-mode procedure is presented in Algorithm 8.5.