Moonwalk: Inverse-Forward Differentiation

We thank the reviewer for their constructive feedback and valuable suggestions, which helped us identify areas for clarification and improvement

1. "The practicality of the method is very limited"

We thank the reviewer for his comment. We would like to add a clarification that all invertible networks are submersive but the opposite is not true. Not all submersive networks are invertible. We added an algorithm with an architecture example to show that we can train a submersive network, which is not invertible. In the case of Linear layers that reduce dimensionality, they are not invertible, but submersive, we provide an algorithm and code snapshots to train such networks. We added more justification, but the main point is that we can use Mixed Moonwalk with linear layers and convolutions, whereas reversible can not operate under these constraints. We would clarify that the main point of our work is to show a novel method for computing gradients. First, it allows us to train the network where Reversible fails to do so. We added Algorithms 2 and 3 to showcase the difference between Mixed and Backprop. We would also to clarify that big O notation to show that the theoretical properties of Moonwalk are similar to backpropagation, whereas optimal performance is heavily dependent on optimization.

2. "I disagree with the authors over their analysis of stability in RevBackprop"

Thank you for raising this point, we do agree that evaluating one activation function is not sufficient to show a comparison between methods. We would like to highlight a few things:

There is a line of work [1] that showcases where reversible networks fail with some coefficients. We did not investigate this particular example, but hypothesize, that moonwalk can solve this issue. We will add in the updated draft more experiments similar to [1].

3. References to activation checkpointing need to be updated with more recent ones which are more general and relevant, such as [1,2]. Thank you, we updated the current draft

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

We would like to highlight that ProjForward does not produce true gradients, and in all experiments, it failed to converge.

5. The method is limited to reversible architectures (and submersive ones), where RevBackprop is already available.

Please, see our point about extension to the subclass of non-invertible subversive networks. (Linear layers and Convolutions)

Questions

1. “The possible differences between “

Thank you for the comment! An example could be shown from algorithm 2 when we need to store For M_x only signs for grads, but for M_\theta we also need activation itself.

2. “Have the authors considered doing ProjForward to compute the input gradient rather than computing it explicitly”

Yes, we did some experiments, without any success at this point. We believe that if we estimate it with multiple directions then when we multiply with inverse matrix the error would quickly grow. We think that more constraints are needed in order to solve this problem. In general, we think that this idea might be another excellent point to prefer moonwalk instead of backprop/reversible.

3. line 276: why omit it when Thanks for pointing out. Both M_x and M_\theta would depend on batch_size. Please, refer to algorithms 2 and 3, in case of adding Batch size, we will have to store N * signs_fo_grads (Which corresponds to M_x), and M_\theta would correspond to variables that we have to store for algorithm 3. In general, adding batch size as an additional parameter would benefit Moonwalk more than backpropagation.

Minor details Updated in the new version.

References: [1] Liao, Make Pre-trained Model Reversible: From Parameter to Memory Efficient Fine-Tuning

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.

Thank you for your point. We added algorithms 2 and 3 to show that our method in general is more efficient than backprop and works on submersive networks, whereas reversible is not compatible with such architectures. We showcase that we only need to store signs of gradients rather than the entire variables

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.

Thank you for the comment! We added algorithm 2 and code snapshots to show that our method works on submersive networks with Linear Layers and 1d Convolutions with Code examples. We would also to highlight that reversible - backprop can’t work with such architectures and we are more memory efficient than backprop in such scenarios.

3. The algorithm is only applicable to very specific architectures, which are rarely used in practice. However this is only a minor weakness, as the use of invertible networks could actually be motivated by algorithms like Moonwalk.

We do agree, but as we show almost any Linear layer with an output size lesser than the input could be right invertible with Gaussian elimination (We also added an algorithm) which is more efficient than using SVD. Please, refer to the updated version of the manuscript where we included new types of layers.

4. I believe the captions must be above the tables according to the ICLR template.

Thank you, we updated the tables.

Questions:

1. 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 random directions. Have you thought that or tried it?

Thank you for the great suggestion! We indeed tried that option but without much success. We believe that if we estimate it with multiple directions then when we multiply with inverse matrix the error would quickly grow. We think that more constraints are needed in order to solve this problem. In general, we think that this idea might be another excellent point to prefer moonwalk instead of backprop/reversible.

