PaperlayerView exact revision on arXiv

Source-linked AI summary

SubTrack++ : Gradient Subspace Tracking for Scalable LLM Training

Sahar Rajabi, Nayeema Nonta, Sirisha Rambhatla

arXiv:2502.01586v3cs.LG

TL;DR

LLM training is resource-intensive, and low-rank optimization brings challenges in tracking changing subspaces and handling discarded gradients. SubTrack++ combines Grassmannian tracking, projection-aware Adam updates, and recovery scaling; it reports state-of-the-art convergence and accuracy across baselines. Its convergence analysis assumes an unchanged subspace, leaving the impact of slowly changing subspaces open.

  • Problem

    LLM training requires substantial resources, while low-rank optimization must track changing gradient subspaces, align optimizer states, and address discarded gradient components.

    Source: §1 · p. 1, §2 · p. 2

  • Method

    SubTrack++ tracks gradient subspaces on the Grassmannian, adapts Adam’s internal statistics to subspace changes, and uses recovery scaling for information lost through low-rank projection.

    Source: §1 · p. 2, §2 · p. 3

  • Results

    SubTrack++ achieves state-of-the-art convergence and accuracy across baselines while retaining low-rank methods’ memory efficiency.

    Source: §6 · p. 10

  • Takeaways & Limitations

    Grassmannian tracking is presented as a plug-and-play module across optimizers, while extending projection-aware optimization beyond the Adam family requires further investigation.

    Source: §6 · p. 10

    The convergence analysis assumes an unchanged subspace; the impact of slowly changing subspaces remains an open problem.

    Source: §3 · p. 5

Abstract

from arXiv · show

Training large language models (LLMs) is highly resource-intensive due to their massive number of parameters and the overhead of optimizer states. While recent work has aimed to reduce memory consumption, such efforts often entail trade-offs among memory efficiency, training time, and model performance. Yet, true democratization of LLMs requires simultaneous progress across all three dimensions. To this end, we propose SubTrack++ that leverages Grassmannian gradient subspace tracking combined with projection-aware optimizers, enabling Adam's internal statistics to adapt to subspace changes. Additionally, employing recovery scaling, a technique that restores information lost through low-rank projections, further enhances model performance. Our method demonstrates SOTA convergence by exploiting Grassmannian geometry, reducing pre-training wall-time by up to 65% and fine-tuning time by 36% compared to existing SOTA methods, while maintaining the same memory footprint.

1 Introduction

SubTrack++ uses Grassmannian subspace tracking, projection-aware Adam updates, and recovery scaling to address training cost and low-rank optimization challenges. The paper reports improved convergence and evaluation loss alongside competitive memory and runtime efficiency.

1 Introduction

  • LLM training and fine-tuning require substantial hardware and time, while low-rank and block-coordinate approaches can trade accuracy for memory savings.Optimizer states also consume a significant portion of memory beyond trainable parameters.Source: §1 · p. 1
  • Figure 1 reports the lowest evaluation loss across methods, memory on par with GaLore and Fira, and markedly faster wall-time than GaLore, Fira, and LDAdam.Peak memory is also significantly lower than APOLLO and LDAdam, with minimal wall-time overhead relative to APOLLO.Source: Figure 1 · p. 2, §1 · p. 2
  • Grassmannian subspace tracking is presented as robust to perturbations and suitable for evolving gradient subspaces in LLM training.Prior work cited in the passage associates it with structurally embedded information, lower computational complexity, and improved performance.Source: §1 · p. 2
  • SubTrack++ tracks subspaces along Grassmannian geodesics, reuses prior subspaces and estimation error instead of repeatedly reconstructing them with expensive SVD, and adapts Adam statistics to subspace shifts.Recovery Scaling uses low-rank optimizer scaling information to recover and scale gradient information lost through projection.Source: §1 · p. 2
  • SubTrack++ supports full-parameter training, maintains GaLore’s memory footprint, and is more time-efficient than SVD- or PowerSGD-based methods.The introduction also reports state-of-the-art convergence and evaluation loss across strong baselines.Source: §1 · p. 2

