Clean: Second-order LLM Training at Linear Memory Cost via Nyström Sketching

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

cs.LGOct 3, 2026

Clean: Second-order LLM Training at Linear Memory Cost via Nyström Sketching

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

Organizations: Mila & DIRO, Université de Montréal · Critical ML, University of Waterloo · Independent Researcher · Zhejiang University

Abstract

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.

Figures & tables

Figure 1: Measured memory consumption for pre-training a LLaMA-13B model on a single 80 GB H100 GPU (maximum sequence length of 1024, batch size of 1). All non-quantized matrices are stored in bfloat16. Activation + Other combines activation memory with an estimated optimizer workspace. Methods marked ∗ ran out of memory, so their optimizer state is estimated.

Figure 2: Clean. (I)

G_{t}

is projected via the estimated left and right preconditioners to yield

G_{t}^{\text{pre}}

, which is then partitioned into four blocks Main , A , B , and C . (II) Adam-style update is applied to each block, maintaining either full or pooled moments depending on the block’s dimensions (only the shaded buffers are stored). (III) The Adam output from each block is projected back into the parameter space and summed to form a full-rank update.

Panel of Figure 3 — Figure 3: Fraction of the gradient norm captured in each of the four blocks. Generally, Main and A are dominant; the contributions of B and C slowly increase only towards the end of training, where the gradients become noisier. Values are averaged across all 24 instances of a layer type in a LLaMA-350M architecture.

Block	Adam’s Input	$1^{st}$ Moment $M$	$2^{nd}$ Moment $V$	Memory
Main	$S_{t}=U_{t,L}^{\mathsf{T}}G_{t}U_{t,R}\in\mathbb{R}^{r\times r}$	$\mathbb{E}_{\beta_{1}}[S_{t}]$	$\mathbb{E}_{\beta_{2}}[S_{t}^{\odot 2}]$	$2r^{2}$
A	$G_{t,A}\in\mathbb{R}^{m\times r}$	$\mathbb{E}_{\beta_{1}}[G_{t,A}]$	$\mathbb{E}_{\beta_{2}}\big[\sum_{i}(G_{t,A})_{ij}^{2}\big]$	$mr+r$
B	$G_{t,B}\in\mathbb{R}^{r\times n}$	$\mathbb{E}_{\beta_{1}}[G_{t,B}]$	$\mathbb{E}_{\beta_{2}}\big[\sum_{j}(G_{t,B})_{ij}^{2}\big]$	$rn+r$
C	$C_{t}\in\mathbb{R}^{m\times n}$	none, raw $C_{t}$	$\mathbb{E}_{\beta_{2}}\big[\\|C_{t}\\|_{F}^{2}\big]$	$1$

Table 1: Moments per block.

\mathbb{E}_{\beta}[\cdot]

denotes an exponential moving average with rate

\beta\in\{\beta_{1},\beta_{2}\}

, and memory counts stored floats. Main keeps both full moments; A and B keep full first moments, but sum their second moments over the long dimension, leaving one value per column or row; C keeps no first moment and a single scalar second moment, since per-entry states would cost

\mathcal{O}(mn)

Optimizer	State size
Adam	$2mn$
SOAP	$2mn+2m^{2}+2n^{2}$
Muon	$mn$
ALSO	$mn+2mr+2nr+~~r^{2}+2r+O(1)$
Clean (Ours)	$2mr+2nr+2r^{2}+4r+O(1)$

Table 2: Optimizer-state memory requirement for an

m\times n

matrix.

Figure 4: Validation perplexity recorded during the training of a LLaMA-1.3B architecture across various optimization methods.

Methods	LLaMA-350M			LLaMA-1.3B
Methods	Perplexity	Opt. Mem.	Peak Mem.	Perplexity	Opt. Mem.	Peak Mem.
AdamW ( Kingma and Ba, 2014 )	19.63	1.40G	47.93G	12.12	5.11G	54.13G
SOAP ( Vyas et al., 2024 )	15.12	7.64G	54.25G	11.52	13.62G	79.02G
Muon ( Jordan et al., 2024 )	16.97	0.83G	47.36G	11.78	2.80G	51.83G
ALSO ( Romashov et al., 2026 )	15.36	1.30G	47.88G	12.19	4.71G	53.82G
Clean (Ours)	15.87	1.25G	47.68G	11.65	4.48G	53.44G

Table 3: Final validation perplexity (

\downarrow

) and optimizer memory when pre-training LLaMA models of different sizes. Clean achieves competitive perplexity compared to full SOAP, while using less memory than AdamW. Clean, Q-Clean and ALSO use rank

