Moonwalk: Inverse-Forward Differentiation

## Abstract

Abstract Backpropagation, while effective for gradient computation, falls short in addressing memory consumption, limiting scalability. This work explores forward-mode gradient computation as an alternative in invertible networks, showing its potential to reduce the memory footprint without substantial drawbacks.
We introduce a novel technique based on a vector-inverse-Jacobian product that accelerates the computation of forward gradients while retaining the advantages of memory reduction and preserving the fidelity of true gradients. Our method, Moonwalk, has a time complexity linear in the depth of the network, unlike the quadratic time complexity of naïve forward, and empirically reduces computation time by several orders of magnitude without allocating more memory. We further accelerate Moonwalk by combining it with reverse-mode differentiation to achieve time complexity comparable with backpropagation while maintaining a much smaller memory footprint. Finally, we showcase the robustness of our method across several architecture choices. Moonwalk is the first forward-based method to compute true gradients in invertible networks in computation time comparable to backpropagation and using significantly less memory.

## 1 Introduction

In recent years, the evolution of deep learning models has been significantly influenced by the advent of automatic differentiation (AD) packages (Paszke et al., ; Bradbury et al., ; Abadi et al., ), facilitating expedited model construction and research. The most commonly used differentiation algorithm is gradient backpropagation (Backprop), which effectively addresses the challenge of time-efficient gradient computation, but fails to tackle the issue of memory consumption (Gomez et al., ; Chakrabarti & Moseley, ). The issue is that, during the construction and reversion of the computation graph, Backprop retains all activations acquired during the forward execution phase. For large networks, the memory footprint scales with the number of activations (Novikov et al., ) which can result in significant memory overhead limiting the ability to scale neural networks.

An alternative to traditional Backprop is forward-mode gradient computation (Forward), a long-established concept in training neural networks (Williams & Zipser, ) but one that remains less widely adopted in practice due to its very high computation time. The notion of employing forward differentiation has recently gained attention as a promising strategy for alleviating layer activation memory constraints inherent in Backprop (Silver et al., ; Baydin et al., ; Fournier et al., ). However, a primary drawback of forward gradient computation lies in its requirement to compute full Jacobian matrices for every layer of the computation graph, proving more computationally expensive than the vector-Jacobian products used in Backprop.

To address this challenge, one potential approach involves projecting true gradients onto a subspace and computing only these projections in forward-mode (Baydin et al., ). Some prior works employ projections onto random subspaces, thereby introducing variance in gradient estimation that limits applicability in large networks (Silver et al., ). Other works predict the gradient direction based on auxiliary networks or past gradients and use it as a preferred projection subspace, showing promise in reducing variance but to-date falling short of the end-to-end training accuracy of Backprop (Fournier et al., ).

In this work, we identify a mathematical identity in computing gradients of invertible networks that allows significant savings in memory and time when computing true gradients in forward-mode without projection.
Our method, Moonwalk, relies on the observation that, once we obtain the gradient of the objective with respect to just the input of the first layer, we can efficiently compute the remaining gradients using a vector-inverse-Jacobian product (VIJP) operator in forward-mode.

To our knowledge, Moonwalk is the first forward-mode differentiation method that can outperform Backprop in both time and memory requirements when computing true gradients in invertible networks.

Our approach computes gradients in two phases.
In the first phase, we compute the gradient of the objective with respect to the input.
This can be done via forward-mode differentiation by computing the full Jacobian of just the input, which is typically much smaller than the Jacobian of the entire network.
This version of Moonwalk is significantly faster than full forward differentiation, and potentially fast enough to be worth the immense memory saving over Backprop, particularly when the dimensionality of the input is small.
When the input is high-dimensional, the first phase can instead compute the input gradient in reverse-mode, which typically still takes much less memory than full backprop.

The second phase then uses the computed gradient vector in a forward pass to obtain the parameter gradients through an operator involving the vector-inverse-Jacobian product with respect to the input of each layer and the Jacobian-vector product with respect to its parameters.

In summary, this work contributes two novel automatic differentiation methods for invertible networks that compute full gradients in forward-mode:

- •
Moonwalk, an entirely forward-mode method that significantly reduces the time requirements of naïve forward differentiation and is the first forward-mode method to feasibly address the memory challenge of Backprop, being particularly fast when the input dimension is very small; and
- •
Mixed-mode Moonwalk, a variant implementing the first phase of Moonwalk in reverse-mode for further acceleration to achieve time complexity comparable with Backprop, while maintaining a memory footprint much smaller than backprop.

## 2 Related Work

Reducing memory. Previous studies have used checkpointing to alleviate the memory footprint of neural networks (Martens & Sutskever, ; Chen et al., ; Gruslys et al., ). This technique reduces the memory consumption in a network with $L$ layers by a factor of $\sqrt{L}$ through the selective storage of activations at intervals of $\sqrt{L}$ layers and their forward-mode recomputation between checkpoints as they become needed during Backprop. Another approach for reducing memory is to avoid storing activations altogether in reversible networks. When the input to each layer can be computed from its output using the inverse function, several methods recompute activations backwards during Backprop (Gomez et al., ; MacKay et al., ; Mangalam et al., ). Bulo et al. () replaced only ReLU and batch normalization layers with invertible variants, reducing memory usage
up to 50%. Jaderberg et al. () used synthetic gradients for efficient activation storage using pre-trained gradient estimators.

Reversible architectures. Recently, reversible architectures have gained significant attention owing to their diverse applications in reducing memory (Gomez et al., ; MacKay et al., ; Mangalam et al., ), enhancing learned representation (Jacobsen et al., ), boosting performance (Kingma & Dhariwal, ), and generative modeling (Dinh et al., ; Rezende & Mohamed, ). Numerous studies have proposed invertible versions of well-known deep neural network architectures such as invertible residual networks (Gomez et al., ; Behrmann et al., ), reversible recurrent neural networks (MacKay et al., ), reversible Vision Transformers (Mangalam et al., ) and $1\times 1$ convolutions (Kingma & Dhariwal, ).

