Moonwalk: Inverse-Forward Differentiation

Authors: Dmitrii Krylov, Armin Karamzade, Roy Fox
Primary area: optimization
Keywords: Forward-mode, Forward Gradients, Automatic Differentiation, Projected gradients, Invertible Networks, Bijective Networks, Jacobian-Vector product, Alternatives to backprop, forwardprop, memory-efficient deeplearning, JAX

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 and right-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 and right-invertible networks in computation time comparable to backpropagation and using significantly less memory.

Decision: reject
Reviewer scores: 5, 5, 6, 3

Reviews and discussion:
<thread>
creation date:
1730382322570

thread_type:
official_review

<thread_object>
creator: reviewer_9rvw

<content>
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.

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.

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730491714426

thread_type:
official_review

<thread_object>
creator: reviewer_gc2t

<content>
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.

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.

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730653141201

thread_type:
official_review

<thread_object>
creator: reviewer_i14z

<content>
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.

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.

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.

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).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730895628802

thread_type:
official_review

<thread_object>
creator: reviewer_zprr

<content>
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.

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.

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.

questions:
Could the author address the weaknesses in general ?
The reviewer would appreciate a detailed explanation on weakness 2 and 4 specifically.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732612382101

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Comments to all reviewers

comment:
We thank all reviewers for their insightful comments and the opportunity to improve our work.

First, we would like to highlight the primary benefit of our method: its applicability to submersive networks. We acknowledge that the previous version of the paper lacked concrete examples of networks that are submersive but not invertible, where reversible backpropagation is not applicable. To address this, we have included illustrative examples in the revised manuscript to clarify these distinctions. 

Second, we emphasize the computational advantages of our method. Algorithms 2 and 3 in the manuscript outline the key differences in the computational graph. Unlike backpropagation, which requires storing all intermediate variables, our method (Moonwalk) reduces storage requirements. Specifically, for gradient computation, Moonwalk only requires the storage of signs for activation functions, providing a clear advantage in terms of memory efficiency.

In the updated version of the manuscript, we added:
Architecture for showcasing a submersive non-invertible network with algorithms to train them using Moonwalk. Algorithm 3,4,5

The section about using submersive networks for training. Sections 6.6 and 6.7.

Updated captions on figures 2-4.

We added examples of SVD and Gaussian Elimination for effective matrix inversion.

We updated references to acknowledge new work in the field of checkpointing.
</content>
</thread_object>

</thread>


Metareview:
metareview:
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.