2. 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?

Thank you for your comment, we did include the examples. Algorithm 2 showcases a network with linear layers. Basically, any network with linear layers or 1d convolutions that have output is smaller than its input would be submersive, but not invertible. Another point is that we need to add constraints like making linear layers upper triangular with ones on the main diagonal to make it stable and invertible with Gaussian elimination.

I thank the authors for answering my questions.

While overall I like the proposed method and find it elegant, I do not find the experimental section convincing enough.

I am keeping my score as it is.

1. "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."

Thank you for your comment. We would like to highlight using an example of an algorithm 2 vs 3 for submersive networks, which we included in the appendix. For further clarification, we also included an example of training such a network with linear layers or with 1d convs. In case when activation gradient computation is very expensive, we show that we need only to store signs, whereas for backprop we have to store the entire variable. This basically means that we are reducing memory footprint from fp16/32 to binary i.e 16x/32x more memory efficient. We would like to highlight, that it would be significantly more efficient for convolutions. We also include a code fo efficient matrix inverse (Avoiding using SVD) with Gaussian - Elimination.

2. "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."

We do agree that some of the experiments were not optimal. We would like to also clarify that the main point is not to show numerical stability, but rather to show that we can train submersive networks, whereas reversible backprop can not do that.

Questions:

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

Thank you for the reference. We do agree that by trading-off memory/computation we can come up with different graphs based on user needs. We would also like to highlight that based on algorithm 3 (appendix) we would only need to store signs instead of full variables for computing gradients wrt input. We would also like to highlight the potential benefit of Moonwalk in the context of multi-gpu training. If we split the model across multiple GPUs we would experience a bottleneck, but in the case of moonwalk the first phase would be faster, and the potential bottleneck would be smaller.

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

We do agree with the reviewer. We are gonna address this point in the draft. The main problem is that JAX on CUDA does not guarantee optimal memory allocation, some times it can construct graphs with less memory, but with more computation.

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

Thank you for the comment, we clarified the figure notation in the updated version.

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

We did not include projForward mostly because of its inaccurate gradient estimation. We would like to highlight that this method does not produce accurate gradients, and in all our experiments it failed to train an end-to-end network.

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

Thank you for your point. We will add hyperparameters to the appendix.

We sincerely thank the reviewer for their constructive feedback and valuable suggestions, which have greatly helped us refine our manuscript and address potential ambiguities.

Firstly, we would like to clarify that the primary focus of our work is to introduce a novel method for computing gradients using forward mode, specifically designed for submersive networks. Please, see our message to all reviewers. To underscore the significance of our approach, we have added an illustrative example in the revised manuscript of a submersive network that cannot be trained using reversible backpropagation but can be successfully trained with Moonwalk.

To clarify further, any linear layer with output dimensionality $k \leq n$can be considered a submersive layer. Notably, reversible networks are incapable of effectively training such networks. Our primary contribution lies in introducing a novel method for efficiently training submersive networks, overcoming the limitations faced by reversible networks.

1. “The dimension of parameters are not mentioned”:

Thank you for your comment. For submersive networks, in general, the input size is $n$, but for all subsequent layers, the output size of each layer can be $k \leq n$. For simplicity, we assumed that all layers have a fixed parameter size of $d$ and input-output size of $n$. We will clarify this assumption in Section 3.1.

2. “An example of architecture”.

Here, we provide a simple yet widely-used architecture as an example. Technically, any sequence of linear layers with decreasing dimensionality can serve as a representative case. To illustrate this, we have included a code snippet that highlights the architecture.
The key point is that in standard backpropagation, the entire activation ($z_2$) must be stored. In contrast, with our proposed Moonwalk method, when computing $h_0$, we only need to store the signs of this activation. If a function like leaky-ReLU is used, this allows us to reduce storage from $fp16$ to binary for every number in $z_2$, leading to a potential 16x memory savings during this phase.

2. "fused optimizer"

Could you please clarify which paper you are referring to? If you are discussing in-place updates, there are several potential issues with such networks that can negatively impact convergence performance. Specifically, if the gradient for $L_{n_1}$ depends on $W_n$, updating $W_n$ prematurely can result in incorrect gradient estimation, which could be a case if we talk about the same weights in the block with residual connections, for example.

3. Could you please clarify this point “c should be lower L/c”?

