PaperlayerView exact revision on arXiv

Source-linked AI summary

Clean: Second-order LLM Training at Linear Memory Cost via Nystr\"om Sketching

Beheshteh T. Rakhshan, Sahar Rajabi, Maziar Sargordi Shikai Fang, Guillaume Rabusseau, Sirisha Rambhatla

arXiv:2610.04204v1cs.LGcs.AIcs.CL

TL;DR

LLM optimizers trade memory efficiency against the curvature information retained for updates. Clean uses randomized Nyström approximations and reintegrates off-subspace components, while Q-Clean compresses optimizer states; the methods are reported to reduce memory while preserving competitive performance and enabling 13B-model pre-training on one 80GB GPU.

  • Problem

    Memory-efficient optimizers can discard cross-parameter curvature, while full-curvature methods such as SOAP can accelerate convergence at prohibitive memory cost.

    Source: §1 · pp. 1–2

  • Method

    Clean approximates SOAP's left and right preconditioners with randomized Nyström sketching and recombines four gradient blocks, while Q-Clean quantizes optimizer states.

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

  • Results

    Clean reaches AdamW's final performance 26% faster in wall-clock time, while Q-Clean cuts optimizer-state memory by more than 50% compared with Muon and retains results on par with full-rank methods.

    Source: §6 · p. 8, §7 · p. 10

  • Takeaways & Limitations

    Clean and Q-Clean are reported to enable pre-training a 13B-parameter model on a single 80GB GPU.

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

    Clean's covariance averaging relies on β2 ∈ [0, 1], for which each update is a convex combination of PSD matrices.

    Source: §3 · p. 4

Abstract

from arXiv · show

Training large language models (LLMs) entails a fundamental trade-off: memory-efficient optimizers such as Adam discard cross-parameter curvature, whereas full-curvature methods such as SOAP can accelerate convergence at prohibitive memory costs. We introduce Clean, a memory-efficient and full-curvature optimizer designed to resolve this bottleneck. Clean leverages the randomized Nystrom method to accurately approximate the left and right preconditioners in SOAP, and to reduce the optimizer's memory complexity from quadratic to linear in terms of model dimensions. We subsequently reintegrate the off-subspace components to capture curvature information beyond the low-rank approximation, preserving rich curvature at minimal memory cost. We further propose Q-Clean, a low-precision variant that aggressively compresses optimizer states. Q-Clean reduces optimizer memory consumption by \textbf{over 50\%} compared to Muon when pre-training a LLaMA-1.3B architecture, all while maintaining strong and competitive predictive performance. Notably, Clean operates with a smaller optimizer-state footprint than standard AdamW while reaching AdamW's final performance \textbf{26\% faster} in wall-clock time. Furthermore, our methods uniquely enable the pre-training of a 13B-parameter model on a single 80GB GPU, providing a scalable, efficient, and accessible approach to large-scale model optimization.

1 INTRODUCTION

Clean targets the memory–curvature trade-off in LLM optimization with linear-memory SOAP-style preconditioning and updates that retain off-subspace information. Clean and Q-Clean are reported to achieve competitive performance with substantially reduced optimizer memory.

1 INTRODUCTION

  • Memory-efficient optimizers can discard cross-parameter curvature, whereas full-curvature methods such as SOAP can accelerate convergence at prohibitive memory cost.The introduction frames balancing curvature fidelity and memory efficiency as a central challenge for efficient LLM training.Source: §1 · pp. 1–2
  • Clean's optimizer-state memory scales linearly with layer dimensions, requiring O((m + n)r) rather than O(mn), while retaining full-rank second-order preconditioning.The introduction presents Clean as a curvature-aware approach to the memory-versus-curvature trade-off.Source: §1 · p. 2
  • Clean uses randomized Nyström approximations for SOAP-style left and right preconditioners, then updates and recombines four orthogonal gradient blocks to retain off-subspace information.Q-Clean is a quantized variant that compresses stored optimizer states.Source: §1 · p. 2
  • Clean and Q-Clean are reported to achieve validation perplexities highly competitive with SOAP while using less optimizer memory than AdamW.The reported comparison also states that Clean reaches AdamW's final performance 26% faster in wall-clock time.Source: §1 · p. 2

2 BACKGROUND AND PRELIMINARIES

The background frames optimizers as preconditioned gradient methods that trade retained curvature information against memory and computation. It reviews diagonal and Kronecker-factored approaches, then introduces randomized Nyström approximation as a low-rank representation for symmetric PSD matrices.

2 BACKGROUND AND PRELIMINARIES

  • A preconditioned gradient method updates parameters using a learning rate and a positive semidefinite preconditioner that encodes curvature information.The preconditioner choice balances curvature accuracy with memory and computational costs.Source: §2 · pp. 2–3