Forward propagation. The concept of learning neural network weights in a forward fashion has been originally explored in the real-time recurrent learning (RTRL) algorithm (Williams & Zipser, ), similar to forward-mode AD. This idea becomes more attractive when directional derivatives (effectively, gradient projection) are employed to eliminate the additional computation time of full forward differentiation. However,
this introduces noise into the gradients (Silver et al., ; Baydin et al., ), which is coming from the random directions or imperfect gradient predictions used in directional derivatives. To address the variance of the forward gradients, Ren et al. () proposed using local greedy loss functions and Fournier et al. () employed local auxiliary networks as the tangent vectors.

## 3 Background

### 3.1 Notation

Consider a neural network $f_{\theta}:\mathbb{R}^{n}\rightarrow\mathbb{R}^{k}$ with $L$ layers and parameters $\theta=\{\theta_{i}\}_{i=1,\ldots,L}$. We will denote the output of layer $i\in\{1,...,L\}$ by $x_{i}=f_{i}(x_{i-1};\theta_{i})$, where $x_{0}\in\mathbb{R}^{n}$ is the input to the network. Let $J(x_{L})=J(f_{\theta}(x_{0}))$ be the scalar loss function, whose gradient with respect to $\theta$ we wish to compute. For simplicity, we will assume that the parameters $\theta_{i}\in\mathbb{R}^{d}$ of all layers have the same dimension $d$ and are layer-wise disjoint and that the outputs $x_{i}\in\mathbb{R}^{k}$ of all layers have the same dimension $k=n$.

Throughout the paper we refer to the Jacobian-vector product, the vector-Jacobian product, and the vector-inverse-Jacobian product as jvp, vjp, and vijp, respectively, and define these operators as

$$ $\displaystyle\text{jvp}(f,\theta,u)$ $\displaystyle=\left(\nicefrac{{\partial f}}{{\partial\theta}}\right)u,$ (1) $\displaystyle\text{vjp}(f,\theta,v)$ $\displaystyle=v\left(\nicefrac{{\partial f}}{{\partial\theta}}\right),\text{ and}$ (2) $\displaystyle\text{vijp}(f,\theta,v)$ $\displaystyle=v\left(\nicefrac{{\partial f}}{{\partial\theta}}\right)^{-1},$ (3) $$

where $u$ is the tangent column vector, $v$ is the cotangent row vector, and the Jacobian, taken here with respect to $f$’s parameters $\theta$, can also be taken with respect to $f$’s input $x$.
vjp and jvp are commonly used operators in AD frameworks, and we use their JAX implementation jax.jvp and jax.vjp (Bradbury et al., ).

### 3.2 Forward-Mode Gradients

Forward-mode differentiation is an alternative to Backprop for computing the gradients(^1^11We use the notation of row-vector gradients, which are more accurately called *total derivatives*.)

$$ $\displaystyle\frac{\partial J}{\partial\theta_{i}}=\frac{\partial J}{\partial x _{L}}\left(\prod_{j=L}^{i+1}\frac{\partial x_{j}}{\partial x_{j-1}}\right) \frac{\partial x_{i}}{\partial\theta_{i}}.$ (4) $$

While Backprop computes the product in (4) from left to right, forward-mode computes it from right to left.
The suffix of the product can be computed during the forward execution of the function, such that the activations $x_{i}$ need not be stored, unlike in Backprop. On the other hand, while the prefixes are vectors of dimension $n$ that can be computed using vjp, the suffixes are matrices of dimensions $n\times d$. To avoid storing these matrices in memory, they are commonly computed column-by-column using jvp, which increases the asymptotic time complexity of forward-mode by a factor of $\min(n,d)L$ compared to Backprop (Table 1).

### 3.3 Projected Forward-Mode Gradients

Projected forward-mode gradients (Silver et al., ; Baydin et al., ) are the directional derivatives computed in a forward fashion. For a unit-length tangent vector $u\in\mathbb{R}^{d}$, the projection of the gradient onto $u$ has length

$$ $\displaystyle\text{jvp}(J,\theta_{i},u)=\frac{\partial J}{\partial\theta_{i}}u$ $\displaystyle=\frac{\partial J}{\partial x_{L}}\left(\prod_{j=L}^{i+1}\frac{ \partial x_{j}}{\partial x_{j-1}}\right)\frac{\partial x_{i}}{\partial\theta_{ i}}u,$ (5) $$

and the projection is $\frac{\partial J}{\partial\theta_{i}}uu^{\intercal}$.
The vector $u$ is usually picked at random from a normal distribution or predicted based on past gradients (Silver et al., ; Baydin et al., ). Now the suffixes of (5) are vectors and they can be computed in forward-mode using a recursive application of jvp with respect to each layer’s input.
Computing projected forward-mode gradients, then, matches the asymptotic time complexity to that of Backprop (Table 1) but introduces some variance into the gradient estimator, which now depends on the choice of $u$.

## 4 Moonwalk

In order to benefit from the memory advantage of forward-mode gradients, while keeping the time complexity similar to that of Backprop and avoiding the introduction of noisy gradients through projection, we restrict our attention to the class of invertible neural networks, in which each layer $x_{i}=f_{i}(x_{i-1};\theta_{i})$ has a unique inverse $x_{i-1}=f_{i}^{-1}(x_{i};\theta_{i})$ with respect to its input for any $x_{i}$ in its image. Because $f_{i}$ is also assumed to be differentiable, this will constrain the output dimension of each layer to equal its input dimension, thus $k=n$.
Note that, downstream of all parametrized layers, the network no longer has to be invertible, and in particular the loss $J$ is still scalar.

Figure: Algorithm 1 Moonwalk

In invertible networks, the Jacobian of each layer with respect to its input is also guaranteed to be invertible with its inverse equal to the inverse-function Jacobian, i.e.

$$ $\displaystyle\frac{\partial f_{i}^{-1}}{\partial x_{i}}=\frac{\partial x_{i-1} }{\partial x_{i}}=\left(\frac{\partial x_{i}}{\partial x_{i-1}}\right)^{-1}.$ $$

