Moonwalk: Inverse-Forward Differentiation
Forward-mode;Forward Gradients;Automatic Differentiation;Projected gradients;Invertible Networks;Bijective Networks;Jacobian-Vector product;Alternatives to backprop;forwardprop;memory-efficient deeplearning;JAX
Yes
No ethics review needed.
Could the author address the weaknesses in general ?
The reviewer would appreciate a detailed explanation on weakness 2 and 4 specifically.
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.
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.
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.
Yes
N/A
No ethics review needed.
- 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).
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.
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.
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.
Yes
No ethics review needed.
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.
- 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.
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.
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.
Yes
No ethics review needed.
**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.
- 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.
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.
- 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.
Moonwalk introduces forward-mode differentiation with a vector-inverse-Jacobian product, efficiently computing true gradients in invertible and right-invertible networks with time complexity like backpropagation and significantly reduced memory.
@inproceedings{
anonymous2024moonwalk,
title={Moonwalk: Inverse-Forward Differentiation},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=97dJ3Jp5P4},
note={under review}
}
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.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Forward-mode Forward Gradients Automatic Differentiation Projected gradients Invertible Networks Bijective Networks Jacobian-Vector product Alternatives to backprop forwardprop memory-efficient deeplearning JAX
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/1a3dc608a6bf354a951bb2d9d920426cd709bf94.pdf
optimization
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Moonwalk: Inverse-Forward Differentiation
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission