Title: Breaking the Frozen Subspace: Importance Sampling for Low-Rank Optimization in LLM Pretraining
Venue: NeurIPS 2025
Accepted: yes
Keywords: low-rank optimization

large language models

Adam

memory efficiency
Abstract: Low-rank optimization has emerged as a promising approach to enabling memory-efficient training of large language models (LLMs). Existing low-rank optimization methods typically project gradients onto a low-rank subspace, reducing the memory cost of storing optimizer states. A key challenge in these methods is selecting suitable subspaces to ensure an effective optimization trajectory. Most existing approaches select the dominant subspace to preserve gradient information, as this intuitively provides the best approximation. However, we find that in practice, the dominant subspace stops changing during pretraining, thereby constraining weight updates to similar subspaces. In this paper, we propose importance sampling for low-rank optimization in LLM pretraining with a provable convergence guarantee, which the dominant subspace approach does not have. Empirically, we demonstrate that our method significantly outperforms previous methods in LLM pretraining tasks.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper identifies and addresses the \"frozen subspace\" problem in low-rank optimization methods used for LLM pretraining. The authors observe that methods like GaLore, which rely on projecting gradients onto a dominant subspace, often suffer from this subspace becoming static early in training, thus limiting the model's learning capacity. To counteract this, the paper proposes SARA (Importance Sampling for Low-Rank optimization), which replaces the deterministic selection of top singular vectors with a stochastic, importance-sampling-based approach. By sampling basis vectors with probabilities proportional to their singular values, SARA encourages the exploration of more diverse subspaces throughout training. The work is supported by a theoretical convergence guarantee and empirical results on LLaMA models, demonstrating that SARA consistently improves performance over dominant subspace baselines and narrows the gap with full-rank Adam.\n\nSTRENGTHS AND WEAKNESSES:\nThis paper presents a contribution to memory-efficient optimization, with several key strengths but also noteworthy weaknesses that temper its overall impact.\n\n**Strengths**\n\n1. Clear Motivation and Problem Diagnosis: The identification and empirical illustration of the \"frozen dominant subspace\" phenomenon provide a compelling rationale for the proposed work.\n\n2. Simple and Practical Method: As a plug-and-play module that modifies only the subspace selection step, SARA is highly practical and can be easily integrated into existing low-rank optimization frameworks like GaLore and Fira.\n\n3. Theoretical Backing: Providing a formal convergence guarantee (Theorem 3.4) is a significant strength. It adds a layer of theoretical rigor that is absent in the popular dominant subspace approach, bolstering the credibility of the method.\n\n4. Good Clarity and Presentation: The paper is well-written and easy to follow. The problem, method, and results are presented logically and supported by clear figures and algorithms.\n\n**Weaknesses**\n\n1. Incremental Novelty: The core idea of analyzing and modifying low-rank subspaces to improve performance is an established paradigm, particularly in work on improving LoRA. This positions the contribution as an effective but arguably incremental advancement rather than a fundamentally new approach.\n\n2. Limited Experimental Scope and Scale: The empirical evaluation, while covering multiple optimizer variants, lacks the scale and breadth expected of top-tier work in this area. \n\ni) The experiments do not extend to larger, more commonly benchmarked models like a LLaMA-7B, and larger datasets, which would be a more convincing test of scalability.\n\nii) The evaluation is confined to perplexity. It critically lacks evaluation on a suite of downstream fine-tuning tasks (e.g. Avg Acc), which is the standard for demonstrating the quality and generalizability of a pretrained model.\n\niii) There is a crucial absence of any ablation or sensitivity analysis on the rank r. (a key hyperparameter)\n\n3. Marginal Performance Gains: The practical significance of the improvements is questionable in some cases. For instance, on the 130M model, SARA's perplexity gain over the Fira-Adam baseline is minimal (22.22 vs. 22.37). Maybe the SOTA baseline has implicitly solved the subspace problem.\n\n4. \"No Overhead\" Claim: The paper claims SARA \"does not bring extra overhead\". However, as the method introduces an explicit weighted sampling step after the SVD. This claim requires empirical validation with wall-clock time and memory usage comparisons to prove that the additional computational cost is indeed negligible.\n\n5. The paper's own theory suggests SARA's convergence rate is slightly worse than GoLore's by a constant factor , yet SARA consistently outperforms GoLore in experiments. This gap is not adequately discussed, leaving open questions about the factors driving the empirical success.\n\nQUESTIONS:\nPlease refer to the Weaknesses.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses a key issue in low-rank optimizers for LLM training: the update directions tend to collapse into the same subspace over time, limiting learning. The authors propose a simple method to fix this while keeping the efficiency benefits of low-rank updates.\n\n**Main contributions**\n\n- Propose **SARA**, a sampling method that selects singular vectors based on their importance to keep update subspaces diverse.  \n- Prove that SARA maintains the convergence rate of standard low-rank methods while preserving memory and compute efficiency.  \n- Show that SARA can be easily added to existing optimizers like GaLore, Fira, and Adafactor without major code changes.  \n- Demonstrate improved results on LLaMA models, reducing the gap between low-rank and full-rank Adam on C4 and SlimPajama.  \n- Provide evidence that SARA improves subspace diversity using overlap and rank-based metrics.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n\n- **Clear motivation.** The paper addresses a real issue in low-rank training: the update directions collapse into the same subspace over time, limiting learning. This is a well-motivated and practically important problem.  \n- **Simple and useful method.** The proposed SARA algorithm is easy to implement, requires no extra memory or compute, and can be applied to many existing optimizers.  \n- **Theoretical support.** While not central to the paper’s impact, the convergence analysis adds technical credibility.  \n- **Strong experiments.** The paper includes extensive experiments on LLaMA models, showing consistent improvements and closing much of the gap with full-rank Adam.  \n\n## Weaknesses\n\n- **Unorganized experiments.** Section 4.3 visualizes subspace metrics for just one parameter group, with the rest of the figures relegated to the appendix. The paper should aggregate and discuss results across all layers and groups in the main text (for example, plotting the average subspace overlap per layer and the average overlap over training steps for each parameter group) to give a clearer, more comprehensive view of the method’s behavior. \n- **No comparison with alternative sampling methods.** The paper does not benchmark SARA against simpler schemes such as uniform, top-k, or random-k sampling, leaving it unclear how much of the gain is due to importance sampling versus just adding noise or diversity.\n- **Writing could be improved.** Some parts of the paper are under-explained, especially the sampling procedure. The design choices, such as sampling probabilities being proportional to singular values, are not well motivated or justified.\n- **No analysis of randomness.** It is unclear whether the randomness introduced by sampling could hurt performance or stability in some cases. A discussion or ablation would help. For instance, does it affect the distributed training/parallelization?\n\nQUESTIONS:\n1. **Why sample by singular values?**  \n   Why are the sampling probabilities proportional to singular values, for instance, why not proportional to exp(.) of singular values to resemble a softmax function? Have you tried uniform sampling? An ablation would clarify this design choice.\n\n> **I understand that running additional experiments might not be feasible, I'd be happy to increase my score if authors can provide some discussion and maybe run some simple experiments with uniform sampling.**\n\n2. **Aggregation of subspace metrics**  \n   Section 4.3 shows overlap for only one parameter group. Please aggregate results over all layers and groups (e.g., average or any reasonable summarization metric for overlap per layer and per training step) to give a complete picture.\n\n> **I am willing to increase my score if authors can aggregate the results for subspace metrics.**\n\n3. **How often is SVD computed?**  \n   How frequently do you compute the SVD and update the sampling probabilities? I could not find details about the hyperparameter $ \\tau $ in the paper. I assume it matches the baseline methods, but please clarify and report it in Appendix B. Also, a rough estimate of the added runtime or FLOP overhead would be helpful.\n\n4. **Is sampling stable across runs?**  \n   Does the added randomness affect training stability or convergence? Showing results across multiple seeds would help. Also, does random sampling affect the optimal value for hyperparameter $\\tau$?\n\n5. **Typo in Figure 4?**  \n   The x-axis in Figure 4 (and similar plots in the appendix) is labeled \"update steps,\" but it appears to show singular value indices. Please clarify or correct the label.\n\nLIMITATIONS:\nAdd a short paragraph acknowledging technical limits: extra SVD cost, possible instability from random sampling, etc.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper tackles a limitation of recent low-rank optimizers such as GaLore and Fira: after a few thousand iterations the dominant gradient subspace stabilises, so subsequent updates accumulate in almost the same low-rank slice and stall expressiveness.\n\nTo break this \"frozen subspace\", the authors introduce SARA—Importance SAmpling for Low-RAnk optimisation. Every $\\tau$ steps they compute the SVD of the minibatch gradient and sample r left singular vectors without replacement in proportion to their singular values (Alg. 2). The resulting projector replaces the usual top-r basis and can be plugged into any gradient-projection optimiser.\n\nKey claims are:\n\n* Theory – SARA's stochastic projector yields an $O(L\\Delta/(\\delta^{2.5}T)+\\sqrt{L\\Delta}\\sigma/(\\delta^{3.5}T))$ convergence rate (Theorem 3.4), matching the expected $O(1/\\sqrt{T})$ for SGD momentum in non-convex smooth objective. In comparison, there is no convergence guarantees for the vanilla GaLore method.\n\n* Practice – On LLaMA-60 M/130 M/350 M/1.1 B pre-training over C4 and SlimPajama, SARA paired with GaLore or Fira achieves slightly better or on-par performance when compared with Galore in terms of perplexity.\n\n* Dynamics – SARA lowers adjacent-subspace overlap (Figure 1, 3) and produces higher-rank weight updates (Figure 4), supporting the motivation.\n\nSTRENGTHS AND WEAKNESSES:\n### Strength\n\n1. Applying sampling to subspace selection in low-rank projection methods is novel, and compelling to me.\n\n2. The theoretical analysis attains the convergence rate expected for momentum-based SGD on smooth objectives, and addresses the non-convergence issue of Galore.\n\n3. The paper is clearly written and straightforward to follow.\n\n4. The experimental results are comprehensive and convincingly support the authors' claims.\n\n### Weakness\n\n1. **Scalability**: Adding experiments on larger language models—for instance, Llama-3.1-8B—would strengthen the empirical evidence for the method’s scalability.\n\n2. **Theory-practice Gap**: While the theorem is proven for SGD momentum, the experiments use Adam for optimization. I understand that it is difficult to prove convergence guarantees for Adam given its non-projection-free property, but it would be beneficial to mention this in the limitations (Appendix C).\n\n3. **Minor presentation issues**:\n\n* Some missing citations [1–4].\n\n* Algorithm 1, line 7: The formatting could be refined for readability.\n\n* Tables 1–4: Sorting the results in descending order would make the comparisons clearer.\n\n* Line 469: The statement should involve inequalities rather than equalities.\n\n### References\n\n[1]: Liu, Xinxin, et al. \"Refining Salience-Aware Sparse Fine-Tuning Strategies for Language Models.\" arXiv preprint arXiv:2412.13488 (2024).\n\n[2]: Pan, Rui, et al. \"LISA: layerwise importance sampling for memory-efficient large language model fine-tuning.\" NeurIPS 2024.\n\n[3]: Hui, Tingfeng, et al. \"Hft: Half fine-tuning for large language models.\" arXiv preprint arXiv:2404.18466 (2024).\n\n[4]: Li, Pengxiang, et al. \"Owlore: Outlier-weighed layerwise sampled low-rank projection for memory-efficient llm fine-tuning.\" arXiv preprint arXiv:2405.18380 (2024).\n\nQUESTIONS:\n* What is the performance of SARA in larger-sized models, such as Llama-3.1-8B?\n\n* Do the improvements in subspace diversity lead to any performance gains for downstream tasks, such as math, commonsense reasoning, or instruction-following?\n\n* I am wondering why the similarity between adjacent projection matrices is large, given the diversity in sampled batches, and the noise is expected to be a dominant term in the gradient in the small batch training setting?\n\nLIMITATIONS:\nAlthough the theoretical analysis is conducted under SGD with momentum, the experimental results are based on the Adam optimizer. Given the challenges in establishing convergence guarantees for Adam, it would be helpful to explicitly acknowledge this discrepancy in the limitations section (e.g., Appendix C).", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the issue of \"frozen\" dominant subspaces in low-rank optimization during LLM pretraining, demonstrating that conventional dominant subspace selection leads to highly overlapping and restrictive update directions. To address this, the authors introduce SARA, i.e., importance sampling for low rank optimization, which randomly samples singular vectors weighted by their importance to construct more diverse update subspaces. Theoretical convergence guarantees are presented and SARA is shown empirically to outperform prior low-rank optimization approaches across several LLM sizes and optimizers, demonstrating reduced perplexity gaps versus full-rank baselines.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\n1.The paper provides a clear and well-motivated analysis of the limitations of dominant subspace selection in existing low-rank optimization methods, supported by empirical observations (see Figure 2) that show increasing and later-stable overlap of dominant subspaces during pretraining.\n\n2.SARA is a simple yet principled modification that introduces importance sampling into the subspace selection process, which reduces overlap between adjacent subspaces (evident in Figure 3) and encourages more effective exploration without adding computational overhead.\n\n3.The methodology is well-integrated with existing low-rank optimizers (e.g., GaLore and Fira), and the plug-and-play nature is likely to make adoption easy within the community.\n\nWeaknesses\n\n1.The evaluation focuses mainly on perplexity during pretraining. While this is the standard criterion for language modeling, it is not evident how SARA’s purported improvements translate into downstream performance (such as finetuning, alignment, or other application-level metrics). The practical significance of reduced PPL gap, especially at scale, is not elaborated on beyond pretraining.\n\n2.While the method is said to be plug-and-play and robust, there is little discussion or systematic ablation regarding the impact of SARA’s hyperparameters (e.g., the subspace refresh frequency τ, the sampled rank r, the total number of candidate singular vectors m) on convergence and final perplexity. This leaves uncertainty on how sensitive SARA is to these settings and may make reproducibility and deployment less straightforward.\n\n3.While the experiments include the model ranging from 60M to 1.1B, there is still a gap between this and the current mainstream LLM size, such as 7B or 32B.  Considering the problem of training cost, if the loss/perplexity changing chart and table under less training tokens can be provided, the paper can also be made more solid.\n\nQUESTIONS:\nPlease see the weaknesses. In general:\n\n1.Beside validation perplexity, more metrics are needed to show the model performance.\n\n2.More analysis of hyperparameter sensitivity is needed.\n\n3.It's better to scale up the model size to prove the generalization ability of the method.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 2, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}