Research map / Optimization / Optimizers
Lean and Mean Adaptive Optimization via Subset-Norm and Subspace-Momentum with Convergence Guarantees
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 Subset-Norm (SN) adaptive step sizes that share AdaGrad step sizes across parameter subsets to cut optimizer memory from O(d) to O(sqrt d), and Subspace-Momentum (SM) that keeps momentum in a low-rank subspace while doing SGD in the orthogonal complement, with high-probability convergence guarantees and LLaMA pretraining/fine-tuning experiments.
Problem
Reducing optimizer-state memory of adaptive optimizers (Adam, AdaGrad, RMSProp) for training large neural networks, ideally with convergence guarantees and without performance loss or heavy tuning.
Approach
Share adaptive step sizes across coordinate subsets (Subset-Norm, interpolating between AdaGrad-Norm and AdaGrad-Coordinate) and apply momentum only inside a low-dimensional subspace while taking plain SGD steps in its orthogonal complement (Subspace-Momentum), so the update stays full-rank and convergence can be proven.
Key results
- AdamSNSM reaches Adam's final validation perplexity of 16.00 on LLaMA 1B in roughly half the tokens (about 6.8B vs 13.1B) and reaches 14.05 at the end of training.
- AdamSNSM reduces Adam optimizer-state memory by more than 80% (e.g., 0.84 GB vs 4.99 GB for 1B model per Table of memory footprint).
- AdamSN improves over Adam on LLaMA 60M/130M/350M/1B pretraining (e.g., 29.75 vs 30.46 at 60M, 14.96 vs 16.00 at 1B).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace Adam's coordinate-wise second moment with Subset-Norm step sizes | Adam · LLaMA pretraining, 60M–1B | validation perplexity: 30.46 -> 29.75; 24.60 -> 22.90; 18.67 -> 17.49; 16.00 -> 14.96 | Helped |
| Add Subspace-Momentum to AdamSN | AdamSN · LLaMA pretraining, 60M–1B | validation perplexity: 29.75 -> 29.74; 22.90 -> 22.43; 17.49 -> 16.91; 14.96 -> 14.05 | Helped |
| Replace coordinate-wise adaptive step size with Subset-Norm | AdaGradm · LLaMA pretraining, 60M | validation perplexity: 30.40 -> 29.73 | Helped |
| Replace full momentum with Subspace-Momentum | AdamSN · LLaMA 7B SFT, UltraFeedback | last validation perplexity: 2.622 -> 2.584 | Helped |
| Use Subspace-Momentum rank 32 instead of rank 64 (ablation) | SNSM · LLaMA 7B SFT, UltraFeedback | last validation perplexity: 2.584 -> 2.580 | Helped |
| Use Subspace-Momentum with the top-k singular-vector projection rather than a random projection (ablation) | SNSM · LLaMA 60M pretraining, fixed subspace | validation perplexity: 40.15 -> 30.65 | Helped |
| Increase the projection update gap to 5% of total training steps (ablation) | AdamSNSM · LLaMA pretraining, 60M–1B | validation perplexity: 29.74 -> 29.84; 22.43 -> 22.71; 16.91 -> 18.43; 14.05 -> 15.28 | Mixed |
| Use a fixed subspace instead of updating the projection (ablation) | AdamSNSM · LLaMA pretraining, 60M–1B | validation perplexity: 29.74 -> 30.65; 22.43 -> 23.65; 16.91 -> 18.94; 14.05 -> 15.16 | Mixed |
| Use Subset-Norm step sizes with Subspace-Momentum rather than coordinate EMA plus full momentum (ablation) | Adam · LLaMA 60M pretraining | final evaluation perplexity: 30.46 -> 29.74 | Helped |
| Add momentum to adaptive optimizer configurations (ablation) | Adaptive step-size methods · LLaMA 60M pretraining | final evaluation perplexity: See Table: Step Sizes and Momentum Choices Full Ablations | Helped |
| Use cosine learning-rate decay instead of a constant learning rate (ablation) | Optimizer combinations · LLaMA 60M pretraining | final evaluation perplexity: See Table: Step Sizes and Momentum Choices Full Ablations | Helped |
| Add gradient clipping (ablation) | AdamSN · LLaMA 60M pretraining | validation perplexity: 30.06 -> 29.75 | Helped |
Surprising findings
- Restricting momentum to a low-rank subspace improves performance of subset-norm adaptive step sizes rather than degrading it.
- Vanilla AdaGrad with subset-norm step sizes closes the performance gap to Adam on LLaMA pretraining at larger scales, and AdaGrad with coordinate-wise steps performs well as model size grows.
Limitations
- Convergence of the combined SNSM method is not established; theory covers SN and SM separately under fixed projections.
- Subset-norm theory only covers equipartition subsets, and the subset-size heuristic is not optimal; SVD projection updates add computational overhead (slower wall-clock in SFT).
Experimental setup
Datasets: C4 UltraFeedback GLUE CIFAR10 CIFAR100 ImageNet
Benchmarks: C4 validation perplexity GLUE CIFAR10 CIFAR100 FID on DiT
Metrics: validation perplexity accuracy Matthews correlation FID peak memory optimizer state memory
Baselines: Adam AdaGrad AdaGradm RMSProp GaLore LoRA ReLoRA FLORA Adam-mini MicroAdam
Builds on: AdaGrad Adam GaLore Adam-mini SM3
Closest papers by meaning
- Enhancing Optimizer Stability: Momentum Adaptation of The NGN Step-size (NeurIPS 2025 · similarity 0.89)
- BROS: Bias-Corrected Randomized Subspaces for Memory-Efficient Single-Loop Bilevel Optimization (NeurIPS 2026 · similarity 0.88)
- SMMF: Square-Matricized Momentum Factorization for Memory-Efficient Optimization (AAAI 2025 · similarity 0.87)
- Regularized Adaptive Momentum Dual Averaging with an Efficient Inexact Subproblem Solver for Training Structured Neural Network (NeurIPS 2024 · similarity 0.86)
- A geometric framework for momentum-based optimizers for low-rank training (NeurIPS 2025 · similarity 0.86)
- Provable and Practical Online Learning Rate Adaptation with Hypergradient Descent (ICML 2025 · similarity 0.86)
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.