2025/06/20 by Glentis, Athanasios, Li, Jiaxiang, Han, Andi +1 · 2 citations
Computer Science · #Artificial Intelligence (cs.AI) #FOS: Computer and information sciences #FOS: Mathematics #Machine Learning (cs.LG) #Natural Language Processing Techniques #Optimization and Control (math.OC) #Speech Recognition and Synthesis #Topic Modeling
paper · pdf · doi:10.48550/arxiv.2506.16659
openalex publication_date 2025/06/20 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28
Training large language models (LLMs) relies on adaptive optimizers such as Adam, which introduce extra operations and require significantly more memory to maintain first- and second-order moments than SGD. While recent works such as GaLore, Fira and APOLLO have proposed state-compressed memory-efficient variants, a fundamental question remains: What are the minimum modifications to plain SGD needed to match state-of-the-art pretraining performance? We systematically investigate this question using a bottom-up approach, and identify two simple yet highly (memory- and compute-) efficient techniques: (1) column-wise gradient normalization (normalizing the gradient along the output dimension), that boosts SGD performance without momentum; and (2) applying first-order momentum only to the output layer, where gradient variance is highest. Combining these two techniques lead to SCALE (Stochastic Column-normAlized Last-layer momEntum), a simple optimizer for memory efficient pretraining. Across multiple models (60M-1B), SCALE matches or exceeds the performance of Adam while using only 35-45% of the total memory. It also consistently outperforms memory-efficient optimizers such as GaLore, Fira and APOLLO, making it a strong candidate for large-scale pretraining under memory constraints. For LLaMA 7B, SCALE outperforms the state-of-the-art memory-efficient methods APOLLO and Muon in both perplexity and memory consumption.