2.1 OPTIMIZERS AS PRECONDITIONERS

  • Diagonal preconditioners such as Adam capture individual parameter scales with O(mn) memory but ignore cross-parameter correlations; full-matrix AdaGrad requires O(m^2n^2) memory.Kronecker-factored Shampoo stores left and right covariance factors, which still incur quadratic memory in layer width.Source: §2.1 · p. 3
  • SOAP projects gradients into the eigenbasis of Shampoo's covariance factors, applies coordinate-wise Adam, and projects the step back, but retains O(m^2 + n^2) memory.The passage reports that SOAP can reduce iterations and wall-clock time compared with Shampoo and AdamW.Source: §2.1 · p. 3

2.2 RANDOMIZED NYSTRÖM PRECONDITIONER

  • Randomized Nyström approximation represents a symmetric PSD matrix with a rank-r surrogate that requires O(nr) rather than O(n^2) storage when r is much smaller than n.A naive construction is numerically unstable; the text points to a numerically stable construction in the appendix.Source: §2.2 · pp. 3–4

3 METHODOLOGY

Clean uses low-rank Nyström factors to approximate SOAP’s left and right covariance preconditioners without forming the full covariance matrices. It retains information outside the low-rank core through block-wise updates, while tailoring stored moments to keep optimizer memory linear in layer dimensions.

Low-rank covariance factors

  • Clean stores rank-r Nyström factors rather than full covariance matrices, reducing their storage cost to (m + n + 2)r for a layer of dimensions m × n.The factors are refreshed every T iterations and reused between refreshes.Source: §3 · p. 4
  • Clean refreshes its Nyström factors every T steps and reuses them between refreshes; averaging covariance updates keeps the factors’ basis evolution gradual.The covariance averages remain positive semidefinite for β2 in [0, 1], so they admit randomized Nyström approximations.Source: §3 · pp. 4, 7

Block-wise low-rank preconditioning

  • The gradient update is split into Main, A, B, and C blocks, preserving components outside the low-rank core instead of discarding them.Main is dominant across nearly all layers, while blocks beyond Main carry a substantial share of the training signal, especially later in training.Source: §3 · p. 5, Figure 3 · p. 5
  • Adam-style moment storage is tailored by block: Main keeps full moments, A and B pool second moments, and C uses no first moment and a scalar second moment.This avoids per-entry optimizer states for the full-sized C block, which would cost O(mn).Source: §3 · p. 5, Table 1 · p. 6

Projection-aware adaptive moments

  • When refreshed bases change, Clean rotates first moments and transports pooled A/B second moments, but leaves Main’s element-wise second moment in place because exact rotation needs unrecorded cross-terms.Warm-start refreshes constrain frame drift to keep the stale-basis error in that second moment sufficiently small.Source: §3 · p. 6

Projection back and parameter update

  • After block-wise Adam-style updates, Clean projects each block back to parameter space and recombines them into the final update.The projection uses bias-corrected momentum terms and Adam scaling factors.Source: §3 · p. 6

Memory efficiency

  • Clean avoids O(mn) optimizer memory while retaining full-rank information, and is the only evaluated optimizer reported to scale linearly in state memory.Its warm-start Nyström variant enables gradual basis evolution during periodic subspace refreshes.Source: §3 · p. 6

4 Q-CLEAN: LOW-PRECISION SECOND-ORDER OPTIMIZATION

Q-Clean compresses optimizer states to reduce memory overhead while retaining Clean’s second-order optimization approach. It applies different quantization schemes to state types and uses stochastic rounding to avoid systematically losing small updates.

Q-Clean

  • Q-Clean stores optimizer states in low precision to further reduce Clean’s memory overhead.The section introduces Q-Clean as a quantized-state extension of Clean.Source: §4 · p. 7

Quantization

  • Stochastic rounding is unbiased in expectation and prevents small but consistent optimizer-state updates from vanishing over time.Round-to-nearest discards updates smaller than half a quantization step.Source: §4 · p. 7

State quantization

  • Q-Clean uses block-wise linear asymmetric quantization for first moments and projection matrices, but nonlinear quantization for the Main second moment.A and B second moments remain full precision because each is a vector of size r and quantizing it would yield no meaningful savings.Source: §4 · p. 7

5 RELATED WORK

Related work spans low-rank gradient projection, second-order preconditioning, low-rank second-order methods, and quantization-based memory reduction. The paper positions Clean among approaches that combine two-sided curvature information with memory-efficient updates.

Low-rank Projection Methods

  • Low-rank gradient projection methods save memory, but one-sided projections can be unstable as the subspace drifts and do not capture row-column covariance geometry.LDAdam (Robert et al., 2024) and Subtrack++ (Rajabi et al., 2025) add projection-aware updates and more robust subspace tracking.Source: §5 · p. 7