We can rewrite each layer’s parameter gradient of the loss as

$$ $\displaystyle g_{i}\coloneqq{}$ $\displaystyle\frac{\partial J}{\partial\theta_{i}}=\frac{\partial J}{\partial x _{i}}\frac{\partial x_{i}}{\partial\theta_{i}}=\frac{\partial J}{\partial x_{i }}\frac{\partial x_{i}}{\partial x_{0}}\left(\frac{\partial x_{i}}{\partial x_ {0}}\right)^{-1}\frac{\partial x_{i}}{\partial\theta_{i}}$ (6) $\displaystyle={}$ $\displaystyle\frac{\partial J}{\partial x_{0}}\left(\prod_{j=i}^{1}\frac{ \partial x_{j}}{\partial x_{j-1}}\right)^{-1}\frac{\partial x_{i}}{\partial \theta_{i}}$ (7) $\displaystyle={}$ $\displaystyle\frac{\partial J}{\partial x_{0}}\left(\prod_{j=1}^{i}\frac{ \partial x_{j-1}}{\partial x_{j}}\right)\frac{\partial x_{i}}{\partial\theta_{ i}}.$ (8) $$

Given the input gradient $\frac{\partial J}{\partial x_{0}}$, we can use (8) to compute the parameter gradient of each layer in forward-mode. To this end, denote the activation gradient of layer $i$ by

$$ $\displaystyle h_{i}\coloneqq\frac{\partial J}{\partial x_{i}}=\frac{\partial J }{\partial x_{0}}\prod_{j=1}^{i}\frac{\partial x_{j-1}}{\partial x_{j}},$ (9) $$

for $1\leq i\leq L$, and the input gradient by $h_{0}\coloneqq\frac{\partial J}{\partial x_{0}}$. Note that $h_{i}$ satisfies the forward recursion

$$ $\displaystyle h_{i}=h_{i-1}\frac{\partial x_{i-1}}{\partial x_{i}}=\text{vijp} (f_{i},x_{i-1},h_{i-1}),$ (10) $$

and from (8) we have

$$ $\displaystyle g_{i}=h_{i}\frac{\partial x_{i}}{\partial\theta_{i}}=\text{vjp}( f_{i},\theta_{i},h_{i}).$ (11) $$

Assuming that we have the input gradient $h_{0}$, we can construct the parameter gradient for each layer on the fly by two operators in equations 10 and 11 and store only $h_{i}\in\mathbb{R}^{k}$ for the next layer’s computation. The complete procedure is given in Algorithm 1 and illustrated in Figure LABEL:fig:grads_inv.

In the next two sections we describe the computation of $h_{0}$ with either forward-mode or Backprop and discuss the trade-offs of using each variant.

### 4.1 Pure-Forward Moonwalk

One way to obtain the input gradient $h_{0}$ is to compute it element-by-element using a projected forward-mode with a standard basis of tangent vectors. Specifically, $h_{0,i}$, the $i$-th component of $h_{0}$, can be computed as in (5) with the $i$-th standard basis vector $e_{i}$ as the tangent vector. This method thus uses forward-mode $n$ times to construct $h_{0}$.

Figure LABEL:fig:h0_fp shows the computation flow of $h_{0}$. Note that we only need to store $h_{0}$ components in the final stage. By computing $h_{0}$ in forward-mode we avoid storing activations. On the other hand, when the input dimension is large, this method takes infeasible time to compute.

### 4.2 Mixed-Mode Moonwalk

Alternatively, we can use Backprop to compute $h_{0}$. While it may seem counterintuitive that a method including this phase could improve on full Backprop, the key insight is that only part of the computation graph is downstream of $x_{0}$. When using Backprop for $h_{0}$, we can therefore avoid storing activations computed from parameters $\theta_{i}$ independent of any $x_{i}$. Another consideration, more practical than principled, is that many Backprop implementations store all gradients before applying any (even in cases such as ours, in which gradients can be applies immediately and discarded), and here unlike full Backprop parameter gradients need not be stored. In many network architectures of interest, these distinctions lead to significant memory savings in Backprop when only computing the input gradient. Figure LABEL:fig:h0_bp shows the flow of this procedure.

Computing $h_{0}$ via Backprop reduces the time complexity by a factor of $n$, compared to forward-mode, but increases the memory footprint, which in the worst case can now be as high as Backprop’s. However, depending on the network architecture, as analyzed in Section 6.3, real-world architectures are often much better than the worst case, leading to significant memory saving over Backprop.

## 5 Complexity Analysis

While estimating the exact time and memory consumption of different methods for computing the gradients hugely depends on the choice of the network’s architecture and the detailed implementation, in this section, we will provide a rough estimate of the time and memory complexity of our methods and compare them with related previous works (Table 1). We simplify the analysis by considering the mini-batch size to be 1 for all methods. We estimate the computational complexity of the following methods:

- 1.
Backprop: Throughout the forward pass, all activations are cached, and subsequently, during a backward pass, gradients for each layer are computed using vjp.
- 2.
Forward: During the forward pass, complete Jacobians for each layer are computed.
- 3.
ProjForward: Gradients of the parameters of each layer, projected in a random or predicted direction, are obtained using jvp concurrently with the forward pass.
- 4.
RevBackprop: Initially, a full forward pass is conducted to obtain the loss and the gradient of the last layer. Subsequently, in a backward pass, the input of each layer is computed from its output and used with vjp to backpropagate gradients.
- 5.
Moonwalk: Initially, the input gradient is computed using Forward. Subsequently, gradients are obtained using vijp and vjp concurrently with a second forward pass. See section 4.1.
- 6.
Mixed: Similar to Moonwalk, but the input gradient is computed using Backprop. This modification significantly reduces computation time at the expense of some memory impact. See section 4.2.

