Research map / Optimization / Optimizers
SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes SWAN, a fully stateless LLM optimizer that preprocesses raw SGD gradients with row-wise normalization (GradNorm) and whitening via a Newton-Schulz variant with diagonal substitution, matching or beating Adam on LLaMA pretraining while using SGD-level memory.
Problem
Adaptive optimizers like Adam store optimizer states that can multiply memory needs, while SGD is stateless but performs poorly for LLM training. The question is whether LLMs can be trained efficiently without storing any optimizer states.
Approach
Compose a row-wise gradient normalization (GradNorm, RMS-style standardization) with a matrix whitening step (GradWhitening, orthogonalization via (GG^T)^{-1/2} G) applied to the instantaneous stochastic gradient, so the update uses only the current gradient; a diagonal-substitution Newton-Schulz scheme makes whitening O(m^2) and a norm rescaling reuses Adam learning rates.
Key results
- On LLaMA 1.3B pretraining on C4, SWAN-ddag reaches validation perplexity 13.56 vs Adam's 16.44 (reproduced) and 15.56 (cited), with SWAN reaching at least 2x token speedup at 350M and 1.3B.
- End-to-end memory for SWAN is on par with SGD, about 50% reduction versus Adam; with per-layer training it reaches about 70% total reduction on top of about 100% optimizer-state reduction.
- With the NSDS whitening scheme, SWAN-ddag achieves raw throughput comparable to Adam (107808 vs 117872 tokens/s on 1B without model parallelism), giving about 2x higher effective throughput.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use SWAN‡ instead of Adam | LLaMA pretraining · C4, LLaMA 350M | validation perplexity: 19.24 -> 16.63 | Helped |
| Use SWAN‡ instead of Adam | LLaMA pretraining · C4, LLaMA 1.3B | validation perplexity: 16.44 -> 13.56 | Helped |
| Use NSDS whitening (SWAN‡) instead of vanilla whitening (SWAN†) | SWAN on LLaMA pretraining · C4, LLaMA 130M | validation perplexity: 22.83 -> 22.61 | Helped |
| Use NSDS whitening without model parallelization instead of vanilla NS whitening without model parallelization | SWAN throughput · LLaMA 1.3B | raw throughput (tokens/s): 58600 -> 107808 | Helped |
| Add GradNorm (ablation) | SWAN-0 · C4, LLaMA 130M | validation perplexity | Helped |
| Add GradWhitening (ablation) | SWAN-0 · C4, LLaMA 130M | validation perplexity | Helped |
| Remove GradNorm and increase the learning rate (ablation) | SWAN-0 · C4, LLaMA 130M | validation perplexity | Hurt |
| Disable learning-rate warm-up (ablation) | SWAN-0 · C4, LLaMA 130M | validation perplexity | Helped |
| Disable learning-rate warm-up (ablation) | Adam · C4, LLaMA 130M | validation perplexity | Hurt |
| Increase NSDS iterations from 2 to 5 or 10 (ablation) | SWAN‡ · LLaMA 130M test | test perplexity: 22.63 -> 22.62; 22.63 -> 22.61 | No clear effect |
| Use 1 NSDS iteration instead of 2 (ablation) | SWAN‡ · LLaMA 130M test | test perplexity: 22.63 -> - (LLM loss diverge) | Hurt |
| Use BF16 instead of FP32 for NSDS (ablation) | SWAN‡ · LLaMA 130M test | test perplexity: 22.63 -> 22.61 | No clear effect |
Surprising findings
- Removing GradNorm and compensating with larger learning rates (up to 1000x) does not close the gap to full SWAN, suggesting GradNorm's noise stabilization is not just an effective learning-rate increase.
- The diagonal-substitution whitening (NSDS) slows early convergence but gives stronger tail convergence and final performance than standard Newton-Schulz whitening.
Limitations
- The compositional effect of GradNorm and GradWhitening is not theoretically understood and is left open.
- Most baseline numbers are quoted from prior papers rather than re-run, and the Adam reproduction differs from the cited numbers.
- Theoretical analysis relies on a simplified transformer block and a Kronecker-structured Hessian assumption.
Experimental setup
Datasets: C4
Benchmarks: C4 validation perplexity
Metrics: validation perplexity speedup ratio in steps/tokens memory footprint (GB) throughput (tokens/s) condition number reduction
Baselines: Adam GaLore Apollo Apollo-mini LoRA ReLoRA Low-Rank Momentum+GradWhitening (Muon-like)
Builds on: Adam Muon GaLore decorrelated batch normalization
Closest papers by meaning
- Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation is Wasteful (NeurIPS 2025 · similarity 0.89)
- Revisiting Zeroth-Order Optimization for Memory-Efficient LLM Fine-Tuning: A Benchmark (ICML 2024 · similarity 0.89)
- Addax: Utilizing Zeroth-Order Gradients to Improve Memory Efficiency and Performance of SGD for Fine-Tuning Language Models (ICLR 2025 · similarity 0.89)
- Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training (ICLR 2024 · similarity 0.88)
- ZO-AdaMU Optimizer: Adapting Perturbation by the Momentum and Uncertainty in Zeroth-Order Optimization (AAAI 2024 · similarity 0.88)
- Gradient Weight-normalized Low-rank Projection for Efficient LLM Training (AAAI 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.