4. "The authors show a superior numerical stability with the use of a TanH function"

The tanh activation function is just one of many examples that illustrate the limitations of reversible models. We are not the first to point out instability in reversible models under certain conditions, as discussed in [https://arxiv.org/pdf/2306.00477]. However, we acknowledge that using a single activation function is insufficient to fully demonstrate the broader stability issues. We would also like to point out that, it is not the main point in the comparison, since we can operate on a wider class of networks. Please, see our message to all reviewers.

5. "Line 417: why do the authors pad the input with zeros in the channel dimension"

The primary goal is to increase the effective dimensionality of the model. A key limitation of bijective invertible networks is that they require the same input and output size for every layer. Here we use reversible architecture as shown in RevNet paper. For example, in an MNIST classification problem with an input size of 64, the network would be constrained by having all layers fixed at size 64, creating a bottleneck. To address this, one approach is to pad the input with zeros, effectively increasing the network’s capacity. Another way to conceptualize this is by projecting the input into a higher-dimensional space. In general, our method is more flexible, since we can have contracting layers.

6. "Albeit authors do investigate time and memory tradeoff"

We want to emphasize that our approach enables the training of submersive networks, which is not possible with RevBackprop. To illustrate this, we have included an example along with a code snippet. Additionally, our method outperforms both standard backpropagation and checkpointing when applied to submersive networks—a broad and versatile class of network architectures.

I thank the authors for their detailed response. I would like to make a few more comments:

1. About parameter dimensionality.

I did not keep the original version of the article so I am not sure what have been updated in section 3.1.

2. About memory savings.

I thank the authors for their simple yet relevant example of an architecture where mixed-mode moonwalk offers substantial memory savings compared to standard backpropagation. Still, memory-wise, the advantage is not clear to the reviewer when considering reversible architectures since reversible backpropagation does not require to store intermediate activations. The memory advantage is however straightforward for submersive architectures.

The reviewer however has some concerns about submersive architectures. While it might be feasible to prove that random $k \times n$ matrices have a high probability of being subjective, some computational tasks might collapse the rank of the matrix below $k$. In practice, matrices are rarely low rank, but often ill-conditioned with many singular values close to zero, which might affect numerical stability when computing a right-inverse. While a thorough ablation of each of the problem associated with numerical stability would be expensive, it would still be nice to spend some time either in the main body, either in the appendix, to emphasize how much wider is the class of submersive networks compared to the class of reversible networks, especially with respect to the current literature. For example, reversible architectures often require to maintain a constant dimensionality; this makes it non-trivial to adapt ResNets to fully-reversible architectures due do downsampling operations.

2. About fused optimizer.

The reviewer were not aware of papers studying fused optimizer specifically, but this reference seems to be the one I would be looking for. Fused optimizer were proposed to optimize memory consumption during training in the Apex library. As reviewers correctly points out, updating the weights too soon might lead to incorrect computations for subsequent gradients, but it is often possible to apply this update way before the end of the full end-to-end backpropagation. Nevertheless, the reviewer does not believe that this is a major concern for mixed-mode Moonwalk specifically and does.

3. "$c$ must be lower than $\frac{L}{c}$"

I am sorry, this was a typo. I meant "$c$ must be lower than $\frac{L}{n}$".

4. About numerical stability.

The reviewer slightly disagree with the author with respect to the importance of their numerical stability claim. The reviewer acknowledge that their method is applicable to a wider class of model, but they did not focus their experiments on submersive architectures that are not reversible. Therefore, the reviewer took the claim of numerical stability seriously. The reviewer think that if numerical stability is such a serious issue, it should be better explained and supported with references or experiments in the paper.

My review update: Overall, the authors answer one of my main concern, which was to emphasize the memory reduction offered by mixed-mode moonwalk compared to backpropagation. The fact that mixed-mode moonwalk is applicable to submersive networks that are not reversible should be emphasized more in the experiments; for example, there is no mention of the cost of the inversion and how it affects training. The reviewer do not see many advantage over reversible backpropagation for reversible networks, apart from the improved numerical stability claim that the reviewer do not find convincing yet.

I am increasing my score from 3 to 5 given that the memory constraint is the main limiting factor. Should the author emphasize better the practicality of submersive architecture or the relevance of improved numerical stability, I'd be willing to improve my score further.