PaperlayerView exact revision on arXiv

Source-linked AI summary

Ori: Drift-Aware Streaming $k$-PCA for Very Low-Rank LLM Training

Xinyu Luo, Jincheng Zhou, Mengbo Wang, Zilin Shen, Nadia Atallah Lanman, Petros Drineas, Bruno Ribeiro, Brian Bullins

arXiv:2610.22794v1cs.LG

TL;DR

Gradient-subspace optimizers must track principal gradient directions that change during training, while existing approaches often recompute subspaces periodically or adapt through a secondary optimization loop. Ori uses drift-aware streaming k-PCA with basis and optimizer-statistic transfer; it narrows the gap to full-rank training, including by up to 69.3% on 1B LLaMA pre-training.

  • Problem

    Gradient-subspace optimizers must adapt to evolving principal gradient directions, while existing methods recompute subspaces periodically or use a secondary optimization loop.

    Source: §1 · p. 1

  • Method

    Ori uses streaming Oja-style k-PCA to track subspaces through warm-started basis updates and transfers optimizer statistics across rotations.

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

  • Results

    Ori narrows the low-rank performance gap to full-rank Adam by up to 69.3% on 1B LLaMA pre-training and is superior in 10 of 16 RoBERTa-Base fine-tuning scenarios.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Across experiments, Ori narrows the gap to full-rank training in the extremely low-rank regime while remaining faster and more memory-efficient than prior gradient-projection methods.

    Source: §6 · p. 10

    The dynamic-regret analysis assumes streaming samples drawn i.i.d. from an unknown distribution and bounded drift between consecutive steps.

    Source: §4 · p. 6

Abstract

from arXiv · show

A prominent family of low-rank LLM optimizers, pioneered by GaLore, performs gradient-subspace projection: at each step, they project the gradient into a low-dimensional subspace and map the update back to the full model. Most existing work either recomputes the gradient subspace at certain intervals, or performs post-hoc adaptation to catch up to the moving subspace. In this work, we explicitly identify in practice the slow, stable manner in which LLM gradient subspaces drift across time. To this end, we propose Ori, a novel low-rank optimizer for training LLMs that captures the shifting principal gradient directions over time to achieve state-of-the-art accuracy under stringent low-rank constraints. Ori employs a streaming $k$-PCA algorithm inspired by Oja++, whose low per-step cost enables high-frequency, on-the-fly subspace tracking. In practice, Ori reduces the performance gap of the state-of-the-art low-rank optimizers to full-rank optimizers by up to 69.3\% on a 1B LLaMA pretraining task, and achieves the best performance on 10 out of 16 task scenarios in Roberta-Base fine-tuning. Together, these contributions allow Ori to further push the performance and efficiency frontiers of gradient-projecting optimizers, narrowing the gap to full-rank training for LLMs. Code can be found at https://anonymous.4open.science/r/Ori-Submission-2132.

1 Introduction

Ori addresses the challenge of tracking changing gradient subspaces by exploiting their observed slow drift with continuous, warm-started streaming updates. It narrows the gap to full-rank training under stringent rank constraints.

1 Introduction

  • Gradient-subspace methods must adapt as principal gradient directions evolve, but GaLore periodically recomputes the subspace with an expensive SVD and OSD uses a secondary non-convex optimization loop.GaLore’s recomputation discards prior subspace information.Source: §1 · pp. 1–2
  • After warm-up, gradient subspaces drift smoothly, with the projector at step t predictably close to its value at t − ∆t for small intervals such as 20 steps.This motivates tracking with small, frequent rotations while preserving and transferring optimizer statistics.Source: §1 · p. 2
  • Ori carries the previous basis forward and uses a few lightweight streaming k-PCA-inspired updates to realign it with the current gradient geometry.Its Oja++ implementation enables high-frequency basis rotations at per-step cost O(dk2) and transfers Adam-style first and second moments between bases.Source: §1 · p. 2
  • Ori reduces low-rank optimizers’ performance gap to full-rank Adam by up to 69.3% on 1B LLaMA pre-training and is superior in 10 of 16 RoBERTa-Base fine-tuning scenarios.On LLaMA, optimizer state rank is 128 versus an embedding dimension of 2048; RoBERTa ranks are 8 or 4.Source: §1 · p. 2
  • Ori improves the performance–memory trade-off over LoRA, GaLore, and OSD in the low-rank regime, narrowing the gap to full-rank Adam.Figure 1 compares evaluation perplexity and peak optimizer state memory on 1B LLaMA pre-training; Ori forms a better Pareto frontier.Source: §1 · p. 2, Figure 1 · p. 2
  • Ori’s contributions include identifying slow gradient-subspace drift, providing sublinear dynamic-regret bounds in the noiseless setting, and using ablations to isolate update-frequency and continuity effects.The bounds have mild dependence on the consecutive drift parameter.Source: §1 · pp. 2–3

