A geometric framework for momentum-based optimizers for low-rank training

We thank the reviewer for their constructive comments and appreciate that they found the paper comprehensive, well structured and that it bridges a gap in the literature.

1. (W1) We fully agree and have added the Algorithmic descriptions for both methods in the Appendix in the revised manuscript and will upload the version as soon as the review process allows it.

   For completeness we point out the differences of the methods below.

   • DLRT w/o projection does not perform the projection update of $S_{\mathcal{V}}$ in Line 7 and 12 of Algorithm 3 in the supplementary material. Thus it creates an inconsistency of the coefficient representation $S$ and the corresponding optimizer states, since they are now spanned by different bases.
   • LoRA pretrain essentially applies Algorithm 4 to the LoRA factors $A,B\in\mathbb{R}^{n\times r}$, where $r$ is fixed, as stated in Line 244.

2. (Q1) The Dynamical Low-Rank Adam method is the main result of the paper, as the (full rank) Adam optimizer is currently the most widely used training method for neural networks - to the best of our knowledge. Thus we focused our evaluation on this setting.

   For completeness we added VGG16 on Cifar10 with SGD using the same hyperparameters as in the experiments in table 1. We choose momentum=0.9 and train a (full-rank) baseline, LoRA pretrain and Algorithm 1. The compression rate of LoRA pretrain is fixed to match the final compression rate of Algorithm 1. We observe similar performance of Algorithm 1 to the Baseline as we saw for Algorithm 2 to the Adam Baseline, whereas LoRA pretrain exhibits a slight drop in accuracy.

3. (Q2) In our experience, and as investigated in recent literature [a], Vision Transformers train relatively slow on small datasets, e.g. Cifar10, and exhibit worse performance than ViT models pretrained on large datasets (ImageNet size). We think that encouraging low-rank solutions by training with Algorithm 2 helps eleviate this issue to some extend - explained by an Occam's Razor argument: A low-rank solution expresses a certain simplicity bias imposed on the very general transformer architecture. One can see that the low-rank initialization leads to a fast loss decay in Figure 2b, whereas the validation loss of the full vision transformer stalls first. Interestingly the LoRA loss, also with a low-rank initialization decays fast at first, but then stalls. This might be attributed to the missing convergence guarantees to a low-rank optimum (see Section 2).

   A similar effect is observed in the GLUE benchmark portfolio, where low-rank finetuning methods (LoRA-based and DLRT-based) outperform full-fine tuning of certain sub-benchmarks.

   To the best of our knowledge, this often seen observation is currently not rigorously understood and future research may shed more light onto this phenomenon.

   [a] H. Zhu, B. Chen, and C. Yang. Understanding why vit trains badly on small datasets: An intuitive perspective, 2023.

I thank the authors for addressing my questions and detailed explanation, my rating remains

We wish to thank the reviewer for their thoughtful and constructive feedback. We are encouraged that the reviewer found our algorithm design to be based on a solid theoretical principle and that our experimental results show significant performance improvements. The reviewer raised three main points regarding the motivation, experimental setup, and the interpretation of our convergence guarantee. We address each weakness (W) and question (Q) below.

1. W1 (On the motivation for the proposed gradient flow):

   The reviewer correctly points out that we showed that satisfying equation (4) is a sufficient condition for convergence to a low-rank optimum, but we did not explicitly demonstrate that naive momentum methods can fail to converge to a point fulfilling $P(W)\nabla_W \mathcal{L}(W)=0$. We thank the reviewer for highlighting this, as a more concrete discussion improves the paper's motivation.

   The potential lack of convergence for a naive application of momentum methods can be proven with a simple but illustrative counter-example. Consider a state where the momentum term is zero, i.e., $\mathcal{V}=0$, by choosing the momentum factors as $U_v = -U$, $S_v = 0$, and $V_v = V$. Now, consider any weight matrix $W$ that is not a low-rank optimum, meaning $P(W)\nabla_W \mathcal{L}(W) \neq 0$, but for which the naive update term happens to be zero: $\hat{P}(W,\mathcal{V})\nabla_W \mathcal{L}(W) = 0$.

   A large class of matrices fulfills this. For instance, any $W=USV^\top$ where the gradient is non-zero in the range of V but zero in the range of U (i.e., $U^\top\nabla_W \mathcal{L}(W) = 0$ but $\nabla_W \mathcal{L}(W)V \neq 0$) would satisfy this. This can be easily verified: Since $\mathcal{V}=0$, we have

but

In this scenario, the naive gradient flow equations would identify $(W, \mathcal{V})$ as a stationary point, since $\dot{W}=0$ and $\dot{\mathcal{V}}=0$. However, this point is not a valid low-rank optimum because the true projected gradient $P(W)\nabla_W \mathcal{L}(W)$ is non-zero. In contrast, $(W, \mathcal{V})$ is not a stationary point of (4) and the gradient flow (4) will provably drive the system to a state where $P(W)\nabla_W \mathcal{L}(W)=0$.