2 SubTrack++

SubTrack++ combines Grassmannian subspace tracking with projection-aware optimizer updates and recovery scaling. Its design addresses subspace drift, optimizer-state alignment, and information discarded by low-rank projection.

2 SubTrack++

  • Low-rank optimization reduces memory but requires adaptive subspace tracking, alignment of Adam’s state with changing coordinates, and recovery of discarded gradient components.SVD-based tracking is computationally expensive, while recovered gradient signals can boost performance.Source: §2 · pp. 2–3
  • SubTrack++ combines Grassmannian tracking that uses estimation error and subspace history, projection-aware Adam state updates, and recovery scaling of discarded gradient components.The overview identifies these as the method’s three core components.Source: §2 · p. 3
  • Subspace estimation is framed as selecting a point on the Grassmannian, refining it with prior subspace information and estimation error through lightweight updates.The method initializes the subspace using the top-r left singular vectors of the step-0 gradient’s SVD.Source: §2 · p. 3
  • Gradients are projected into a reduced subspace, and every k steps the subspace is updated by minimizing its distance to the current gradient.The residual from the least-squares fit lies in the subspace’s orthogonal complement and informs the Grassmannian tangent vector.Source: §2 · p. 3
  • To control subspace shifts, SubTrack++ uses a rank-1 approximation of the gradient of the estimation error and moves along a Grassmannian geodesic.The update preserves orthonormality and retains the portion of the previous subspace not updated in that step.Source: §2 · p. 3
  • At subspace updates, the projection-aware optimizer modifies Adam’s momentum rules to reflect coordinate changes in its statistics.The stated goal is to keep optimizer statistics aligned as the subspace evolves.Source: §2 · pp. 3–4
  • Recovery scaling adds a correction for gradient information lost through low-rank projection, using scaling information from the low-rank optimizer.A clipping-inspired mechanism limits the correction’s growth rate by factor ζ.Source: §2 · pp. 4–5
  • The authors report robust, stable training and state-of-the-art performance while maintaining minimal memory footprint and wall-time.The operation sequence is presented in Algorithm 1.Source: §2 · p. 5

3 Theoretical Analysis

The analysis derives SubTrack++’s subspace update from Grassmann geometry and establishes convergence under stated assumptions. It also explains the trade-off between tracking changing subspaces and the conditions used in the convergence proof.

Theoretical Analysis

  • For gradient matrix dimensions m≤n, SubTrack++ uses left projection; otherwise it switches projection sides, and ρ_t′ is the optimizer’s entry-wise regularizer.S_t′ denotes the projection matrix that projects the gradient onto the subspace.Source: §3 · p. 5
  • Under the theorem’s stated assumptions, including fixed projection matrices and ρ_t≡1, P_t→0 and the subspace-tracking process converges.The assumptions also require L-continuous gradient-component functions, bounded weights, and min(κ_t) > (L_A + 2L_B L_C M^2).Source: §3 · p. 5
  • The convergence proof assumes an unchanged subspace; unlike GaLore, SubTrack++ uses rank-1 updates to avoid drastic changes and allow more frequent updates, while slowly changing subspaces remain open to deeper analysis.GaLore must limit its updates to ensure convergence because each update can potentially change the entire subspace.Source: §3 · p. 5
  • The subspace update is derived from Grassmann geometry: the exponential map sends a tangent direction to the endpoint of its originating geodesic.A Grassmann point represents a subspace, while the exponential-map definition relates tangent directions to geodesics.Source: §3 · p. 6
  • SubTrack++ uses a constant step size to track subspace changes, with convergence guaranteed when changes are controlled to maintain the stable-subspace assumption.The analysis notes that decreasing the step size can impair accurate tracking and adaptation to subspace changes.Source: §3 · p. 6

4 Experiments and Results

Across pre-training and fine-tuning experiments, SubTrack++ combines faster training with strong performance while retaining low-rank methods’ memory efficiency. Ablations and robustness tests examine the contributions and operating conditions of its subspace-tracking approach.