We evaluate time based on the standard cost of matrix multiplication (product of their two outer dimensions and shared inner dimension) without considering optimization tricks. To evaluate memory, we define $M_{x,i}$ to be the required memory to store the necessary information to compute $\frac{\partial x_{i}}{\partial x_{i-1}}$, and $M_{\theta,i}$ the added memory to also compute $\frac{\partial x_{i}}{\partial\theta_{i}}$. For simplicity, we assume that these values are the same across layers and omit the layer index. We refer to memory consumption as the extra amount of memory needed to compute gradients without reflecting the memory to store the parameters or gradients themselves after computation.

For Backprop, we have to store activations required for both input and parameter gradients for every layer, which results in $O(M_{x}L+M_{\theta}L)$ memory complexity. For Mixed, we only need to store $M_{x}$ for every layer in the first phase, in order to compute the input gradient $h_{0}$, and then we can reuse $M_{\theta}$ in the second phase after computing each parameter gradient, for a total memory complexity of $O(M_{x}L+M_{\theta})$.
All other methods can discard activation information after each layer, for a memory complexity of $O(M_{x}+M_{\theta})$.

Backprop computation consists of computing two vector-Jacobian products in each layer, $\text{vjp}(f_{i},x_{i-1},\frac{\partial J}{\partial x_{i}})$ and $\text{vjp}(f_{i},\theta_{i},\frac{\partial J}{\partial x_{i}})$, which accounts for per-layer time complexity of $O(n^{2})$ and $O(nd)$, respectively, and for a total of $O(n^{2}L+ndL)$. RevBackprop additionally needs to evaluate the inverse function $f_{i}^{-1}(x)$, which does not impact the overall complexity in the terms we consider.

In Forward, each single parameter $\theta_{j,\ell}$ in layer $j$, of the total $dL$ parameters, generates a pass to compute its gradient, in which we compute $\text{jvp}(f_{i},x_{i-1},\frac{\partial x_{i-1}}{\partial\theta_{j,\ell}})$ in each layer $i>j$, for a total of $O(n^{2}dL^{2})$.

ProjForward with tangent $u=\{u_{i}\}_{i=1,\ldots,L}$ is similar to Forward with just a single pass taking $O(n^{2}L)$ time, but initialized in each layer from $\text{jvp}(f_{i},\theta_{i},u_{i})$ for a total of $O(n^{2}L+ndL)$, which happens to coincide with the time complexity of Backprop.

The first phase of Moonwalk computes $\text{jvp}(f_{i},x_{i-1},\frac{\partial x_{i-1}}{\partial x_{0}}e_{\ell})$ in each layer for each input element $\ell$ , for a total time complexity of $O(n^{3}L)$. The second phase computes $\text{vijp}(f_{i},x_{i-1},h_{i-1})$ and $\text{vjp}(f_{i},\theta_{i},h_{i})$ in each layer for $O(n^{2}L+ndL)$ more, and a total of $O(n^{3}L+ndL)$.
Mixed replaces the first phase with Backprop for just the input gradient, incurring time complexity $O(n^{2}L)$, in addition to $O(n^{2}L+ndL)$ from the same second phase as Moonwalk, for a total of $O(n^{2}L+ndL)$ time complexity.

The complexity of memory consumption in forward-mode method only depends on the largest layer in the network, while methods involving Backprop require storing all the activations. Thus the memory consumption grows linearly with the number of layers in these methods.

Note that while Table 1 only shows the order of growth, the actual time and memory consumption of these methods have been empirically studied in the following section.

**Table 1: Asymptotic complexity and key characteristics of Moonwalk, its Mixed-mode variant, and four prior methods.**
| Method | Time | Memory | Stochastic | Numerically Stable | Requires Backprop |
| --- | --- | --- | --- | --- | --- |
| Backprop | $O(n^{2}L+ndL)$ | $O\left(M_{x}L+M_{\theta}L\right)$ | ✗ | ✓ | ✓ |
| Forward | $O(n^{2}dL^{2})$ | $O(M_{x}+M_{\theta})$ | ✗ | ✓ | ✗ |
| ProjForward | $O(n^{2}L+ndL)$ | $O(M_{x}+M_{\theta})$ | ✓ | ✓ | ✗ |
| RevBackprop | $O(n^{2}L+ndL)$ | $O(M_{x}+M_{\theta})$ | ✗ | ✗ | ✓ |
| Moonwalk | $O(n^{3}L+ndL)$ | $O(M_{x}+M_{\theta})$ | ✗ | ✓ | ✗ |
| Mixed | $O(n^{2}L+ndL)$ | $O(M_{x}L+M_{\theta})$ | ✗ | ✓ | ✓ |

## 6 Experiments

Figure: Figure 1: Maximum allocated memory during training. The input is padded to 32x32x8, which corresponds to allocating memory for 32x32x8 $\cdot$ number of layers $\cdot$ number of blocks parameters.
Refer to caption: x1.png

Figure: Figure 2: Maximum allocated memory during training for a larger network. The input is padded to 32x32x18. The total number of parameters is 200k $\cdot$ number of layers per block.
Refer to caption: x2.png

Figure: Figure 3: Time comparison for computing one batch of 512 with an input size of 32x32x18 across three blocks, each with a varying number of layers per block.
Refer to caption: x3.png

Figure: Figure 4: Train accuracy of three models trained with RevBackprop, Backprop, and Mixed gradient methods for 100 epochs, averaged over 20 runs.
Refer to caption: x4.png

Figure: Figure 5: L2 gradient error between true gradients computed through Backprop and gradients computed through one of the methods, averaged over 20 runs.
Refer to caption: x5.png

### 6.1 Experimental Setup

Model and architecture.
We adopt a modified RevNet architecture (Gomez et al., ) with three blocks. Each block of the network consists of coupling layers that partition the input into two subsets $x_{1}$ and $x_{2}$, and output

$$ $\displaystyle y_{1}$ $\displaystyle=x_{1}$ (12) $\displaystyle y_{2}$ $\displaystyle=x_{2}+\mathcal{F}(x_{1}).$ $$

The function $\mathcal{F}$ can be any arbitrary function, not necessarily invertible. In our case, it is represented by a series of convolutional layers, each followed by a ReLu activation (Agarap, ) and BatchNorm (Ioffe & Szegedy, ) layers.
The input to each layer can be reconstructed from it is output, without the need to cache activations, by

