Exact, Tractable Gauss-Newton Optimization in Deep Reversible Architectures Reveal Poor Generalization

We thank the reviewer for the thorough and detailed comments. We provide answers to questions and weaknesses below.

I am wondering how transferable the results actually are to architecture that are typically used and do not include random feature projections such as those used in the proposed work.

This is a great question. As the reviewer noted, we used inverted bottlenecks to “ensure linear independence”, i.e. to ensure that the model is adequately overparameterized such that the efficient generalized inverse we propose is valid. In other words, inverted bottlenecks ensure that the scalable GN weight updates (Eqs. 16-17) do implement gradient flow in function space (Eq 3; the essence of the GN method), such that our results are not potentially confounded by broken theoretical assumptions. Nevertheless, we agree that our GN updates can still be applied in the absence of inverted bottlenecks: even though they may not enjoy the same theoretical guarantees, they could still lead to similar training behaviour and it is worth investigating this empirically. We did this on the CIFAR10 dataset, with results shown in Figure 2 of the one-page rebuttal PDF. We followed the same experimental procedure as in the paper, only removing all inverted bottlenecks, and we tuned the learning rate for each optimizer. In the full-batch setting, GN is still performing much better than Adam and SGD. In the mini-batch setting we observe a very similar trend to what is shown in our main paper: GN leads to an early saturation of the loss, which instead does not appear in Adam and SGD. We thank the reviewer for this suggestion, and we will include these results in the final version of the paper.

[...] are there any comparable results for the setting of squared loss settings such as on UCI regression data sets?

We thank the reviewer for the valuable suggestion to apply our method directly to regression tasks with a squared loss. We have added some new experiments on two of the UCI Regression datasets (Wine Quality & Superconductivity), with results shown in Figure 4 of the one-page rebuttal PDF. These results corroborate our main findings in the cross-entropy / classification setting: in the full-batch case GN is significantly faster than SGD and Adam, while in the mini-batch case there is an apparent stagnation of the test and train losses under GN.

We hope that the additional analyses address the reviewer's concerns and that they may consider raising the score of our paper.

We thank the reviewer for the quick reply. We are glad our rebuttal could help clarify doubts, and that the reviewer is happy to raise the score. We however notice that the previous score of 6 has not been updated, so we kindly ask the reviewer to edit the score.

We thank the reviewer for their positive appraisal of our work, and for the constructive feedback. Below, we provide answers to their questions and address the weaknesses they have identified.

The experiments section could be more full; for example, exploring the batch size dependence more fully (or, extending the full batch examples to more datapoints). I would also be interested to see how GN training performs in a setting where SGD works at least as well as ADAM.

We followed the reviewer's suggestion to extend the full-batch example to 2048 samples, which can be found in Figure 1 in the one-page supplemental PDF. We used the same architecture and experimental procedure as in the main paper, including tuning the learning rate for each optimizer. The overall trend remains unaltered: in the full-batch setting, GN is significantly faster than both Adam and SGD; in the mini-batch setting, GN decreases the training and test loss faster initially, but then causes them to saturate early, while they continue to decrease for Adam (these effects on train/test losses are also reflected in classication accuracies). In general, our experiments have not revealed any important difference in behaviour when changing the batch size. We thank the reviewer for highlighting this aspect, and we will include this new result in the final version of the paper.

The reviewer will also be interested in our response to Reviewer 9txE, where we show results of additional experiments on two of the UCI regression datasets (Wine Quality and Superconductivity). For these experiments, we used a squared error loss to complement our other classification experiments which used a cross-entropy loss. The results are shown in Figure 4 of our one-page supplemental PDF, and they do corroborate our main conclusions. Interestingly, Adam and SGD perform similarly on the Wine Quality dataset in the mini-batch setting -- a situation which the reviewer wished we had explored --, and even in this setting, GN exhibits similar overfitting characteristics to those reported in our original submission.

What happens in L2 regularization is added to the experiments?

That is a great question -- L2 regularization being a very standard way of mitigating overfitting, it had also occurred to us that it could be important in addressing the overfitting behaviour of GN that we describe. Experiments with L2 regularization can be found in Appendix K of the main submission; following the standard prescription of AdamW (Loshchilov & Hutter, 2017) we implemented L2 regularization as weight decay in all three optimizers. We found that adding L2 regularization has almost no effect on the overfitting behaviour of GN (we have tuned the amount of weight decay over 5 runs with different order of magnitudes). We also tried additional forms of regularization in Appendix L and came to similar conclusions.

It would be helpful if more intuitions about the form of the pseudoinverse of J are brought into the main text

We agree that providing higher-level intuitions about the form of our generalized inverse of J would be useful. We had attempted some of that already through our word-only description of the pseudoinversion on lines 234-240, but we will rewrite this to be clearer and more specific. In addition to that, we will mention that, in the overparameterized setting, the Moore-Penrose pseudo-inverse would determine a parameter update with minimum Euclidean norm, among those providing descent in function space. Instead, the update given by our right inverse does not have minimum Euclidean norm, but it may minimize a different type of norm that is currently unknown and will be the subject of future studies.