This demonstrates a fundamental geometric inconsistency in the naive approach: it can prematurely halt at suboptimal points that do not correspond to true low-rank critical points. We added this counter-example and discussion to the appendix of the revised manuscript to make our motivation clearer and more compelling.

2. W2 (On the experimental setting and runtime comparison):

   The details for the numerical experiments are provided in the Appendix in Table 3 and 5, respectively. We set an epoch limit for all experiments in the paper, where we have chosen the epoch number such that all methods have sufficient iterations to converge. This is a default test case setup.

   To underline that Algorithm 2 does not require a significantly increased training time, we provide a wall clock comparison of finetuning Llama2 7b-chat-hf on BoolQ and PICA benchmarks with standard LoRA [Hu et al., 2021] and Algorithm 2 in the table below. The wall-time is reported for three epochs with batch size 12 and maximal sequence length of 640 tokens on a single NVIDIA H100. We observe singificant outperformance of Algorithm 2 at a negligible wall-time overhead for all test cases.

Hyperparameters are determined by initial hyperparameter sweeps (which were performed for LoRA and Algorithm 2 individually). Hyperparameters for Algorithm 2 are reported below:

Hyperparameters for LoRA:

2. W3 (On the derived error bound):

   We appreciate the reviewer's detailed question regarding Theorem 3. We wish to note that our algorithm can be straightforwardly modified to provably control the term $\widehat \vartheta$ (as underlined in the comment following Theorem 3 in our manuscript in Line 431). In this case, $\widehat{\vartheta}$ becomes a user-determined parameter. We chose to present the main algorithm without this explicit modification because our analysis and empirical results suggest it is unnecessary and adds computational overhead for no tangible benefit. To name an example, we obtain 84.15% (instead of 84.09%) validation accuracy when finetuning Llama2 7b-chat-hf on the BoolQ benchmark when controlling the projection term by choosing $\widehat{\vartheta}=0.1$, which is fulfilled automatically for almost all iterations.

   Intuitively, the influence of this term is naturally limited, as it is expected to diminish during training. This expectation is based on the dynamics of the optimization process: as the optimizer converges to a local minimum, the momentum term $\mathcal{V}$ itself must converge toward zero. This is a direct consequence of Theorem 1 in the Appendix. Consequently, the difference in its projection onto two closely related subspaces will also naturally tend to zero. This intuition is strongly backed by the consistent high performance of our method across all experiments, where this term was not explicitly controlled.

   To perhaps clarify the context of our response, we feel it's helpful to underline the role of Theorem 3 within our work. We included this analysis in the appendix to provide readers with theoretical intuition for our algorithmic design—specifically, to show how our integrator navigates the numerical instability associated with inverting the $S_v$ matrix. It serves as a motivational stepping stone for our final algorithm, which is the paper's primary contribution. Our goal is a novel, efficient method derived from a geometric framework, and the ultimate validation for such a method lies in its empirical performance. We believe the current algorithm, which, as the reviewer noted, shows significant performance improvements, is the strongest version to present, as it achieves these results at minimal cost without unneeded operations.

   We will revise the discussion around Theorem 3 in the appendix to better articulate this interplay between theoretical motivation and practical algorithm design.

3. Q1 (Applicability to matrix completion/sensing):

   Absolutely. While our paper focuses on neural network training, the proposed geometric framework is general. The core of our method is a momentum-based optimizer for problems with low-rank constraints. Applying it to classic low-rank optimization problems like matrix completion and matrix sensing is a promising and natural extension. We hypothesize that our approach could offer similar benefits in convergence speed and stability for these tasks, particularly in challenging, ill-conditioned scenarios.

   We remark that we do not aim for a state-of-the-art matrix completion algorithm as there exists a rich corpus of specialized methods that are tailored to the specific structure of a matrix completion problem and can therefore outperform more general methods in this specific task.

We hope these clarifications, along with the planned additions to the appendix, fully address the reviewer's concerns. We are confident that these improvements will strengthen the paper and hope the reviewer will reconsider their rating.

Thank you very much for the encouraging feedback and for your useful comments and constructive suggestions. We believe the reviewer’s comments and suggestions have helped greatly in improving our manuscript. We will upload the revised version as soon as the review process allows.

1. We have corrected that typo.

2. $\gamma$ denotes the momentum decay parameter. A clarifying remark has been added after Eq. (3) in the revised manuscript, and Section 7 in the appendix now includes a table with all notation definitions.

3. All occurrences of (num) have been replaced by Eq. (num) in the revised manuscript.

4. We chose to included theorems in the appendix to provide interested readers with theoretical intuition for our algorithmic design while not hindering readability of our work. However, our main contribution is the statement of Algorithm 2 and its empirically demonstrated performance. If you believe it increases readability to directly state theorems in the main paper, we are happy to include these.

