Clean optimizer cuts memory use and speeds training for large language models
Clean: Second-order LLM Training at Linear Memory Cost via Nyström Sketching
Machine LearningArtificial IntelligenceComputation and Language
Summary
Training very large language models usually requires a trade-off between memory use and the speed of learning. The authors introduce Clean, a new optimizer that uses a clever math trick called Nystrom sketching to keep important information about the model’s parameters without needing too much memory. This method helps the training run faster and with less space than existing methods like Adam. They also created a smaller, compressed version called Q-Clean, which further reduces memory needs while maintaining strong performance.
What this means in practice
- •For machine learning engineers: Train very large language models on limited GPU memory by using Clean to reduce optimizer memory needs and speed up training.
- •For cloud ai service providers: Offer faster and more memory-efficient model training options to clients by integrating Clean and Q-Clean optimizers in distributed training platforms.$Commercial implications: Enables selling enhanced AI training services that handle larger models on fewer resources, lowering client costs and cloud usage.
Authors
Beheshteh T. Rakhshan, Sahar Rajabi, Maziar Sargordi Shikai Fang, Guillaume Rabusseau, Sirisha Rambhatla
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.