$$ $\displaystyle x_{1}$ $\displaystyle=y_{1}$ (13) $\displaystyle x_{2}$ $\displaystyle=y_{2}-\mathcal{F}(x_{1}).$ $$

In our experiments, the network consists of three sequential blocks. The number of layers in each block is set in each experiment to one of {1, 3, 5, 10}. Following each block, we apply an invertible down-sampling operator $\psi$, as proposed by Jacobsen et al. (). The operator $\psi$ effectively divides the spatial dimensions into the channel dimension, resulting in a fourfold increase in the feature-space dimension for each successive block.

Data and Hyper-parameters.
We implement all methods using the JAX framework (Bradbury et al., ) and conduct experiments on training a neural network on the CIFAR-10 dataset (). The experiments are performed on an RTX A4500 GPU with a batch size of 512. Computation time is measured after code compilation is completed. Memory consumption is tracked using the nvidia-smi utility every second, with the maximum value recorded. Data preprocessing involves padding with zeros in the channel dimension, expanding the input dimension from 32x32x3 to to one of {32x32x18,32x32x8}. We use the Adam optimizer (Kingma & Ba, ) with a learning rate of $10^{-3}$, and we use an exponential decay scheduler unless specified otherwise.

First, we demonstrate that Forward is markedly more memory efficient than Backprop by measuring the memory consumption of the network during the training of each minibatch. Second, we show that Moonwalk is substantially faster than the Forward. By imposing constraints on the network architecture, we substantially reduce the overall gradient computation time. Third, we illustrate that Moonwalk can be combined with Backprop, further reducing the computation time while preserving significant memory savings. Finally, we establish that Mixed exhibits greater robustness compared to RevBackprop.

### 6.2 Implementation

We implement all methods using JAX (Bradbury et al., ). Forward involves applying the jax.jvp operator sequentially at each layer, with the number of applications corresponding to the number of parameters in that layer. For each layer of the network, we define both forward and inverse functions. To enhance computation speed, we flatten all layer parameters and create an identity matrix used for projection. Utilizing a batched function map accelerates computation, although it comes at the cost of an increased overall memory footprint for forward-based methods. An alternative to mitigate the increased memory footprint is to compute the jax.jvp product one by one without creating a large identity matrix. While this approach reduces the memory footprint, it also significantly increases compilation time and affects overall computation, prohibiting the use of jax.vmap in such scenarios. The depicted performance of Mixed in the following subsections is not optimized, though it is already only twice as slow compared to Backprop. Some of the factors, including data transfer between compiled and non-compiled environments, significantly affect its performance. In theory, it might be possible to make it faster than Reversible.

### 6.3 Memory reduction

For our first experiment, we pad the input size to 32x32x8. The results summarized in Figure 1 suggest that the memory consumption of Backprop increases linearly with the number of layers per block. Backprop necessitates the most substantial memory allocation among the compared approaches, resulting in a twofold disparity at ten layers per block for Moonwalk, Mixed, and RevBackprop. For ProjForward, the disparity is even larger at 2.4 times. The disparity growth for Forward is a bit slower due to its initial compilation overhead.
Additionally, it is noteworthy to observe an initial compilation gap, resulting in increased memory utilization when the architecture comprises only one layer per block.
The memory footprint in Backprop scales with layer size, while Forward memory growth is confined solely to the number of layers. This reduction is achieved by discarding activations that are no longer needed for the forward pass. The resulting memory footprint for Forward consists only of the memory allocated for the storage of gradient updates until they are applied to the model parameters.
Furthermore, we evaluated methods on another type of network, where we padded the input to the size of 32x32x18. In such cases, memory is mostly dominated by $M_{x}L$. Our demonstration in Figure 2 reveals that such architecture choice results in an expanded disparity between Backprop and forward methods. Notably, Mixed experiences an escalation in memory consumption as its allocation begins to be dominated by the activation preserved for gradient computation of the input.

### 6.4 Computation Time

While Forward showcases a substantial reduction in memory usage, it encounters challenges in terms of computation time (Figure 3). The approach demands a considerable amount of computation, rendering it impractical in many scenarios. In contrast, Moonwalk significantly reduces computation time, achieving efficiency up to several orders of magnitude compared to Forward. We evaluate all methods using the same architecture, incorporating varying numbers of layers per block.

The results, illustrated in Figure 3, accentuate the notable efficiency gains achieved by Moonwalk. In the case of six layers, Forward takes 1839 seconds per batch, whereas Moonwalk completes the task in 67 seconds, representing a substantial 27-fold reduction in computing time. Extending the analysis to a model with 60 layers, Forward demands more than 80000 seconds, while our method accomplishes the same task in 700 seconds, manifesting a remarkable time reduction factor of 110.

These outcomes underscore the significant impact on training duration, revealing that employing Moonwalk for training a full model on the CIFAR-10 dataset requires approximately five days—an appreciable improvement over Forward, which necessitates around 300 days.

### 6.5 Time-Memory trade-off

While achieving noteworthy time reduction using Moonwalk, we can further optimize efficiency by precomputing the gradients of the input, denoted as $h_{0}$, in our method using backpropagation. Although the memory reduction in this approach is constrained by the memory required for activations, it significantly decreases overall memory consumption compared to Backprop. Two pivotal components contribute to this reduction:

Elimination of Activations Storage: We refrain from storing activations utilized in computing the gradients of the weights $M_{\theta}L$. However, it is important to note that in this scenario, we are still bound by the activations necessary for computing $h_{0}$. In instances where input is padded to 32x32x8 (see Figure 1), we do not incur a memory disadvantage compared to alternative methods. In scenarios featuring larger input size 32x32x18, as depicted in Figure 2, a noticeable increase in memory usage over Forward becomes apparent.

Sequential Memory Usage: There is no longer a need to store gradients and activations simultaneously, as Backprop does. After precomputing $h_{0}$, we discard all activations, freeing up memory for gradient computations.

