Research map / Architectures and efficiency / Pruning and efficient tuning
Adam-mini: Use Fewer Learning Rates To Gain More
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 Adam-mini, an optimizer that replaces per-parameter second-moment learning rates with one averaged learning rate per Hessian-structure-based parameter block, cutting optimizer memory by about 50% while matching AdamW.
Problem
Adam stores a per-parameter second-moment vector v, costing roughly 2x model size in optimizer state; it is unclear whether that per-parameter learning-rate resolution is needed, and prior memory-efficient alternatives (Adafactor, CAME, SM3) often underperform AdamW.
Approach
Using the near-block-diagonal Hessian structure of neural nets, partition parameters into dense Hessian sub-blocks (Transformer-specific rules: query/key by heads, value/attn.proj/mlp by output neurons, embed/output by tokens) and use one learning rate per block, set from the mean of Adam's v within that block.
Key results
- Adam-mini matches or slightly beats AdamW in pre-training from 39M to 1B Llama 2 models while using about 50% less optimizer memory (e.g., Llama 2-7B: 53.92 GB to 26.96 GB).
- On 2x A800-80GB, Adam-mini reaches 49.6% higher throughput than AdamW for Llama 2-7B pre-training, cutting wall-clock time by 33.1%.
- On MT-Bench with GPT-4 as judge, Adam-mini scores 5.68 vs 5.54 for AdamW on RLHF of Llama 2-7B.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace coordinate-wise second-moment rates with one blockwise rate using mean(v) | AdamW on Llama 2 scaling pre-training · Llama 2 models, 39M–1B | validation perplexity: 40.795 -> 40.407; 29.319 -> 29.014; 24.670 -> 24.192; 20.360 -> 20.172; 17.178 -> 17.035; 12.452 -> 12.372 | Helped |
| replace coordinate-wise second-moment rates with one blockwise rate using mean(v) | AdamW on Llama 2-7B SFT (LoRA) · MT-Bench | GPT-4 evaluation score: 4.23 -> 4.41 | Helped |
| replace coordinate-wise second-moment rates with one blockwise rate using mean(v) | AdamW on Llama 2-7B SFT · MT-Bench | GPT-4 evaluation score: 5.37 -> 5.40 | Helped |
| replace coordinate-wise second-moment rates with one blockwise rate using mean(v) | AdamW on Llama 2-7B RLHF · MT-Bench | GPT-4 evaluation score: 5.54 -> 5.68 | Helped |
| use Adam-mini instead of AdamW | Llama 2-7B pre-training on 2× A800-80GB · pre-training | throughput: +49.6% | Helped |
| use Adam-mini instead of AdamW | Llama 2-7B pre-training · 70B tokens | GPU hours: 5219.16 -> 3489.55 (↓ 33.1%) | Helped |
| replace mean(v) with 2-norm(v), 1-norm(v), max(v), or min(v) for block rates (ablation) | Adam-mini · Adam-mini ablation experiments | performance | Hurt |
| use the PyTorch default parameter partition instead of the Hessian-based partition (ablation) | Adam-mini on Llama 2-1B · pre-training | training stability | Hurt |
| partition Transformer parameters by the smallest dense Hessian sub-blocks (ablation) | Adam-mini on Llama 2-1B · pre-training | training stability | Helped |
Surprising findings
- The PyTorch default layer-wise partition causes training instability on 1B-scale LLMs, so the Hessian-based partition is necessary.
- Adafactor, CAME, and SM3 consistently underperform AdamW in these experiments despite saving similar memory.
Limitations
- Partition rules are derived for Transformers and tested on a limited set of architectures; may fail on new architectures.
- The mean-of-v learning rate is a cheap heuristic and may not be optimal; better block learning rates (e.g., via eigenvalues) are expensive.
- Theoretical justification (Hessian near-block-diagonal analysis) is informal and lower bounds on preconditioned condition number are not proven.
Experimental setup
Datasets: OpenWebText C4 ultrafeedback ImageNet CelebA OGB-arxiv
Benchmarks: MT-Bench ImageNet CelebA OGB-arxiv
Metrics: validation loss validation perplexity throughput (tokens/s) GPU hours memory (GB) MT-Bench score FID Inception Score validation accuracy
Baselines: AdamW Adafactor CAME SM3 Lion LAMB
Builds on: AdamW Adafactor SM3 LAMB ReMax
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)
- Adam Reduces a Unique Form of Sharpness: Theoretical Insights Near the Minimizer Manifold (NeurIPS 2025 · similarity 0.89)
- Understanding the Generalization of Stochastic Gradient Adam in Learning Neural Networks (NeurIPS 2025 · similarity 0.88)
- Mini-batch Coresets for Memory-efficient Language Model Training on Data Mixtures (ICLR 2025 · similarity 0.88)
- Addax: Utilizing Zeroth-Order Gradients to Improve Memory Efficiency and Performance of SGD for Fine-Tuning Language Models (ICLR 2025 · similarity 0.87)
- Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training (ICLR 2024 · 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.