2 Background and Related Work

Low-rank optimization includes weight-subspace and gradient-subspace approaches, with gradient projection addressing optimizer memory. Ori differs from prior online approaches by using streaming k-PCA rather than a secondary optimizer for projector updates.

The Landscape: Low-Rank Optimization

  • Weight-subspace methods constrain model weights, while gradient-subspace methods project gradients onto a low-rank basis; GaLore pioneered the latter with periodic SVD.Other gradient-subspace variants use quantization, structured sparsity, random projections, or importance sampling.Source: §2 · p. 3

The Mechanism: Online Subspace Tracking

  • Prior online subspace methods use auxiliary optimization, momentum in a low-dimensional subspace, block power iteration, or Grassmannian geodesic updates.The cited approaches include Subset-Norm and Subspace-Momentum [23], block power iteration [28], geodesic updates [27], and a secondary optimizer [15].Source: §2 · p. 3
  • Ori uses streaming Oja++ updates instead of a secondary optimizer, combining gap-free static convergence assurances with dynamic-regret bounds under drift.Its lower computational cost than manifold gradient descent enables more frequent projector updates.Source: §2 · p. 3

3 Ori Optimizer

Ori treats gradient-subspace estimation as continuous tracking, motivated by empirical evidence that subspaces drift gradually after warm-up. It uses a warm-started Oja++ update to track this change, and also adapts the approach for LDAdam.

3.1 The Phenomenon of Gradient Subspace Drift

  • After warm-up, gradient-subspace overlap stabilizes between 0.55 and 0.75 across ranks 8, 32, and 128, above random-subspace overlap but below 1.0.The measurements use a 60M LLaMA model during C4 pre-training, comparing subspaces 10 steps apart.Source: §3.1 · p. 4, Figure 2 · p. 4
  • The overlap remains stable across ranks while staying below 1.0, indicating gradual but nonzero subspace change that online tracking can exploit.The observed temporal coherence motivates continuous tracking rather than treating each subspace estimate as independent.Source: §3.1 · p. 4

3.2 Online Subspace Tracking via Oja++ with Moment Transfer

  • Oja-style streaming k-PCA incrementally rotates the current basis using new gradient-derived samples, followed by re-orthogonalization to maintain the basis properties.Ori uses Oja++ for its gap-free convergence property, which is relevant when eigengaps may be small in LLM training.Source: §3.2 · pp. 4–5
  • Ori initializes each subspace update from the previous basis, interpolated with a random Gaussian matrix using drift coefficient ρ.The resulting Qinit is the starting point for the update rule; Oja++ uses the previous basis as a warm start and has gap-free convergence, independent of the eigengap.Source: §3.2 · p. 5
  • O-LDAdam adds multi-step stochastic power iteration to LDAdam's update routine to support subspace tracking at longer update intervals.LDAdam's performance degrades when its update interval increases from T = 1 to T = 200; the authors attribute this to single-step power iteration being insufficient over large intervals.Source: §3.2 · p. 5

4 Dynamic Regret Analysis with Oja

The analysis asks whether online Oja can track a changing low-rank comparator, rather than converge to a fixed principal subspace. Under a bounded-drift setting, it establishes sublinear regret in the noiseless case.

4 Dynamic Regret Analysis with Oja

  • The analysis frames streaming k-PCA as tracking a drifting low-rank comparator under concept drift, with dynamic regret as the performance measure.The environment has S steps of T updates, and each streaming sample is drawn i.i.d. from an unknown distribution.Source: §4 · p. 6
  • In the noiseless setting, regret is sublinear in the total ST updates, and its dependence on S degrades as drift λ increases.The guarantee concerns noisy Oja under λ-drift, with the noiseless case specified by Etot = 0.Source: §4 · p. 6

5 Experiments

The experiments compare Ori with full- and low-rank optimizers across LLaMA pretraining and RoBERTa fine-tuning, emphasizing constrained ranks. They also test Ori’s design choices and whether its subspace-tracking approach can extend to LDAdam.

5.1 Baselines and Experimental Setup

  • The experiments compare Adam-based full-rank, weight-subspace, and gradient-subspace optimizers across LLaMA model sizes and ranks, using validation perplexity for pretraining and GLUE task metrics for fine-tuning.The pretraining evaluation spans 60M–1B LLaMA models and three ranks per model; the GLUE evaluation reports task-specific metrics over three random-seed runs.Source: §5 · p. 6, §5.1 · pp. 6–7, §5.2 · p. 7