Second-order Preconditioning

  • Shampoo and K-FAC use two-sided curvature preconditioning but require full-sized covariance matrices; SOAP (Vyas et al., 2024) retains full-matrix accumulators.SOAP combines Shampoo preconditioning with first-order updates.Source: §5 · pp. 7–8

Low-rank Second-order Methods

  • Recent low-rank second-order methods approximate full-rank optimizers or SOAP, while COSMOS (Liu et al., 2025) applies SOAP only to leading eigensubspaces.ALSO compensates for residual blocks outside its core low-rank projection.Source: §5 · p. 8

Low Precision Training with Quantization

  • Quantization work includes PTQ and QAT; examples include Q-Galore (Zhang et al., 2024) and 4-bit Shampoo.The passage identifies training instability as a major QAT challenge and notes SWALP’s memory overhead for large foundation models.Source: §5 · p. 8

6 EXPERIMENTS

The experiments compare Clean and Q-Clean with established optimizers in LLaMA pre-training, finding competitive validation perplexity with lower optimizer-memory costs. Clean also reaches AdamW’s final performance faster, while ablations examine rank and preconditioner-refresh settings.

Setup

  • Models are trained from scratch on C4 at 350M and 1.3B parameters, using 1,750 and 5,000 steps, respectively, with up to four H100 80GB GPUs.Comparisons include AdamW, SOAP, Muon, and ALSO; the 1.3B learning-rate sweep is shortened due to its cost.Source: §6 · p. 8

Results

  • Clean achieves competitive final performance compared to SOAP, while Q-Clean cuts optimizer-state memory by more than 50% compared to Muon and remains on par with full-rank methods.Table 3 summarizes final validation perplexity and optimizer memory across LLaMA model sizes.Source: §6 · p. 8, Table 3 · p. 9

Efficiency

  • Clean reaches AdamW’s final validation loss in 64 of 81 hours, a 26% speedup, and uses 77% of the training token budget; Q-Clean uses 87%.The runs use four H100 GPUs with the same batch size, schedule, and token budget.Source: Figure 5 · p. 9, §6 · p. 9

Ablations

  • On LLaMA-350M, validation perplexity improves as rank rises to 256; r = 128 is competitive, while r = 512 adds no gains, and T = 10 balances refresh frequency.Both T = 1 and T = 100 negatively affect final performance; the ablations vary rank or update interval while holding the other at its default.Source: §6 · p. 9

7 CONCLUSION

Clean uses randomized Nyström sketching to make SOAP-style two-sided preconditioning more memory efficient while retaining off-subspace training signals. The conclusion reports competitive validation perplexity, faster convergence than AdamW, and single-GPU pre-training of a 13B model.

Conclusion

  • Clean approximates left and right preconditioners with randomized Nyström sketching, reducing two-sided preconditioning memory complexity from quadratic to linear and reintegrating off-subspace residual blocks.Q-Clean is a low-precision variant that compresses optimizer states.Source: §7 · p. 9
  • Clean and Q-Clean achieve validation perplexities highly competitive with SOAP while using less optimizer memory than standard AdamW.The methods enable pre-training a 13B-parameter architecture on a single 80GB GPU.Source: §7 · p. 10
  • Clean reaches AdamW’s final performance 26% faster in wall-clock time.The conclusion also reports that the methods enable single-GPU pre-training of a 13B-parameter model.Source: §7 · p. 10

AI USE STATEMENT

The supplied material documents AI use and describes a warm-started Nyström procedure alongside Q-Clean’s optimizer-state quantization choices. Block-wise scales limit the precision loss caused by outliers to the affected block.

AI USE STATEMENT

  • The authors report AI assistance with methodology feedback, implementation, literature analysis, figures, and draft revision, and say they reviewed the work and take responsibility for its final content.They state that AI was not used for the listed other tasks requiring disclosure, including generating synthetic datasets or developing theoretical models.Source: AI USE STATEMENT · p. 10

A NYSTRÖM ALGORITHM

  • The Nyström algorithm uses a numerically stable alternative and warm-starts with prior eigenvectors plus fresh Gaussian noise, power iterations, periodic orthogonalization, and sign alignment.Fresh noise explores new gradient directions; sign alignment supports temporal consistency across optimization steps.Source: A NYSTRÖM ALGORITHM · p. 13

B QUANTIZATION DETAILS

  • Q-Clean uses low-precision optimizer states, with asymmetric affine quantization for projection and first-moment matrices and nonlinear symmetric quantization for the main second moment.The second moments of blocks A and B remain full precision because their size-r vectors have negligible memory footprint.Source: B QUANTIZATION DETAILS · p. 14, §B.1 · p. 14

B.1 BLOCK-WISE n-BIT QUANTIZATION

  • Block-wise quantization assigns each contiguous block its own scale and, for asymmetric quantization, zero-point, isolating outlier-driven precision loss to the affected block.For the second moment, the nonlinear codebook provides greater precision for small values while still representing large outliers.Source: §B.1 · p. 14