In Figure 3, Mixed is only about twice as slow as Backprop. The primary advantage of this method lies in its substantial time reduction compared to Forward.

### 6.6 Numerical Stability

It is worth noting that RevBackprop (MacKay et al., ), while requiring a bijective network with similar computation time, encounters challenges in terms of numerical stability. As demonstrated by Gomez et al. () and Jacobsen et al. (), invertible architectures with finite precision encounter numerical instability issues, leading to divergence from true gradients and convergence to a different local minimum than Backprop. In this study, we highlight that RevBackprop fails to converge when an activation function is added to the output of the layer. Employing the same network and dataloader for simultaneous training, Fig 4 illustrates that both Backprop and Mixed achieve perfect accuracy, while RevBackprop fails to converge. This divergence is attributed to the error accumulating with each update.

Fig 5 displays the gradient error estimation between the algorithm’s gradients and oracle gradients at each step using the same parameters. The experiment reveals the accumulating error causing the network weights to drift away from true gradients. Mixed exhibits more stable convergence due to lower gradient error. Such disparity is a result of different approaches involved in gradient computation. In order to compute the gradients, RevBackprop requires computing the inverse of the function, which might be extremely unstable, as in the case with a tanh activation function. As opposed to RevBackprop, Moonwalk and Mixed methods only require computing inverse-vector Jacobian product and avoid computation of the inverse function itself. In some cases, as we have shown with tanh, this tends to be the more stable approach to computing gradients.
The impact of "gradient drifting" in the reversible method becomes noticeable after 30 epochs.

### 6.7 Stochasticity of ProjForward

ProjForward demonstrates the most conservative memory usage; however, it stands as the sole method not generating true gradients. Instead, it introduces variance in its estimation owing to the adoption of an arbitrary random distribution for projecting the parameter space. However, it has been substantiated to be advantageous in specific scenarios, particularly when a reliable estimate for the gradient direction is available (Silver et al., ; Baydin et al., ; Fournier et al., ).

## 7 Conclusion and Future work

The efficiency of memory utilization in Forward surpasses that of Backprop, thereby alleviating the challenges associated with activation storage constraints. However, this enhanced memory efficiency comes at a considerable computational cost, rendering Forward computationally impractical. To address this limitation, we augmented Forward processing through the incorporation of invertible networks. Our experimental findings reveal that our proposed method, Moonwalk, exhibits a marked acceleration, surpassing Forward by several orders of magnitude. Furthermore, we proposed an approach to mitigate the computational overhead associated with our method. By integrating Backprop for the computation of input gradients, we achieved a substantial acceleration of overall gradient computation.

In this paper, we restricted our attention to the class of invertible neural networks in order to compute the gradients for the next layer given the previous one. Future work would study the effect of using projected forward-mode gradients instead of the full forward-mode. Also, one could investigate the class of neural networks composed of locally invertible layers to shed light on the applicability of Moonwalk for a larger family of architectures.

## Peer Reviews

confidence: 5
contribution: 2
initial_rating: 3
presentation: 4
questions: Could the author address the weaknesses in general ?
The reviewer would appreciate a detailed explanation on weakness 2 and 4 specifically.
soundness: 2
strengths: 1. The paper is very well-structured, easy to read and understand. Apart from very few notations that could be improved, this is a very appreciated feature of the paper.
2. The authors carefully compare their proposed method with relevant alternatives aiming at trading time and memory complexities for neural network training.
3. Experimental setup are clearly detailed and figures are very easy to read, which is highly appreciated.
summary: The presented paper proposes a new algorithm for estimating gradients of reversible neural networks, able to decrease the memory footprint compared to standard backpropagation. The authors further claim a superior numerical stability than reversible backpropagation, hence showing a potential advantage of their method for reversible architectures. The paper is generally well-written, easy to understand, with a careful time and memory complexity analysis with hypothesis clearly stated, from both a theoretical and practical point of view.
weaknesses: 1. Section 3.1 gives explicit notations for the different quantities involved along with their dimensions. However, the dimension of parameters $\theta$ or $\theta_i$ is not mentioned. The reviewer assumes that for all $1 \leqslant i \leqslant L$, the parameters $\theta_i$ have dimensionality $d$. This might be important when discussing complexity issues down the line. Similarly at line 153-154, the reviewer assumes that “the suffixes are of dimension $n_i \times d$. Note that in section 5, authors assumes that all these quantities are the same for all $i$ for simplification. Either this simplification can be done in section 3.1, either all quantities should depend on $i$ in section 3.1 and then simplified in section 5.

2. In section 4.2, authors argue that using backprop to compute the gradient $h_0$ only can lead to substantial memory savings. This claim stems from two separate arguments:
- The first is that if we only want to compute $h_0$ , we do not need to store any activation computed from parameters $\theta_i$ independent from $x_i$. While this is true, I would like the authors to give some examples of architectures where this argument would be relevant.
- The second argument is that parameter gradients can take up a substantial portion of memory during backpropagation. This argument needs to take into account concrete details about the computational task at hand. While it is true that gradients can take up substantial memory space, one can resort to “fused optimizer” as a way to decrease the peak memory footprint of each iteration, which consists in applying parameter update for layer $i$ before continuing backpropagation on layer $i-1$. This would however only be possible if no gradient accumulation is needed, otherwise, gradients would need to be stored between different micro-batches. However, if gradient accumulation is needed, then mixed-mode random walk would also need to preserve gradients between micro-batches, rendering the argument ineffective. The only remaining argument would be to show that the memory reduction of mixed-mode moonwalk is such that we can increase the batch size to a point where no gradient accumulation would be needed, while it would still be necessary in standard backpropagation or RevBackprop. Overall, this makes the memory reduction argument of mixed-mode moonwalk fairly weak.

3. In the paragraph Memory complexity with checkpointing in section 5, the notation $n$ representing a bound on each layer’s size should be named differently, as it refers to the dimension of $x_0$ in section 3.1. Furthermore, $c$ must be lower than $\frac{L}{c}$, and it is not clear to the reviewer how the best tradeoff $c$ stated at line 309 can be guaranteed to be lower than $L$, unless $M_x + M_{\theta} \leqslant n$ or at least $M_x + M_{\theta} = \mathcal{O}(n)$.