5.2 Main results: LLaMA pre-training

  • At rank 16, Ori reaches validation perplexity 43.61, versus 60.29 for GaLore and 54.05 for OSD on 130M LLaMA pretraining.Across model scales, Ori’s advantage is strongest at severely constrained ranks; it is on par with GaLore at some less restrictive ranks.Source: §5.2 · pp. 7–8, Table 1 · p. 7
  • At T = 200 and rank 32 on LLaMA-1B, O-LDAdam reaches perplexity 15.66, versus 20.27 for LDAdam at T = 1.Raising vanilla LDAdam’s interval to T = 200 increased perplexity to 45.66; O-LDAdam applies Ori-inspired stochastic power iteration.Source: Table 2 · p. 7, §5.2 · p. 8

5.3 Finetuning

  • Ori leads 10 of 16 GLUE task/rank combinations and has the highest average score among low-rank methods at both tested ranks.At rank 8, its average score is 86.30 versus 84.34 for GaLore, 79.68 for LoRA, and 85.82 for full-rank Adam.Source: §5.3 · p. 8, Table 3 · p. 8

5.4 Ablation Study

  • Across ranks, ρ = 0.99 performs best; both a random restart (ρ = 0.0) and direct carryover (ρ = 1.0) perform worse.Ori uses ρ = 0.99 in its pretraining and fine-tuning experiments.Source: §5.4 · p. 9, Figure 3 · p. 9
  • More frequent projector updates yield lower loss, with stale-projector costs especially pronounced at rank 8.Ori uses T = 20 as a balance between performance and computational overhead; the ablation varies the period from 10 to 200.Source: §5.4 · p. 9, Figure 3 · p. 9
  • Oja++ consistently outperforms standard Oja across all tested ranks as the gradient-projector update subroutine.The comparison is reported for the 60M LLaMA model pretrained on C4 for 0.44B tokens.Source: §5.4 · p. 9, Figure 3 · p. 9

6 Conclusion

Ori continuously tracks the gradient subspace with streaming k-PCA and preserves optimizer statistics across projector updates, narrowing the gap to full-rank training in the extremely low-rank regime while remaining faster and more memory-efficient than prior gradient-projection methods.

6 Conclusion

  • Ori continuously tracks the gradient subspace with streaming k-PCA and preserves optimizer statistics across projector updates, mitigating staleness while keeping memory and compute low.The conclusion also identifies behavior under heterogeneous training data and explicit optimizer-state rotation across basis changes as limitations and future directions.Source: §6 · p. 10

A Proofs

  • Theorem 2 converts the drift assumption into an angle bound between subspaces.The proof links steps through drift in angle space.Source: A Proofs · p. 13

A.1 Proof of Lemma 2: Distance Identity

  • The projector distance identity expresses squared Frobenius distance between rank-k subspaces through their principal angles.For orthonormal bases, the squared singular values of Q1ᵀQ2 are the squared cosines of those angles.Source: A Proofs · p. 13, §A.1 · p. 14

A.2 Detailed proof of equation 8 and equation 9(one-step potential drop for Oja)

  • Oja’s update yields a one-step potential-drop inequality tied to instantaneous regret, with second-order error controlled under bounded input data.The analysis attributes the higher-order term to the pre-retraction update and QR-retraction curvature.Source: A Proofs · p. 13, §A.2 · pp. 14–15

A.3 Proof of Theorem 1

Theorem 1’s proof bounds regret per iteration, sums that bound within steps, and links steps through subspace drift.

A.3 Proof of Theorem 1

  • Rearranging Equation 21 gives a per-iteration regret bound, which is summed over the iterations within each step.The proof then sums the resulting step bounds across steps using drift linking.Source: §A.3 · p. 16
  • Across steps, the comparison is against a local comparator, and the proof uses projector distance and the drift assumption to link successive steps.With orthonormal initialization, the bound’s initialization constant changes from k² to k.Source: §A.3 · pp. 16–17

B.1 Low-Rank Training Paradigms

The related-work discussion distinguishes parameter-subspace training from gradient-subspace projection, then situates Ori among online tracking methods and concept-drift theory.

B.1 Low-Rank Training Paradigms

  • Parameter-subspace methods learn, refresh, or constrain low-rank weight bases, whereas gradient-subspace methods project gradients to reduce memory and computation.The discussion covers learned factors and structured decompositions as well as methods that keep bases fixed or update them periodically.Source: §B.1 · pp. 17–18