r=256

at 350M and

r=512

at 1.3B. Excluding SOAP, the best results are shown in bold and the second best are underlined .

Panel of Figure 5 — Figure 5: Convergence efficiency on pre-training LLaMA-1.3B, all runs on four H100 GPUs with the same batch size, schedule and token budget. (a) Wall-clock: Clean reaches AdamW’s final validation loss in $64$ of the $81$ hours of AdamW requirements, demonstrating 26% speedup. (b) Token efficiency: Clean reaches that loss in $77\%$ of the training budget and Q-Clean in $87\%$ .

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 6 — Figure 6: The plots compare non-excess kurtosis, dynamic range, and scale under two quantization granularities: per-block (a–c) and per-matrix (d–f). The shift from a per-matrix to a per-block approach significantly reduces the heavy-tailed distributions for most matrices, rendering them suitable for linear quantization. However, the second moment matrix ( $V$ , in purple) remains a prominent outlier, exhibiting persistent heavy tails even in the per-block view, which motivates the use of a non-linear quantizer specifically for $V$ .

Figure 7: Left: total memory for pre-training LLaMA-13B on a single 80 GB H100 as a function of rank; the shaded region is infeasible for all of our baselines at any rank. Right: optimizer state per parameter across model sizes at

r\approx d/8

. Every baseline carries an

mn

term (Table 2 ) and so stays flat or rises; Clean and Q-Clean scale as

r(m+n)

and are the only methods whose per-parameter cost decreases with model size.

Figure 8: Measured optimizer state on LLaMA. Baselines carry an

mn

term; ours replaces it with

r(r+2)

, so our state is smaller whenever

mn>r(r{+}2)

. The margin widens with scale because

mn

grows with the model (

2.8

M to

70.8

M from 350M to 13B) while

r(r{+}2)

is fixed by the rank. Left: the shaded gap is the omitted

mn

term at

r{=}32

. Right: reduction over ALSO,

15.7\times

r{=}32

and

3.8\times

r{=}512

on 13B. Batch 1, seq 1024, bf16 state; filled markers measured, hollow closed-form where the run exceeded 80 GB in ALSO.

Panel of Figure 9 — Figure 9: Ablations of Clean on pre-training LLaMA-350M. (a) Rank $r$ : Performance improves as rank increases up to 128 with marginal improvements afterward. (b) Update interval $T$ : $T=10$ is optimal; updating too often ( $T=1$ ) or too rarely ( $T=100$ ) negatively impacts performance.

	LLaMA-350M	LLaMA-1.3B
Architecture
Parameters	368M	1.34B
Hidden dimension	1024	2048
Intermediate dimension	2736	5461
Attention heads	16	32
Layers	24	24

Appendix

Table 4: Model architectures and training setup shared by all optimizers. Gradient accumulation makes up the global batch. Peak memory in Table 3 is measured on a single GPU at this micro-batch size, without activation checkpointing.

Method	Learning rate
AdamW	$7\times 10^{-4}$	$\mathbf{1.4\times 10^{-3}}$	$2.8\times 10^{-3}$
SOAP	$7\times 10^{-4}$	$1.4\times 10^{-3}$	$\mathbf{2.8\times 10^{-3}}$
Muon	$2\times 10^{-3}$	$6\times 10^{-3}$	$\mathbf{1.0\times 10^{-2}}$
ALSO	$1.4\times 10^{-3}$	$\mathbf{4.4\times 10^{-3}}$	$1.4\times 10^{-2}$
Clean (Ours)	$1.4\times 10^{-3}$	$\mathbf{4.4\times 10^{-3}}$	$1.4\times 10^{-2}$
Q-Clean (Ours)	$1.4\times 10^{-3}$	$\mathbf{4.4\times 10^{-3}}$	$1.4\times 10^{-2}$

Appendix

Table 5: Learning rates swept at 350M. Three values per optimizer; the selected value, used in Table 6 and reported in Table 3 , is shown in bold .

Method	Hyperparameter	LLaMA-350M	LLaMA-1.3B
AdamW	Learning rate	$1.4\times 10^{-3}$	$1.4\times 10^{-3}$
AdamW	$(\beta_{1},\beta_{2})$	$(0.95,0.95)$
SOAP	Learning rate	$2.8\times 10^{-3}$	$2.8\times 10^{-3}$
	$(\beta_{1},\beta_{2})$	$(0.95,0.95)$
	Shampoo $\beta$	0.95
	Preconditioning frequency	10