4. The authors show a superior numerical stability with the use of a TanH function. First, the reviewer does not understand why the authors resort to this activation function w.r.t model performance. In absence of a better justification for the use of that function, it looks like this activation has been chosen to favor their method against others. Please provide an argument for the relevance of this activation function. The first reviewer’s guess is that authors wants to use a reversible activation function, but they could as well resort to softplus or leakyRelu as an alternative to ReLU. Furthermore, the reviewer is not aware of reversible residual networks whose activation function is applied on both streams. Instead, in most reversible architectures, the non-linearity are all embedded into the function $\mathcal{F}$. As far as the reviewer understands, RevBackprop has a lower memory footprint than Mixed-mode Moonwalk, which means that the numerical stability argument would be the only remaining argument to justify the use of the proposed method. Thus, it would be nice to elaborate to what extent this numerical stability advantage would be crucial within modern architectures.

5. Line 417: why do the authors pad the input with zeros in the channel dimension ? This increases the input dimension by ~2.6 or 6. The reviewer would like some explanation. Is it solely to experiment with different kinds of input size "synthetically" ?

6. Albeit authors do investigate time and memory tradeoff, their focus on simple datasets and architectures does make it easier to analyze their method in detail on the presented examples, but it would be nice to summarize the practical potential of the proposed method. For example, it would be convenient to summarize the scaling potential offered by the proposed methods against RevBackprop on use cases where memory footprint is the crucial limiting factor. While this might not represent substantial text, it might be highly valuable to the reader.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: - In Section 4.2, the authors mention that Backprop typically retains some information that could be discarded. It is not due to Backprop itself but rather to optimize computation pipeline utilization (see, e.g., [1]). I am curious whether, when the authors consider pipeline scheduling efficiency across multiple data batches for Moonwalk, similar retention of additional information might occur, as observed in Backprop.

- The tests on time and memory overhead require more careful execution. In Figures 4 and 7, when each block contains three layers, the time consumption deviates from a monotonic trend, which seems unexpected and lacks explanation.

- Figures 2-7 contain instances where figure captions do not match the content, and references in the text are incorrect or entirely missing, significantly hindering the readability of Section 6.

- The experiments involve up to five baselines, so why do most results include only two or three of them? Except for the vanilla forward algorithm, which may be prohibitively costly, the remaining methods should be testable within a reasonable timeframe.

- The learning curves in Figures 3 and 7 lack key hyperparameter descriptions, raising concerns about whether the results are consistent under alternative experimental settings.

>[1] Narayanan, D., Harlap, A., Phanishayee, A., Seshadri, V., Devanur, N. R., Ganger, G. R., ... & Zaharia, M. (2019, October). PipeDream: Generalized pipeline parallelism for DNN training. In Proceedings of the 27th ACM Symposium on Operating Systems Principles (pp. 1-15).
soundness: 2
strengths: This paper presents a simple yet practical approach to managing the cost of forward differentiation, highlighting a direction distinct from subspace projection and offering the potential for further exploration. The authors provide a clear and comprehensible description of the methodology and present detailed theoretical properties to demonstrate its advantages.
summary: The paper explores alternative gradient computation strategies in the context of invertible neural networks, leveraging the properties of surjective differentials to reformulate the chain-rule recursion for parameter derivatives. A two-stage semi-forward-mode gradient computation algorithm, named Moonwalk, is proposed to reduce memory overhead in gradient computation. This approach requires multiple forward passes to determine the input gradient and perform forward differentiation recursion. The authors provide a complexity analysis of different gradient computation methods to highlight the potential of Moonwalk. Experiments on classical image classification tasks are conducted to evaluate the algorithm's performance comprehensively.
weaknesses: My main concern is a potential paradox: if activation gradient computation in Backprop is much more expensive than parameter gradient computation, the Mix algorithm’s first step yields minimal savings; otherwise, activation storage costs can become negligible. This could place the Mix variant, which is essential for showcasing Moonwalk's advantages, in an awkward position.

The paper attempts to demonstrate the advantages of the proposed algorithm through experiments from multiple perspectives. However, certain aspects of the experimental setup and presentation are suboptimal, which affects the demonstration of the algorithm's effectiveness. Please refer to the questions section for further details.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 5. While I understand that the paper focusses on exact computation of the gradients, it would be a great addition to discuss more about estimations (like the ProjForward algorithm). There are for instance the forward-only algorithms (Forward-Forward, DFA, PEPITA…). In particular when computing the gradient wrt the input, it seems natural to try to estimate it as in ProjForward using vjp with $k$ random directions. Have you thought that or tried it?
6. Although it seems nice to extend the applicability of Moonwalk to a larger class of functions, I find it hard to get intuition of what this changes in the context of deep learning models. Do you have examples of layers which would be submersive but non invertible?

