arXiv preprint ·
Clean: Second-order LLM Training at Linear Memory Cost via Nyström Sketching
- Published
- Source
- arXiv preprint, posted before peer review
- Original
- Read it at arXiv
- Topics
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.
The authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.
Read it on ReadSearch
Each opens after you create an account, and brings you back to this paper.
More in Machine Learning
Principled Top-$k$ Selection for Language Models with Hybrid Gradients
Agent Behavior as Code: Efficient and Robust LLM Agents with Programmatic Specifications
DimSteer: Steering LLM Authoring with Automatically Discovered Stylistic Controls
Viva La Vida: Verification and Accumulation Failures in Multi-Agent Proof Search
Language Model Activations Inhabit Privileged Error-Correcting Basins
Get the few papers worth your time in Machine Learning every morning.
ReadSearch reads the new papers in your field each day and brings you the few worth reading, each with its key finding in a sentence.