Appendix

Table 6: Optimizer-specific hyperparameters; learning rates are peak values of the schedule in Table 4 .

Explore similar work

May 9, 2026cs.LG

Navigating LLM Valley: From AdamW to Memory-Efficient and Matrix-Based Optimizers

Training large language models requires optimization algorithms that are not only statistically effective, but also computationally and memory efficient at extreme scale. Although Adam remains the dominant optimizer for large-scale language-model pretraining and fine-tuning, recent work has revisited nearly every component of the optimization stack: adaptive moment estimation, decoupled weight decay, memory footprint, curvature approximation, sign-based updates, large-batch stability, low-rank gradient structure, and matrix-wise orthogonalized updates. This survey reviews optimizer design for large language models through a systems-and-optimization lens. We organize the literature into classical first-order optimizers, adaptive optimizers, memory-efficient variants, second-order and curvature-aware methods, sign-based and discovered optimizers, low-rank and projection-based methods, and matrix-based optimizers such as Muon. We also discuss benchmarking methodology, including hyperparameter fairness, scale dependence, wall-clock efficiency, token efficiency, memory overhead, and downstream evaluation. We argue that optimizer research for LLMs is entering a new phase: moving from single-algorithm speedup claims toward rigorous, scale-aware comparisons that jointly evaluate convergence, stability, memory, and implementation complexity.

Aditya Ranganath

Center for Applied Scientific Computing Lawrence Livermore National Laboratory Livermore, CA, USA

Deep Learning Optimization Memory-Efficient Optimization Matrix Optimization LLM Training

Ask PDF

May 30, 2026cs.LG

Memory-Efficient LLM Training with Dynamic Sparsity: From Stability to Practical Scaling

Dynamic Sparse Training (DST) offers a promising paradigm for improving the training and inference efficiency of deep neural networks; however, we find that in large language model training, DST can suffer from optimization instability, manifested as loss spikes after topology updates. In this work, we show that the naive use of standard Adam-based optimizers leads to a cold-start issue for newly regrown parameters, resulting in excessively large updates and disrupted training dynamics. To address this issue, we propose Sparse Memory-Efficient Training (SMET), which stabilizes DST with optimizer warm-up and improves training progress through density-aware learning-rate scaling. SMET further reduces memory consumption by storing gradients and optimizer states only for active parameters. We provide a theoretical analysis of the update behaviors under SMET, showing improved optimization stability. Extensive experiments demonstrate that SMET enables stable, scalable, and memory-efficient sparse pre-training of LLMs, paving the way for sparse training as a practical alternative to dense training. Our code is publicly available at: https://github.com/QiaoXiao7282/SMET.

Qiao Xiao, Boqian Wu, Patrik Okanovic +6

Eindhoven University of Technology · University of Luxembourg · University of Twente +1

Language Model Pretraining Efficient Language Model Training LLM Training Memory-Efficient Training

Ask PDF

Aug 16, 2026cs.LG

SubZero+: Memory-Efficient Adaptive Zeroth-Order LLM Fine-Tuning in Random Subspaces

Zeroth-order (ZO) optimization with SGD in random subspaces enables memory-efficient fine-tuning of large language models without backpropagation. However, high gradient estimation noise fundamentally undermines adaptive optimizers like Adam. We propose SubZero+, which achieves practical adaptive ZO optimization through a carefully designed dual low-dimensionality strategy: (i) multi-query forward-difference gradient estimation in periodically refreshed random subspaces to mitigate noise amplification in moment buffers, and (ii) Adam updates with periodic restarts performed directly in low-dimensional space rather than full-parameter space. In experiments, this dual design retains memory overhead comparable to momentum-free ZO methods while achieving stronger optimization performance than the evaluated ZO baselines. Theoretically, in the exact-directional limit,

K

-query averaging preserves conditional unbiasedness, while the coefficient estimator's covariance and mean-squared error, as well as query-induced second-moment inflation, scale exactly as

1/K

. Extensive experiments across SuperGLUE with models from 1.3B to 32B parameters under both full fine-tuning and LoRA schemes demonstrate consistent improvements over competing ZO methods. SubZero+ significantly narrows the performance gap with first-order optimization while preserving ZO's inference-time memory efficiency.

Ziming Yu, Shuyao Xiao, Xingyu Zhao +6

Beijing Normal University · Singapore Management University · Xiaomi Inc. +2

Memory-Efficient Fine-Tuning Zeroth-Order Optimization LLM Fine-Tuning Memory-Efficient Optimization

Ask PDF