5. We fully agree that it was unclear how additional operations increase the overall runtime compared to standard LoRA training. Therefore, we decide to provide a wall clock comparison of a bigger benchmark to demonstrate that aditional operations do not significantly increase the runtime. We provide a wall clock comparison of finetuning Llama2 7b-chat-hf on BoolQ and PICA benchmarks with standard LoRA [Hu et al., 2021] and Algorithm 2 in the table below. The wall-time is reported for three epochs with batch size 12 and maximal sequence length of 640 tokens on a single NVIDIA H100. We observe singificant outperformance of Algorithm 2 at a negligible wall-time overhead for all test cases.

Hyperparameters are determined by initial hyperparameter sweeps (which were performed for LoRA and Algorithm 2 individually). Hyperparameters for Algorithm 2 are reported below:

Hyperparameters for LoRA:

6. First, we remark that we only make claims about the empirical convergence rate. To that end, we have provided a convergence rate comparison in Figure 2, where we train a ViT model on Cifar-10 from scratch with

   (a) full-rank (baseline) Adam (Algorithm 4),

   (b) the LoRA-based Adam, i.e. Algorithm 4 applied to all factors of the low-rank factorization individually,

   (c) the proposed Adam-DLRT (Algorithm 2), and

   (d) a naive implementation of Adam for DLRT, that does not change the mo- mentum states.

As observed in Figure 2, the proposed Algorithm 2 has a fast initial convergence, and a higher final accuracy than the comparison methods. For empirical soundness, we have added the standard deviation to the image in the revised manuscript.

7. We provide numerical results for the ImageNet-1k (1.2 Million im- ages), where we compress the ViT-L32 vision transformer (304 Million parameters). We first conduct a hyperparameter sweep to determine standard hyperparameters as learning rate, weight-decay and batch size: ViT-L32 is trained on ImageNet with a batch size of 256, learning rate 0.001, for 20 epochs using AdamW optimizer and L2 regularization of 0.0001; DLRT uses a relative truncation tolerance of 0.013 and initial rank 200.

   We compare then the baseline model, LoRA-based simulateneous descent pretrianing, and Algorithm 2 in the table below. We observe that Algorithm 2 is able to recover the baseline accuracy up to a small margin whereas LoRA pretraining exhibits decreased top1 and top5 accuracy.

   Method	c.r. [%]	Top-1 Acc. [%]	Top-5 Acc. [%]
   Baseline	0	74.37	92.20
   Algorithm 2	61.45	72.27	90.19
   LoRA Pretrain	60.00	63.20	84.81

   Finally we wish to remark that the slightly lower compression rate is expected since the hidden dimensions of ViT-L32 of 1024 is close to the number of ImageNet classes (1000), thus there is less redundancy in the model compared to other reported benchmarks.

8. We unified notation across momentum and Adam equations, using consistent symbols for the first moment terms in the revised manuscript. A Notation table is now included in the appendix. We believe that this makes it much easier to follow our manuscript and we thank the reviewer for pointing this out.

Dear reviewer,

as encouraged by the email by the program chairs, we wanted to use this opportunity to initiate a discussion. We believe we have addressed your concerns in our answer above. Specifically, we have provided extensive additional numerical experiments that showcase the viability and scalability of our method on

a) compression of models on large datasets (ViT-L32 vision on ImageNet) and

b) large model finetuning (Llama2 7b-chat-hf on BoolQ and PICA)

We hope this fully addresses your concerns. Should any questions remain, we would of course be happy to clarify further.

Thank you again for your thoughtful review and for helping us improve the paper.

We thank the reviewer for their insightful comments.

1. Wall-clock Timing (W1):

   We fully agree that it was unclear how additional operations increase the overall runtime compared to standard LoRA training. Therefore, we decided to provide a wall clock comparison of a bigger benchmark to demonstrate that aditional operations do not significantly increase the runtime. We have added wall clock comparisons for finetuning Llama2 7b-chat-hf on BoolQ and PICA benchmarks with standard LoRA [Hu et al., 2021] and Algorithm 2 in the table below. The wall-time is reported for three epochs with batch size 12 and maximal sequence length of 640 tokens on a single NVIDIA H100. We observe singificant outperformance of Algorithm 2 at a negligible wall-time overhead for all test cases.

Hyperparameters are determined by initial hyperparameter sweeps (which were performed for LoRA and Algorithm 2 individually). Hyperparameters for Algorithm 2 are reported below:

Hyperparameters for LoRA:

2. Convergence Results Transfer (Q1):
   The convergence result for Theorem 2 directly carries over to the discrete setting of Algorithm 1. The reason for this is that a first order time discretization will not add error terms that are asymptotically bigger than the current error rates. The situation of Algorithm 2 become much more challenging, since there does not exist a direct interpretation of ADAM as a time dis- cretization of a gradient flow. We point this out in Section 4, Line 185. Theorem 1 does not directly translate into a discrete setting and a discrete version would, most probably, require additional conditions on the learning rate.

3. Typo Fix (Q2):
   You're right — it should be Windermere. We've corrected it in the revised manuscript.