B.3 Concept Drift in Online Learning

  • The paper frames gradient-subspace drift using concept-drift theory and presents its dynamic-regret analysis as the first formal analysis of a low-rank optimizer in that framework.The discussion connects this framing to prior evidence on low-dimensional, slowly evolving gradient subspaces.Source: §B.1 · p. 18, §B.3 · p. 18

B.2 Online Subspace Tracking

  • Online subspace methods use k-PCA tracking or refreshes to maintain a low-dimensional basis, with prior work reporting lower pretraining perplexity and a narrower gap to full-rank training.Subspace-Momentum confines momentum to a fixed subspace and uses SGD in its orthogonal complement; Online Subspace Descent continuously adapts its projection matrix.Source: §B.2 · p. 18

C.1 AdamW with Ori

  • Ori’s projector update subroutine is not tied to AdamW and can be integrated with other optimizers, which the paper leaves for future exploration.The illustrated algorithm equips AdamW with Ori projector updates.Source: §C.1 · p. 19

C.2 Ori Hyperparameters and Recommended Search Strategy

  • The recommended settings include drift coefficient ρ = 0.99 and projector update period Tperiod = 20, while warmup T warmup = 20 is reported as stable.The guidance also recommends η = 1.0 and Nupdate = 2 across scenarios.Source: §C.2 · p. 20

C.3 Amortized Complexity Analysis

Ori’s per-step projection and lifting costs scale as O(dnr), while its periodic projector-update overhead is amortized over the update interval.

C.3 Amortized Complexity Analysis

  • Each training step’s gradient projection and lifting cost O(dnr), while the Adam update in projected space costs O(rn).The setup defines d × n as weight dimensions and r as the rank.Source: §C.3 · p. 20
  • At rank 128 with d = n = 2048, Ori uses 9.4% of full-rank Adam’s memory: 0.2506 GB versus 2.4942 GB.The stated per-layer storage is O(r(d + 2n)).Source: §C.3 · p. 20
  • Despite updating 10× more frequently, Ori is the fastest gradient-subspace method across model sizes and ranks.The passage points to Table 11 for the comparison.Source: §C.3 · p. 20

D Stochastic Power Iteration for LDAdam

This section adapts LDAdam’s projector update to Ori-style online tracking, addressing the sensitivity of single-step power iteration to long update intervals. The resulting O-LDAdam combines stochastic subspace updates with LDAdam’s projection-aware optimization framework.

D Stochastic Power Iteration for LDAdam

  • LDAdam uses warm-started block power iteration to compute gradient subspaces and projection-aware rules to transfer Adam’s first- and second-moment states.It also uses generalized error feedback to compensate for information lost when compressing gradients and optimizer states.Source: D Stochastic Power Iteration for LDAdam · p. 21
  • With T = 200, LDAdam may fail to converge on 1B LLaMA, exceeding 160 perplexity at both ranks r = 512 and r = 128.LDAdam typically requires per-step projector updates to maintain accuracy, which can be computationally prohibitive.Source: D Stochastic Power Iteration for LDAdam · p. 21
  • O-LDAdam replaces LDAdam’s single-step full-block update with stochastic power iteration, sampling a row or column vector from the input matrix at each step.The update is treated as an Ori-style Oja rule with an identity carry-over term; stochastic power iteration sets α = 0.Source: D Stochastic Power Iteration for LDAdam · p. 21
  • Table 4 lists the accuracy metrics used for the RoBERTa-Base fine-tuning results in Table 3.The O-LDAdam hyperparameter table covers its Ori-inspired stochastic power-iteration update for the GLUE fine-tuning experiments.Source: Table 4 · p. 22

E.1 Dataset and Training Details

The appendix documents datasets, training configurations, compute resources, and procedures used to visualize gradient-subspace drift, run ablations, and measure optimizer memory. It also identifies the tables containing Ori and O-LDAdam hyperparameters.

E.1 Dataset and Training Details

  • LLaMA pretraining uses C4 without data repetition and reports final validation perplexity; RoBERTa-Base fine-tuning uses GLUE and reports final accuracy.GLUE task metrics are averaged over three random-seed runs, with standard deviations reported.Source: §E.1 · p. 22

E.2 Compute resources

  • The 350M and 1B LLaMA pretraining runs use nodes equipped with four H100 or four MI210 GPUs.The experiments also use a range of other GPU hardware, including RTX, A100, and MI-series devices.Source: §E.2 · p. 22

E.3 Ori hyperparameters

  • Tables 5 and 6 list Ori hyperparameters for LLaMA-C4 pretraining and RoBERTa-Base-GLUE fine-tuning, respectively.Pretraining uses an effective batch size of 512, sequence length 256, cosine learning-rate schedule, and bfloat16.Source: §E.3 · p. 22, §E.1 · p. 22, Table 6 · p. 23

