Title: Finding Low-Rank Matrix Weights in DNNs via Riemannian Optimization: RAdaGrad and RAdamW
Venue: NeurIPS 2025
Accepted: yes
Keywords: Riemannian Gradient Descent

Low-rank matrix weights

Fine-tuning

AdaGrad

AdamW
Abstract: Finding low-rank matrix weights is a key technique for addressing the high memory usage and computational demands of large models. Most existing algorithms rely on the factorization of the low-rank matrix weights, which is non-unique and redundant. Their convergence is slow especially when the target low-rank matrices are ill-conditioned, because the convergence rate depends on the condition number of the Jacobian operator for the factorization and the Hessian of the loss function with respect to the weight matrix. To address this challenge, we adopt the Riemannian gradient descent (RGD) algorithm on the Riemannian manifold of fixed-rank matrices to update the entire low-rank weight matrix. This algorithm completely avoids the factorization, thereby eliminating the negative impact of the Jacobian condition number. Furthermore, by leveraging the geometric structure of the Riemannian manifold and selecting an appropriate metric, it mitigates the negative impact of the Hessian condition number. Ultimately, this results in our two plug-and-play optimizers: RAdaGrad and RAdamW, which are RGD with metrics adapted from AdaGrad and AdamW and restricted to the manifold. Our algorithms can be seamlessly integrated with various deep neural network architectures without any modifications. We evaluate the effectiveness of our algorithms through fine-tuning experiments on large language models and diffusion models. Experimental results consistently demonstrate that our algorithms provide superior performance compared to state-of-the-art methods. Additionally, our algorithm is not only effective for fine-tuning large models but is also applicable to deep neural network (DNN) compression.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper provides a riemannian perspective on low-rank training. And introduces preconditioners for the weight gradients to improve convergence of low-rank equivalents of ADAM and AdaGrad. \n\nThe method is evaluated on LLM finetuning.\n\nSTRENGTHS AND WEAKNESSES:\npro: \n- The paper formulates low‑rank matrix training as optimization on the Riemannian manifold of fixed‑rank matrices, providing clear definitions of the geometry, providing an intuitive but lightweight introduction into the topic. \n- The paper is well structured and easy to read\ncontra:\n- The paper somewhat reinvents the wheel with robust low-rank optimization. There are prior work (see questions below) on this.\n\nQUESTIONS:\nQuestion\n- Is G_t the gradient w.r.t W (i.e. a nxn matrix) or somehow split into gradients w.r.t USigmaV^T? \n- Similarly, is H_r an SVD on \\sigma or on W? \n- How sensitive are RAdaGrad and RAdamW to choices of step‑size, decay rates, and smoothing constants? [Actionable item: provide a brief hyperparameter search (can be a smaller model, not gpt2, if compute is a bottleneck). \n\nComments:\n\n- Line 42: Citation [31] does not factorize W=PQ^\\top, but instead factorizes W=USV^\\top. Furthermore, it has been shown [a], that the convergence of the method (DLRT) is not dependent on the condition of W. More generally, the underlying dynamical low-rank approximation method [b] is robust w.r.t the curvature of the manifold, and ill conditioning of the matrix W. This should be rectified. \nIn particular, the issue of ill conditioning of the gradient would only apply to \\nabla_U L and \\nabla_V L for DLRT, but the special treatment in the basis augmentation step prevents the issue.  \n\n[a] https://arxiv.org/abs/2305.19059\n[b] DISCRETIZED DYNAMICAL LOW-RANK APPROXIMATION IN THE PRESENCE OF SMALL SINGULAR VALUES: EMIL KIERI, CHRISTIAN LUBICH, HANNA WALACH\n\nLIMITATIONS:\nI think scope and limitations are well addressed.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies training models with low-rank weight matrices. These may arise in a number of contexts like LoRA. Prior approaches (like maintaining the entries of a low-rank factorization of the matrix) lead to ambiguities (the same matrix has many representations) and dependences on the condition numbers of the Jacobian and Hessian in the convergence. \n\nThis work suggests a new approach using Riemannian Optimization. Roughly, one does gradient descent on the manifold of low-rank matrices. There are two variants presented, building on AdaGrad and AdamW.\n\nThe paper first describes the mathematical details of the approach, including its computational complexity. It then describes experiments, showing that this approach leads to better correctness in various settings, including LoRA for GPT-2.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n- Natural idea for getting around the issues with prior approaches for low rank weights\n- Strong theoretical analysis, and very clear presentation\n- Experiments show real improvements by just making this change to LoRA and in other settings.\n- The approach is general and can be \"plugged in\" to many different models and frameworks.\n\nWeaknesses\n- I have a couple questions about the experiments and running time; see below.\n\nQUESTIONS:\n- In the experiments, were the Jacobian and Hessian condition numbers large? Is the main source of improvement that the algorithm can not make use of matrices with large condition numbers? I was expecting the experiments to show faster convergence, rather than actually better quality scores, so I'm wondering if you could explain this more.\n\n- Could you elaborate on \"it is easy to find an orthonormal basis...\" on line 152? I was expecting calculations in the manifold to be slower than in prior approaches, but it seems to be roughly just as fast?\n\n- Could you compare more precisely the computational complexity with prior work? What are the running times of prior algorithms?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes learning of the low-rank weights using Reimanian gradient descent method. To this end the paper adapts Adam and RMSprop accelerated gradient descent methods to be compatible with Reimanian gradient descent. The authors provide in depth mathematical formulation for the approach and back the method with solid experimental results.\n\nSTRENGTHS AND WEAKNESSES:\nThe paper’s strength lies in its principled foundation in optimization theory and its connection to a rich body of prior work. However, this also seems to be its main weakness.\n\nThe Riemanninan manifold optimizaton has been applied to the neural networks in many different ways and produced plethora of works. I think the paper would strongly benefit from thorough comparison to the previous works and having clear statement on what makes this paper different from preivous one. I especially would like to see comparisons to the ICLR-2019 paper of \"Riemannian Adaptive Optimization Methods\" that propose modifications of Adam, Adagrad, and Amsgrad with respect to Riemannian optimization. There are many more relevant works(if you follow the references), but I would list couple more here:\n- Geometry Aware Constrained Optimization Techniques for Deep Learning, CVPR2018\n- Averaging Stochastic Gradient Descent on Riemannian Manifolds, COLT 2018\nThis is by no means a full list of relevant papers, there is probably dozens of papers proposing some accelerationg scheme + Rimeannian manifold; \n\nGiven the amount of research happened in the field, I think what would put this paper apart would be the actual practical implementation (most of the Riemanian mehtods require SVD or other full decomposition on every gradient step), however, the paper is very scarce on the practical implementation. There is the part of the pseudocode in supplementary materials and the complexity analysis of the steps (lines 270-274), but it is hard to understand what is the actual runtime implications. In practice, many of the efficient factorizations (like QR, or computing inverse) take considerably longer times (have higher constants) wrt to similarly complex operations, and as such, the time cost of every update might be considerably longer. Previous works typically circumvent this by doing proper projections ever K steps. If authors can share the wallclock of every projection and compare it to vanilla SGD updates that would better idea about pracitcal usability of the method. Also, it seems this method favors smaller ranks (due to cost of updates), however, when training  LoRA adpaters you rarely see rank-1 and rank-4 apdaters, typically I would say it is rank 10 at the very least. \n\nAdditionally, no experiments are presented for pure compression (i.e., training a smaller model directly), which would broaden the applicability of the method beyond adapters.\n\nThere are also issues with citation accuracy. For example:\n-  on lines 31-32, when referring to low-rank  being used in dnn compression paper cites 14, 24, 38, but none of these papers in fact use low-rank compression.\n- on line 232 the reference 26 is great (I love this book), but can you pinpoint exact chapter/equation you are citing?\n\nSome claims are imprecise:\n- on line 30, papers says that \"these methods often rely on specialized hardware\" where \"these\" referes to \"quantization\" and \"pruning\", however, quantization is more-or-less hardware agnostic these days (as long as your hardware supports integer multiplications & has specialized kernels/code) and for instance structured pruning is definitely hardware agnostic. Therefore, the statement is imprecise and should be corrected.  \n- on line 81 \"Current mainstream methods\" supposed to refer to \"general low-rank optimization methods\", however only cite low-rank adapter research (i.e., learing w+uv with frozen w), which by no means would be \"mainstream\". Probably, you meant to write \"current mainstream low-rank adaptation methods\"?\n\nOverall, this is a promising paper, but it would benefit greatly from a deeper discussion of prior work, clarity on practical implications, and more careful citation and phrasing.\n\nQUESTIONS:\nPlease address the weaknesses above^\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-04"}

comments: SUMMARY:
The authors study memory efficient tuning of large models through low rank optimization. They propose a framework based on the optimization of the Riemannian manifold. This approach is able to eliminate the negative impact of the Jacobian condition number. Specifically, they propose RAdaGrad and RAdamW. The empirical result shows the effectiveness of the proposed method.

STRENGTHS AND WEAKNESSES:
Strengths:

The strength of the paper is that it seems to propose a framework that generalizes the existing methods. The empirical result also supports the effectiveness of the proposed framework.

Weakness:

However, the relation with the existing works is not discussed in detail. Specifically, the difference between the proposed method and ScaledGD and ScaledAdamW is not clearly stated, which are also methods based on Riemannian optimization. Additionally, there is LoRA-RITE[1], which also seems to avoid the Jacobian condition number, while adopts adaptive preconditioning similar to AdamW and Shampoo.

[1] Yen, J. N., Si, S., Meng, Z., Yu, F., Duvvuri, S. S., Dhillon, I. S., ... & Kumar, S. (2024). LoRA Done RITE: Robust Invariant Transformation Equilibration for LoRA Optimization. The Thirteenth International Conference on Learning Representations (ICLR 2025)

QUESTIONS:
My main question is the same as the one mentioned in the weakness section. What is the difference between the proposed method, ScaledAdam, and LoRA-RITE, which all seems like attempt to avoid the condition number problem and adopts adaptive preconditioning?

LIMITATIONS:
The limitation section feels a bit short.
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a Riemannian optimization algorithm for the problem of low rank training of deep networks. The authors claim that existing methods are vulnerable to the condition number of the jacobian operator for the factor and their proposed method avoids this trap. The authors derive riemannian equivalents of adamW and adagrad by proposing alternative riemannian metrics. The paper empirically demonstrates improved performance in LM finetuning and pre-training tasks.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: The paper studies an important problem of LoRA finetuning. The empirical results show that their training algorithms improve performance.\n\nWeaknesses: This paper fails to adequately demonstrate the existence of a convergence problem in LoRA finetuning, and further demonstrate that their method adequately addresses the problem. \n\n1. The authors do not include any examples of failed convergence in training low rank factorization models, and only argue that the condition number of W is a problem through a derivation of the hessian. They do not compare convergence rates that the condition number of the jacobian plays a factor in the poor convergence of SGD in the factored parameterization.\n\n2. The proposed method is Riemannian gradient descent/Adam using the manifold of low rank matrices. This does not improve on the time complexity of the factorized parameterization, and might involve more computation since an SVD is required at each step.\n\n3. The empirical evaluation only establishes improved performance in terms of BLEU, etc. The authors do not show that their convergence times are better.\n\n4. Prior work has already developed riemannian adaptive gradient methods, which the authors do not cite. https://arxiv.org/abs/1810.00760\n\nQUESTIONS:\nSee weaknesses\n\nLIMITATIONS:\nNone", "TITLE": "Official Review", "RECOMMENDATION": 2, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-09"}