We hope that the additional analyses address the reviewer's concerns and that they may consider raising the score of our paper.

We thank the reviewer for their time reviewing our paper -- we are glad they found it clear to follow.

The reviewer wondered what is gained by replacing the Moore-Penrose pseudoinverse by a generalized inverse. To the best of our knowledge, there is no known tractable way of computing the Moore-Penrose pseudoinverse for deep nonlinear networks. This not only precludes its use in applications, but also makes it difficult to investigate the properties of GN at scale. Thus, the main “gain” of our generalized inverse is that it offers a computationally tractable expression for a weight update that implements gradient flow in function space (the core principle underlying the GN method, as summarized around Eq. 4).

The reviewer rightly noted that our generalized inverse applies to reversible network archictures only, which we understand can appear restrictive. Coupling layers (Dinh et al, 2015; the form of reversible blocks that we use here) were originally introduced as a technical work-around for efficient computation of Jacobians and inverses (the same reason we use them here). They have since been very influential in the area of generative modelling (normalizing flows, diffusion models, ...) where they achieve SOTA results, and in other settings for learning flexible isomorphisms. The reversible vision transformer (Mangalam et al., 2022) also achieves near-SOTA results, and Liao et al (NeurIPS 2024) have shown that even large pre-trained LLMs can be made reversible by inserting specific adaptors, leading to memory-efficient fine-tuning. In fact, we believe that extending our efficient GN updates to these other classes of models, in addition to developing new theory to address the overfitting behaviour we have uncovered, is a promising future direction.

With the CKA results, it is expected that the Gauss-Newton learning schemes, which becomes enabled by the reversibility of the model, is more effective when the model gets deeper. Then, what happens with the performance when the reversible model gets deeper or shallower?

As the model grows deeper, we find no significant performance increase for GN training, whereas Adam and SGD improve by a small amount towards convergence. As the model gets shallower, we find that performance drops in all methods by a similar amount.

The reviewer's mention of CKA analysis in this context prompted us to think about alternative ways of studying the emergence of representations in RevMLPs. In the paper as it currently stands, our CKA analysis is performed separately for each half-coupling layer (i.e. separately for Eq. 14 and Eq. 15). We wondered what the CKA similarity would look like at the level of the 'full' coupling layers (i.e. the hidden representation jointly contained in the concatenation of $x_\ell^{(1)}$ and $x_\ell^{(2)}$ for each block $\ell$). We computed these CKA similarities (w.r.t. before-training representations at initialization) and found that they remain very high throughout training (Figure 3 of one-page rebuttal PDF). These results show that the large (albeit late) changes in half-coupling-layer representations observed in our original analysis were in fact coordinated between the two half layers in each block, such that the compound representation which they jointly hold does not actually change. This finding is of course completely in line with our main conclusions, and to some degrees even simplify the story (we no longer need to appeal to the fact that the observed changes in (half-layer) representations happen after GN has already reached its final performance). We are grateful to the reviewer for prompting us to think along those lines.

Is assuming RevMLP, a reversible neural network, setting feasible to use? It seems that the accuracy with CIFAR-10 dataset has a scale of 70%, which is quite lower than our consensus.

The accuracy we obtained using RevMLPs on CIFAR10 is in fact on par with classical MLPs (~60-70% test accuracy without data augmentation). MLPs, whether reversible or not, are just very prone to overfitting on this type of vision problems. Prior to submission, we did experiment with our own reversible version of MLP-Mixer (Tolstikhin et al, 2021; essentially using RevMLPs instead of MLPs in each mixer block), for which we were able to derive an analogous Jacobian generalized inverse. For these rev-MLP-Mixers, performance on CIFAR10 was overall much better ($> 90\%$ test accuracy with Adam), similar to what was reported in the non-reversible version, showing that the specific reversibility property does not really affect performance. Nevertheless, we had found that GN displayed identical overfitting properties in rev-MLP-Mixer as it did in plain RevMLPs, but training runs took much longer. For this reason we decided to simplify both the exposition and our experiments by focusing our paper on plain RevMLPs. We will add a paragraph of discussion on this in our revision.

I'd like to confirm that the reason of overfitting with respect to each minibatch. What happens if the minibatches are re-shuffled after each iteration?

In our experiments in the minibatch setting, we can confirm that the entire dataset is re-shuffled at the beginning of each epoch before getting divided into minibatches; thus, each minibatch is (statistically) unique; what we show is that GN tends to overfit to each minibatch, and this isn't rescued by learning slowly over many randomized minibatches.

How is the V_t defined? I assumed that it is a class of random matrices, such as random Gaussian or random Fourier transform matrix.

Apologies for this omission; we have now added a full description of how these inverted bottleneck weights are drawn: i.i.d. from a normal distribution (as the reviewer suspected) with zero mean and variance $2/d$ (the input to those have dimensions $d/2$, hence the factor of 2).

We kindly notify the reviewer that the discussion period is reaching an end. We believe we have addressed all the raised concerns, which has also helped us improve the paper, and we would be happy to receive feedback and answer any remaining doubt