Pre-Training Experiments

  • 43% faster than LDAdam on the 1B model and 67% faster on the 3B model, SubTrack++ sets state-of-the-art results across model scales.The comparison identifies LDAdam as the strongest prior baseline.Source: §4 · p. 7
  • At 100k training steps, SubTrack++ reaches evaluation loss 3.37, while GaLore reaches 4.64 under the same conditions.In the 1B pre-training comparison, SubTrack++ also converges fastest across training steps and wall-time.Source: §4 · p. 7
  • In Llama-1B pre-training, full-gradient norms fall from 0.46 to 0.08 and projected-gradient norms from 0.45 to 0.05, following nearly identical trajectories.SubTrack++ occasionally outperforms full-rank training; the authors suggest low-rank projection’s implicit regularization may contribute.Source: §4 · p. 7

Ablation Studies

  • Grassmannian tracking alone substantially reduces wall-time versus frequent SVD updates; adding the projection-aware optimizer and recovery scaling lowers loss from 6.53 to 5.43 and 5.28, respectively, and their combination reaches 4.51.These gains require only minimal increases in runtime and memory overhead.Source: §4 · p. 8, Figure 4 · p. 8
  • Rank-1 updates achieve the best performance among tested ranks, while increasingly frequent updates help only up to a point and excessive frequency may impede convergence.The authors associate controlled, small subspace adjustments with stability and rank-1 updates with avoiding stagnation in a low-rank region.Source: §4 · p. 8, Figure 5 · p. 9

Fine-Tuning Experiments

  • On Llama-2-7B-chat-hf Alpaca fine-tuning, SubTrack++ achieves 36% lower wall-time and better performance than GaLore.The supervised fine-tuning experiment ran for one epoch.Source: §4 · p. 8

Time and Space Complexity

  • SubTrack++ uses Grassmannian tracking at the same subspace-update frequency as GaLore and Fira, with optimizer-state memory equivalent to GaLore and other baselines.The authors state that its update-time complexity explains its greater efficiency than SVD-based methods.Source: §4 · p. 8, Table 2 · p. 8

Robust Subspace Tracking

  • The subspace-drift norm on Llama-350M pre-training falls from 0.06 to below 0.0002 within 5–7 updates and remains near zero thereafter.The authors report that GaLore’s SVD can make erratic jumps on Ackley; at scale factor 1 it does not reach the global minimum within 100 steps, while scale factor 3 increases jump length.Source: §4 · p. 9

5 Related Works

Related work addresses LLM training efficiency through parameter-efficient tuning, gradient and optimizer-state compression, and evolving low-rank subspaces. SubTrack++ is situated among approaches that track subspaces geometrically and methods that use projections or optimizer-state adjustments.

Parameter-Efficient Training

  • LoRA [Hu et al., 2021] reduces fine-tuning memory with low-rank trainable matrices, while quantization, activation compression, and block-coordinate methods offer other efficiency strategies.The section also describes full-parameter fine-tuning via random projections of the gradient matrix.Source: §5 · p. 9

Gradient Low-Rank Projection

  • GaLore projects gradients into a low-rank subspace and back for full-parameter tuning; related methods compress gradients, use sparse projections, or update selected parameters.Fira adds norm-based scaling, LDAdam [Robert et al., 2025] uses projection-aware updates and generalized error feedback, and Online PCA is used for evolving projections [Liang et al., 2024].Source: §5 · pp. 9–10

Geometric Subspace Updates

  • On the Ackley function, GaLore’s SVD fails to reach the global minimum within 100 steps at scale factor 1; scale factor 3 reaches it but increases jump length, unlike the controlled Grassmannian updates.The comparison uses a subspace update interval of 10.Source: Figure 6 · p. 10
  • Grassmannian methods track evolving subspaces, including incremental updates for partially observed data and geodesic-based subspace evolution [Balzano et al., 2011, Blocker et al., 2023].LORO [Mo et al., 2025] applies Riemannian optimization to fixed-rank matrices for low-rank pretraining.Source: §5 · p. 10

6 Discussion and Conclusion