B.2 MAINTAINING HIGH-PRECISION TRAINING WITH STOCHASTIC ROUNDING ACCUMULATED QUANTIZATION

INT8 optimizer-state storage reduces memory, but round-to-nearest can suppress updates that are too small. The method instead uses stochastic rounding for asymmetric quantization, which is unbiased in expectation.

B.2 MAINTAINING HIGH-PRECISION TRAINING WITH STOCHASTIC ROUNDING ACCUMULATED QUANTIZATION

  • Stochastic rounding probabilistically rounds a value to one of its adjacent quantization levels, making the quantized value unbiased in expectation.The passage contrasts it with round-to-nearest, which can mitigate updates that are too small.Source: §B.2 · p. 15

B.3 QUANTIZER CHOICE

The quantizer choice depends on how statistics are measured: block-wise quantization reduces heavy tails for most matrix families, but the second-moment matrix V remains an outlier and needs nonlinear quantization.

B.3 QUANTIZER CHOICE

  • The analysis tracks scale and dynamic-range statistics across pooled matrices from the LLaMA-350M model optimized with SOAP.It includes weights, gradients, Adam moments, SOAP covariance matrices, and their eigenbases.Source: §B.3 · p. 15
  • Per-matrix statistics summarize each whole matrix, whereas per-block statistics are computed separately for blocks and pooled.The distinction corresponds to sharing one quantization scale across a matrix versus quantizing blocks individually.Source: §B.3 · p. 15
  • Block-wise statistics make W, G, M, L, R, QL, and QR suitable for asymmetric linear quantization, while V still requires a symmetric nonlinear quantizer.Linear quantization of V leads to divergence, unlike for the other matrices.Source: §B.3 · p. 15

C SOAP

SOAP applies Adam to gradients in the covariance eigenbasis, but explicitly stores full covariance matrices and Adam moments, imposing high memory costs. Its plotted matrix statistics also show why V receives a distinct nonlinear quantizer.

C SOAP

  • SOAP constructs and stores full L and R covariance matrices while maintaining full Adam moments, creating prohibitive memory requirements.The gradient is projected onto their eigenbasis for the Adam update.Source: C SOAP · p. 15
  • Per-block statistics reduce heavy tails for most matrices, but V remains a prominent outlier and motivates a nonlinear quantizer.The plotted quantities are non-excess kurtosis, dynamic range, and scale under per-matrix and per-block granularities.Source: Figure 6 · p. 16
  • SOAP projects G into the L and R eigenbasis, updates Adam moments there, and applies the resulting step back in the original parameter space.The algorithm also updates L and R from G and refreshes their eigenvectors periodically.Source: C SOAP · p. 16

D LINEAR MEMORY DISCUSSION

Clean and Q-Clean scale optimizer memory more favorably than the baselines, and they enable LLaMA-13B pre-training on one 80 GB H100. The section also reports rank and update-interval ablations and specifies the shared and optimizer-specific training hyperparameters.

D LINEAR MEMORY DISCUSSION

  • Clean and Q-Clean are the only methods shown to pre-train LLaMA-13B on a single 80 GB H100 GPU.The figure compares total memory by rank; its shaded region is infeasible for all baselines at any rank.Source: D LINEAR MEMORY DISCUSSION · p. 16, Figure 7 · p. 17
  • With rank r=d/8, Clean and Q-Clean are the only methods whose optimizer-state cost per parameter decreases with model size.The figure compares these costs across model sizes; Clean keeps no O(mn) or O(m2 + n2) matrix, unlike the other optimizers' O(mn) component.Source: Figure 7 · p. 17, D LINEAR MEMORY DISCUSSION · p. 17
  • 15.7× at r=32 and 3.8× at r=512 are the reported optimizer-memory reductions over ALSO on 13B.At r=32, the figure attributes the widening savings gap to omitting the mn term; it reports batch 1, sequence length 1024, and bf16 state.Source: Figure 8 · p. 17

E ABLATING HYPERPARAMETERS

  • On LLaMA-350M, performance improves with rank through 128 and only marginally thereafter, while T=10 is optimal; the reported selected configuration is r=256 and T=10.The figures show learning curves and validation perplexities across ranks and update intervals.Source: Figure 9 · p. 18, E ABLATING HYPERPARAMETERS · p. 18

F HYPERPARAMETERS

  • Table 4 specifies shared architectures and training setup, Table 5 lists three learning rates swept at 350M, and Table 6 gives optimizer-specific hyperparameters.Peak memory is measured on one GPU at the micro-batch size without activation checkpointing; selected learning rates are used in Tables 3 and 6.Source: Table 4 · p. 19, Table 5 · p. 19, Table 6 · p. 20

Related papers

Loading 2610.04204v1…