Title: A geometric framework for momentum-based optimizers for low-rank training
Venue: NeurIPS 2025
Accepted: yes
Keywords: Low-Rank

Compression

Finetuning

Optimization

Manifold
Abstract: Low-rank pre-training and fine-tuning have recently emerged as promising techniques for reducing the computational and storage costs of large neural networks. Training low-rank parameterizations typically relies on conventional optimizers such as heavy ball momentum methods or Adam. In this work, we identify and analyze potential difficulties that these training methods encounter when used to train low-rank parameterizations of weights. In particular, we show that classical momentum methods can struggle to converge to a local optimum due to the geometry of the underlying optimization landscape. To address this, we introduce novel training strategies derived from dynamical low-rank approximation, which explicitly account for the underlying geometric structure. Our approach leverages and combines tools from dynamical low-rank approximation and momentum-based optimization to design optimizers that respect the intrinsic geometry of the parameter space. We validate our methods through numerical experiments, demonstrating faster convergence, and stronger validation metrics at given parameter budgets.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a new momentum based optimization framework for low rank pre-training and fine-tuning. The proposed method combines dynamic low-rank approximation and momentum methods. Compared to existing layer factorization method DLRT, the proposed method extend the applicable optimizers from SGD to Heave-ball momentum gradient descent and Adam.\n\nThe proposed framework is supported with theoretical convergence and error bound guarantees, the authors also provide extensive experimental results on pre-training, fine-tuning and transfer learning to demonstrate the performance of proposed method generally matches the full rank baseline.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n1. This paper extends DLRT to momentum optimizers, including both Heavy-ball momentum gradient descent and Adam, bridging an existing gap in the literature.\n2. The authors provide a comprehensive treatment of the proposed framework, including detailed derivations, clear computational complexity and memory analysis, as well as theoretical convergence and error bound guarantees.\n3.  This paper is well structured and easy to follow\n\nWeaknesses\n1. To enhance the readability of the paper for general readers who may not be familiar with LoRA or DLRT, it would be beneficial to include more explanations of the benchmark methods used in the experiment results, adding algorithmic descriptions for DLRT w/o proj and LoRA pretrain\n\nQUESTIONS:\n1. This paper mainly provides experimental results for Algorithm 2, how does the performance of Algorithm 2 compare to that of Algorithm 1?\n2. In the experiment \"Pretraining from scratch of ViT-small on Cifar10\", algorithm 2 exhibits faster convergence and smaller validation loss compared to the full rank baseline. Could the authors provide some insights or explanations for this observed behavior?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the momentum-based algorithm when applying to low-rank weight optimization in neural network training. The paper identifies a condition where the heaby-ball method converges to the low-rank optimum solution, and shows that directly applying momentum to the factorization form of the weight does not satisfy this condition. Based on this condition, the paper derives the heavy-ball and Adam method for low-rank problems. The paper presents experimental results in the main text, and a convergence guarantee in the appendix to corroborate the effectiveness of the proposed algorithm.\n\nSTRENGTHS AND WEAKNESSES:\n**Strength**\n\n1. The algorithm design comes from a solid theoretical principle and looks quite convincing.\n\n2. Per experimental result, the performance is better than the chosen baseline for a significant amount.\n\n**Weaknesses**\n\n1. The motivation of using the proposed algorithm is a bit weak. The paper points out the issue with directly applying momentum to DLRT by showing that the direct application does not satisfy (4). However, in the appendix it is only shown that (4) is a sufficient condition, but it is not clear whether (4) is also a necessary condition. Therefore, I believe that the explanation of why direct application of the momentum to DLRT is inconclusive.\n\n2. I am not sure whether the experimental result in Table 1,2 are measured in epochs or in total training time. If the performance is based on a fixed number of epochs, then the comparison can be unfair as the proposed algorithm imposes a constant times of additional per-iteration cost. \n\n3. Convergence guarantee in Theorem 3 is not strong enough. In particular, Theorem 3 bounds the difference between the algorithm trajectory and an ideal trajectory that converges to the low-rank optimum. However, the upper bound depends on $\\hat{\\nu}$, which is the gap between the projection of the momentum onto the current basis and the augmented basis. Intuitively, the augmented basis should explore more possibility, so this difference can be quite large.\n\nQUESTIONS:\nDoes the proposed algorithm also work for tasks like matrix completion and matrix sensing?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nMomentum-based updates for low-rank training needs to be geometry-aware for correctness. The paper shows how naive implementation of LoRA with heavy ball method (independently updating each of the factors in Equation 3) does not ensure that update dynamics stay on the low-rank manifold. The paper gives a framework for the evolution of the factorized weights $W$ and factorized momentum terms $\\mathcal{V}$. Using this framework, the paper provides update steps and algorithms for two methods: low-rank heavy ball method and low-rank Adam method. The paper also addresses two numerical challenges in their methods: evolving $W$ in factorized form to maintain low-cost and avoiding stiffness in the dynamical system through the avoidance of inverting $S_v$ and $S_v^\\intercal$. Experimental results show that the proposed Algorithm 2 (low-rank Adam) gives both higher accuracy and higher compression rate than existing methods and naive implementation of low-rank momentum.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n - The paper addresses a short-coming of other LoRA-style methods: the updates may leave the low-rank manifold. It does so by developing a rigorous framework for momentum-based optimization methods on the manifold of low-rank matrices.\n - Proposes low-rank versions of two popular methods: Adam and Heavy Ball. The algorithms keeps inner product calculations only on low-ranked factorized matrices and thus is efficient.\n - Empirical results look good: the proposed method achieves high accuracy with high compression rates.\n\nWeaknesses\n - The experiments do not show any wall-clock time measurements. Although theoretically, all operations on done on much smaller factorized matrices, but even small operations (e.g. the truncation step in Algorithm 2) could have significant cost in the aggregate.\n\nQUESTIONS:\n- Do the convergence results given in Theorems 1 and 2 for continuous-time dynamics carry over to discrete algorithms (such as Algorithms 1 and 2)? \n- Is there an extra e in Lady Winderemere's fan on line 430?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper theoretically demonstrates that conventional momentum-based optimizers like Heavy Ball and Adam fail to converge to local optima in low-rank neural network training due to ignoring the geometric constraints of the low-rank manifold's tangent space. To address this, it extends the Dynamic Low-Rank Training (DLRT) theory and propose a geometry-aware optimization framework that explicitly enforces momentum updates to remain within the tangent space. It achieves 95% parameter compression in both image classification (pretraining) and natural language processing benchmarks (finetuning) with comparable performance.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n- Theoretical Analysis of Limitations: The paper theoretically demonstrates that conventional momentum-based optimizers fail to converge to low-rank optima due to geometric constraints.\n- Memory Efficiency with Performance: Achieves comparable accuracy to full-rank baselines at huge compression ratio across multiple benchmarks (image domain and NLP).\n\nWeakness\n- Type (Line 67): `optimizer` should be corrected to `optimizer`.\n- Clarity on gamma in equations: I think that gamma represents the damping coefficient but there is no explicit definition.\n- Ambiguous Reference to equations: It’s better to use expressions such as Eq. (5) or Equation (5) rather than (5). It makes use be confused\n- Theorem Placement: Theorems are briefly mentioned in the main text without explanation; integrating key theorems into the main body could improve readability.\n- Missing Wall Time Comparison: The paper analyzes the computational efficiency, but real-world training time (wall clock time) is absent.\n- Theoretical Convergence Rate: A direct convergence rate comparison with existing methods helps the claim of faster convergence mentioned in the abstract.\n- ImageNet Pretraining: Experiments are conducted on smaller datasets (CIFAR, UCM), but validation on large-scale datasets like ImageNet is missing. Pretraining, in this context, refers to training on large datasets to enhance generalization performance. Training on large-scale datasets like ImageNet could strengthen the persuasiveness of the paper.\n- The overall notation explanation needs to be improved for better clarity and readability.\n\nQUESTIONS:\nMentioned in Strengths And Weaknesses.\n\nLIMITATIONS:\nMentioned in Strengths And Weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-01"}