I remain open to discussion and may improve my grade in the future.
soundness: 3
strengths: - The proposed algorithm is mathematically founded, and seems intuitive and natural.
- Moonwalk, and its Mixed version, are new interesting algorithms that propose different time/memory complexity tradeoffs from existing alternatives to backpropagation.
- The time and memory analysis is thorough, and is done for the proposed algorithms and several other existing algorithms.
- Empirical benchmarks on RevNet models validate the claims, showing noticeably less memory usage compared to backpropagation.
summary: The authors introduce Moonwalk, an algorithm for computing the gradient of an invertible network (or more generally submersive networks). Compared to backpropagation, it does not require storing intermediate hidden states in memory, thus allowing for a more memory-efficient training, at the cost of an increased computation time. A variation of Moonwalk is also proposed to have the same time complexity as backpropagation but still with noticeable memory savings.
weaknesses: 1. The main issue I find in Moonwalk is computing the gradient w.r.t. the input, which is very expensive. As seen in Fig. 4, it is a couple orders of magnitude slower than backpropagation. This is extremely unpractical. However, the authors are aware of this limitation and propose the Mixed variant to mitigate it, which I find much more convincing.
2. Unlike what is stated in the abstract ("Finally, we showcase the robustness of our method across several architecture choices."), the algorithms are only tested on RevNet with 3 blocks. Only the number of layers in the blocks, the number of input channels, and the activation between blocks, are changed. I could be nice to see Moonwalk work on other inversible architectures, which would in particular make the time and memory benchmarks more convincing.
3. The algorithm is only applicable to very specific architectures, which are rarely used in practice. But this is only a minor weakness, as the use of invertible networks could actually be motivated by algorithms like Moonwalk.
4. I believe the captions must be above the tables according to the ICLR template.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: **Questions**
- The possible differences between $M_x$ and $M_θ$ are unclear in the text. What are these values in some standard layers? ($W$ and $X$ for a linear layer for instance). This makes it very unclear if one value can overshadow the other in practice.
- Have the authors considered doing ProjForward to compute the input gradient rather than computing it explicitly?
- line 276: why omit it when  $M_θ$ (for a linear layer) will depend on the batch size value but not $M_x$?

**Minor details**
- Wrong use of citet/citep on some occasions (line 500 for instance).
- Figures 2-4 captions are hard to read. It is also very hard to compare the non-forward methods in Figure 4.
- The caption of Figure 3a is wrong, and the one of Figure 4 lacks a dot.
- line 176 misses a space.
- The use of the term "Pure-Forward Moonwalk" or just "Moonwalk" varies during the paper.
soundness: 2
strengths: - This method is novel and allows for an exact forward-mode gradient computation, as soon as the input gradient is computed.
- In the case of pure-forward Moonwalk, the method improves over standard Forward-mode AD, but not necessarily over ProjForward, depending on its convergence.
- The paper is well-written and clear.
summary: This article presents an alternative to backpropagation named Moonwalk. It first computes the loss gradient w.r.t the input, either using forward-mode automatic differentiation or a standard backward pass. Then, the parameter gradients are computed during a forward pass, using reversible (or submersive) layers to compute the vector-inverse-jacobian of each. The authors implement their method in a RevNet, and compare its theoretical and practical memory and time overheads. They also experimented on the stability of the standard RevBackprop approach to reversible networks.
weaknesses: - The practicality of the method is very limited and over-claimed for the Mixed approach. RevBackprop is more effective on all points (see next weakness regarding the stability of RevBackprop). The use of a big O notation to hide the (approximately) doubling of the time complexity in Moonwalk is not acceptable: for almost all methods (except Forward and Forward-Moonwalk), the only difference in execution time is up to a constant, which matters. I found it surprising to claim that RevBackprop is not applicable to submersive networks, considering that all invertible networks are submersive (line 130), and the authors never give an example of a submersive network that is not invertible, considering only a RevNet. Since this is the only real advantage over RevBackprop, it requires more justification and a real network example.
- I disagree with the authors over their analysis of stability in RevBackprop. They only provide a single seemingly hand-picked example of numerical instability of RevBackprop when adding a tanh activation function. This is a very rare activation, not used in RevNets or Transformer-based models like Revformers. Other works like i-revnet et RevVit showed that stability was a non-issue and that the approximation error was equal to $10^{-6}$ in the worst cases. Without more precise examples, this example is not a demonstration of the instability of RevBackprop. Furthermore, it remains to be made clearer why computing the vector-inverse-jacobian should be more stable than the inverse function directly used in RevBackprop, since it also uses the inverse function, as explained in line 141.
- References to activation checkpointing need to be updated with more recent ones which are more general and relevant, such as [1,2].
- Despite the improvement over standard Forward-mode AD, the method is not compared in practice to the convergence of ProjForward, making it hard to choose one over the other. 
- The method is limited to reversible architectures (and submersive ones), where RevBackprop is already available.



[1] Efficient rematerialization for deep networks. NeurIPS 2019, Kumar et al.
[2] Rockmate: an efficient, fast, automatic and generic tool for re-materialization in pytorch. ICML 2023, Zhao et al.

Meta-review: The paper develops Moonwalk, a forward-mode differentiation method for *submersive* networks, i.e. networks with surjective Jacobians. In particular, all invertible networks are submersive. The method requires computing the gradient with respect to the input first, but then the rest of the computation is efficient in forward mode. In theory, forward-mode differentiation does not require storing the network activations, and can substantially reduce the memory footprint of gradient computation compared to standard backpropagation. Apart from computing the gradient with respect to the input, Moonwalk also matches backpropagation in terms of time complexity. In the experiments, the authors compare the proposed method to several baselines (Backprop, RevBackprop, ProjForward) on both runtime and memory, with the mixed method (use Backprop for input gradient and Forward mode for the rest) showing good results.

Strengths:
- Novel forward propagation method with interesting properties.
- The authors identify a broader class of applicability for their method compared to RevBackprop: submersive networks.  
- In the experiments, the Moonwalk shows good performance on both runtime and memory compared to baselines.

Weaknesses:
- The fully-forward variation of MoonWalk is still impractical (five days to train on CIFAR-10)
- The mixed version involves Backprop to compute gradient with respect to the input. The authors show that in practice the backprop for just the input gradient can be cheaper in terms of memory compared to full backprop.
- Compared to RevBackprop, the main advantage is the applicability to submersive but non-invertible networks. However, the experiments focus on invertible RevNet models.
  + For the comparison to RevBackprop, the authors show a numerical instability on the RevBackprop; however, the setup involves an unusual activation function that may be specifically chosen to increase RevBackprop instability
- The method is not generally applicable, it requires the network architecture to be submersive.

Decision recommendation: The paper makes an interesting methodological contribution. Mixed Moonwalk is potentially an interesting method for submersive but non-invertible models. However, currently the experiments fail to highlight the performance on that model class, as they focus on invertivle models. For invertible models, it is unclear if the method has advantages over RevBackprop. I believe the authors should emphasize results on submersive but non-invertible models. In the current form, I am leaning towards rejecting the paper.

Decision: Reject