SubTrack++ combines Grassmannian subspace tracking, projection-aware optimization, and recovery scaling for low-rank gradient training. It reports state-of-the-art convergence and accuracy across baselines, while noting limits to optimizer-family extension and model scale.

6 Discussion and Conclusion

  • SubTrack++ combines Grassmannian subspace tracking, projection-aware Adam updates, and recovery scaling to use gradient information lost through low-rank projection.The method projects gradients into a low-rank subspace while reflecting subspace changes in Adam’s internal statistics.Source: §6 · p. 10
  • The method achieves state-of-the-art convergence and accuracy across the reported baselines.Grassmannian tracking preserves the computed subspace while incorporating gradient components from its orthogonal complement.Source: §6 · p. 10
  • Projection-aware optimization beyond the Adam family requires further design and investigation, and benchmarking above 7B parameters was infeasible given the authors’ time and resources.Source: §6 · p. 10

A Convergence of SubTrack++

The convergence proof derives a recursive gradient form and bounds its error under stated smoothness, norm, projection, and learning-rate conditions. Under the constant-projection assumption, it concludes that the error term converges to zero, including with a left projection matrix.

A Convergence of SubTrack++

  • The stated convergence condition requires constant projections, bounded weight norm, L-continuous functions, and min(κ_t) > (L_A + 2L_B L_C M^2).The theorem passage specifies learning rate μ and ρ_t ≡ 1 as well.Source: A Convergence of SubTrack++ · p. 14
  • The proof assumes fixed projection matrices during training to derive the recursive form of g_t.Source: A Convergence of SubTrack++ · pp. 14–15
  • The proof concludes that P_t converges to zero when the left projection matrix is used.The derivation states that the recursive formulation remains unchanged and the modified P_t still satisfies the required bounds.Source: A Convergence of SubTrack++ · pp. 14, 16

B Grassmann Exponential

The appendix describes the Grassmann exponential as the geodesic generated from a subspace point in a tangent direction. It then explains how the projection-aware optimizer transforms Adam’s moment estimates when the tracked subspace changes.

B Grassmann Exponential

  • The Grassmann exponential gives the geodesic emanating from subspace P in tangent direction ∆.The appendix represents ∆ through its horizontal lift and thin SVD before stating the exponential map.Source: B Grassmann Exponential · p. 16

C Projection-Aware Optimizer

  • For changing subspaces, S_t^⊤S_{t−1} represents the basis change used to transform Adam’s first-moment estimate.The optimizer also changes basis for the second-moment estimate; the derivation assumes independent gradient coordinates to approximate covariance.Source: C Projection-Aware Optimizer · pp. 17–18
  • The appendix derives projection-aware optimizer update rules by applying equations (38) and (41).Source: C Projection-Aware Optimizer · p. 18

D Time Complexity Analysis

SubTrack++ avoids an SVD of the full gradient matrix, yielding an update complexity dominated by O(mnr). Its Llama-7B pre-training comparison also reports substantially lower wall-time than GaLore.

D Time Complexity Analysis

  • For 100k Llama-7B pre-training iterations, Figure 7 compares training loss against steps and wall-time and reports substantially lower wall-time than GaLore.Source: Figure 7 · p. 19
  • SubTrack++’s subspace update is dominated by O(mnr) matrix multiplications and avoids the O(nm^2) SVD of the full gradient matrix.The complexity analysis assumes an m × n gradient matrix and rank-r projection, with r ≪ m ≤ n.Source: D Time Complexity Analysis · pp. 18–19, Table 3 · p. 19

E Pre-Training Llama-Based Architectures

The authors pre-trained six Llama-based architectures for 10k iterations and also evaluated SubTrack++ and GaLore on 7B models for 100k iterations, where SubTrack++ maintained faster convergence and superior performance.

