Clean Cuts Second-Order Optimizer Memory to Linear via Nyström Sketching, Enabling 13B Pre-training on a Single 80GB GPU

Clean replaces SOAP's explicitly accumulated full gradient covariances with randomized Nyström approximations, storing only thin factors so preconditioner memory scales linearly rather than quadratically with layer dimensions. Prior SOAP, Shampoo, and several efficient optimizers carry at least one optimizer state growing quadratically with model dimensions; Clean is described as the first second-order SOAP-like optimizer with linear memory. The paper provides a state-size comparison table and reports optimizer memory in LLaMA-350M and 1.3B pre-training: Clean at 1.25G and 4.48G, SOAP at 7.64G and 13.62G, AdamW at 1.40G and 5.11G.

Clean does not discard gradients outside the low-rank core; it splits the update into four orthogonal blocks Main, A, B, C, applies full Adam to Main, pooled second moments to A and B, and raw gradient with a scalar second moment to C, then projects all blocks back to form a full-rank update. Low-rank projection methods typically keep only the core subspace; Clean explicitly reintegrates residual blocks and uses projection-aware moment transport across basis changes. The authors measure each block's share of gradient norm on LLaMA-350M, finding Main dominates in nearly all layers while C contributes less, yet non-Main blocks carry substantial signal especially in later iterations, motivating the per-block moment design.

Q-Clean quantizes optimizer states: first moments and projection matrices use INT8 block-wise asymmetric linear quantization with stochastic rounding, the main second moment uses bitsandbytes non-linear quantization, and A and B second moments stay in full precision. This further amplifies Clean's memory advantage while stochastic rounding prevents small updates from vanishing under round-to-nearest. In LLaMA-1.3B pre-training, Q-Clean uses 1.28G optimizer memory versus Muon's 2.80G, a reduction over 50%, with perplexity 12.10, on par with full-rank methods.

Linear memory makes Clean and Q-Clean the only evaluated optimizers that allow pre-training a 13B-parameter model on a single 80GB H100, and Clean reaches AdamW's final performance 26% faster in wall-clock time. This offers a path to training larger foundation models on more accessible hardware, whereas the low-rank second-order baseline ALSO still exceeds 80GB in the same setting. The paper reports Clean reaching AdamW's final performance in 64 hours versus 81, consuming only 77% of the token budget, with Q-Clean at 87%; the 13B single-GPU claim comes from the memory analysis in Appendix D.

The results target from-scratch pre-training of LLaMA architectures on C4, covering 350M and 1.3B parameters with sequence length 1024 and effective batch size 4096 under the Chinchilla rule; the 13B single-GPU claim comes from the memory analysis in Appendix D rather than a full pre-training run. The method applies to matrix weight layers, while 1-D parameters still receive Adam updates. For researchers and engineering teams seeking to cut optimizer-state memory while retaining two-sided curvature, this offers a reusable block-wise moment design and quantization recipe.

The 1.3B learning rate used a shortened sweep and only the best configuration was resumed, so hyperparameter sensitivity at larger scale remains to be seen; how much error the pooled moments and approximate transport across basis changes introduce is discussed mainly via Figure 3 and the appendix, without per-block error bounds; the 13B single-GPU claim rests on memory analysis with no full-scale pre-training perplexity; and quantization stability under longer training or different architectures remains an open question.