E.4 O-LDAdam hyperparameters

  • Tables 7 and 8 list the O-LDAdam hyperparameters for LLaMA-C4 pretraining and RoBERTa-Base-GLUE fine-tuning; other settings follow baseline LDAdam.Fine-tuning runs for 30 epochs with one gradient-accumulation step, sequence length 512, and a linear learning-rate schedule.Source: §E.1 · p. 22, §E.4 · p. 23, Table 7 · p. 23

E.5 Baseline hyperparameters

  • Baseline hyperparameters are described in Section 5.1.This appendix section directs readers to the main experiment section for those settings.Source: §E.5 · p. 23

E.6 Visualizing subspace drift

  • The drift visualization measures mean gradient-subspace overlap across layers during 1,000 Adam training steps on 60M LLaMA, using Δt = 10 and ranks 128, 32, and 8.Random matrices matched in shape and Frobenius norm provide comparison baselines for the overlap values.Source: §E.6 · p. 23

E.7 Ablation study details

  • The ablation study reuses the 60M LLaMA-C4 setting; for fairness, standard Oja receives 16, 12, or 8 updates at ranks 128, 32, or 8, respectively.Ori uses Nupdate = 2; the Oja baseline’s update count matches the total inner updates in the default Oja++ configuration.Source: §E.7 · pp. 23–24

E.8 Measuring optimizer state peak memory

  • Optimizer-step peak memory is measured separately by subtracting memory allocated after backward from the peak memory recorded during the optimizer step.The procedure resets the peak-memory counter after measuring post-backward allocation.Source: §E.8 · p. 24

F.1 Memory and Computation Time

Ori’s efficiency results compare optimizer peak memory and wall-clock time on LLaMA-C4 pretraining, while O-LDAdam is compared with the error-feedback baseline LDAdam. Ori’s frequency and memory advantages are reported alongside trade-offs in speed and error-feedback overhead.

F.1 Memory and Computation Time

  • At rank 32, Ori uses 0.2487 GB of optimizer peak memory, 12.1% less than OSD (0.2832 GB) and 26.5% less than GaLore (0.3385 GB).Ori consistently uses the least optimizer peak memory across ranks among the standard non-error-feedback baselines.Source: Table 9 · p. 25, §F.1 · p. 25
  • Despite updating its projector 10x more frequently (TOri = 20 vs. TGaLore/OSD = 200), Ori is the fastest standard non-error-feedback method.ReLoRA is faster at lower ranks, but the passage notes worse model performance; the comparison is presented as a speed-performance trade-off.Source: Table 9 · p. 25, §F.1 · p. 25
  • O-LDAdam has the same memory footprint as LDAdam at T = 200, and its extra power-iteration time is described as not prohibitively large.Compared with LDAdam’s per-step schedule (T = 1), O-LDAdam provides a speedup while balancing performance and computational cost in the extremely low-rank regime.Source: Table 10 · p. 25, §F.1 · pp. 25–26
  • Table 11 compares memory and wall-clock time across LLaMA model sizes and reports Ori as generally efficient across the tested settings.The accompanying efficiency discussion covers both optimizer peak memory and pretraining wall-clock time.Source: §F.1 · pp. 25–26, Table 11 · p. 26

G Limitations

The authors identify limits in the evaluation scope, optimizer-state handling, and rank selection, while describing broader societal effects as uncertain and long-term. The discussion points to heterogeneous-data evaluation, explicit state projection, and adaptive rank as future directions.

G Limitations

  • Evaluation covers single-domain C4 pretraining and GLUE fine-tuning, leaving Ori’s performance on highly heterogeneous data mixtures empirically unvalidated.The authors note that prevalent domains could dominate top-k directions and suppress rarer domain-specific signals, while acknowledging this remains an open question.Source: G Limitations · p. 26
  • Ori does not explicitly rotate Adam moments when the basis changes, instead relying on slow drift and exponential averaging to decay stale information.The authors suggest explicit state projection could improve performance, particularly with faster subspace drift or lower ρ.Source: G Limitations · pp. 26–27
  • Ori keeps rank r fixed throughout training, so the rank must be specified in advance.Combining streaming tracking with rank-adaptive methods is proposed as a future extension.Source: G Limitations · p. 27

H Broader Social Impact

  • The authors foresee no immediate negative societal impacts, while noting possible long-term implications involving energy use or downstream applications.Source: H Broader Social Impact · p. 27

Related papers

Loading 2610.22794v1…