E Pre-Training Llama-Based Architectures

  • Across six Llama-based architectures, pre-training ran for 10k iterations using the hyperparameters reported in Table 4.The 7B comparison used 100k iterations; its hyperparameters matched Table 4 except for the iteration count.Source: E Pre-Training Llama-Based Architectures · p. 19
  • For 7B models trained for 100k iterations, SubTrack++ maintained faster convergence and superior performance than GaLore.The authors describe this comparison as demonstrating the method’s generalizability.Source: E Pre-Training Llama-Based Architectures · p. 19
  • The 100k-iteration 7B run used Table 4 hyperparameters, except for the number of iterations.Figure 8 presents a bar-chart version of Figure 1.Source: E Pre-Training Llama-Based Architectures · p. 19

F Memory and Time Comparison

For pre-training, SubTrack++ is reported to match or exceed baselines in memory and runtime efficiency while achieving the lowest evaluation loss in the Figure 8 comparison. The timing comparisons use a fixed number of subspace updates, with hardware and iteration settings varying by model size.

F Memory and Time Comparison

  • SubTrack++ requires on-par or better peak memory than GaLore while achieving state-of-the-art results.Table 5 compares peak memory consumption across pre-training baselines.Source: F Memory and Time Comparison · p. 19
  • Wall-time comparisons include 10 subspace updates per method; 60M–3B runs use 2,000 iterations on A100 GPUs, while 7B runs use 5,000 iterations on RTX A6000 GPUs.The update intervals are 200 for 60M–3B models and 500 for 7B models.Source: F Memory and Time Comparison · p. 19
  • SubTrack++ achieves the lowest evaluation loss across methods while matching or exceeding baselines in memory and runtime efficiency in the 1B comparison.Figure 8 reports memory on par with GaLore and Fira and lower than APOLLO and LDAdam; runtime is minimal overhead relative to APOLLO and markedly faster than GaLore, Fira, and LDAdam.Source: Figure 8 · p. 20

G Fine-Tuning Experiments

The experiments fine-tune RoBERTa-base and RoBERTa-large on GLUE and SuperGLUE, and Llama-2-7B-chat-hf on Alpaca. In the Alpaca experiment, SubTrack++ is reported to use less wall-time and achieve better performance than GaLore.

G Fine-Tuning Experiments

NeurIPS Paper Checklist

The checklist summarizes the paper’s reporting on limitations, theoretical proofs, reproducibility details, ethics and impact, statistical testing, and LLM use. It notes that some ethics and impact questions and LLM usage are not applicable, and that statistical significance tests were not performed because of resource and time constraints.

Limitations / Main claims

  • [Yes] The paper says its limitations and future directions are discussed, and computational complexity is reported.It also answers [Yes] that the abstract and introduction reflect the method, results, and experimental scope.Source: NeurIPS Paper Checklist · pp. 24–25

Theory assumptions and proofs

  • [Yes] The paper says theoretical proofs are provided and aligned with prior work, with support from sources on Grassmannian geometry.Source: NeurIPS Paper Checklist · p. 25

Experimental result reproducibility / Open access to data and code

  • [Yes] The paper says its supplemental material provides resources and hyperparameters for reproducing the proposed method and baselines.It also says an anonymized code repository is included in the supplemental material and that code will be included in the paper after submission.Source: NeurIPS Paper Checklist · pp. 25–26

Experiment statistical significance

  • [No] The paper reports no statistical tests because its experiments were resource- and time-consuming.The authors say they included a variety of experiments to ensure generalization.Source: NeurIPS Paper Checklist · p. 27

Societal impacts / Safeguards / Ethics

  • [NA] The paper treats societal impact and safeguards as not applicable, describing the work as foundational research and leaving safety investigation for future work.It also says the research conforms to the NeurIPS Code of Ethics and that crowdsourcing and human-subject research do not apply.Source: NeurIPS Paper Checklist · pp. 28–30

Licenses for existing assets / New assets

  • [Yes] The paper says existing assets are properly cited and new code assets in the anonymous supplemental material are documented.Source: NeurIPS Paper Checklist · p. 29

Declaration of LLM usage

  • The checklist marks LLM usage as not applicable because LLMs are not important, original, or non-standard components of the core method.The guidelines state that such components would need to be described.Source: NeurIPS Paper Checklist · p. 30
Loading 2502.01586v3…