Adam-mini: Use Fewer Learning Rates To Gain More

2024/06/24 by Yushun Zhang, Congliang Chen, Zhang, Yushun +14 · 22 citations
Computer Science · Engineering · #Advanced Adaptive Filtering Techniques #Artificial Intelligence (cs.AI) #FOS: Computer and information sciences #Machine Learning (cs.LG) #Music and Audio Processing

paper · pdf · doi:10.48550/arxiv.2406.16793

openalex publication_date 2024/06/24 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

We propose Adam-mini, an optimizer that achieves on par or better performance than AdamW with 50% less memory footprint. Adam-mini reduces memory by cutting down the learning rate resources in Adam (i.e., 1/√(v)). By investigating the Hessian structure of neural nets, we find Adam's v might not function at its full potential as effectively as we expected. We find that ≥ 99.9% of these learning rates in v could be harmlessly removed if we (1) carefully partition the parameters into blocks following our new principle on Hessian structure; (2) assign a single but good learning rate to each parameter block. We then provide one simple way to find good learning rates and propose Adam-mini. Empirically, we verify that Adam-mini performs on par or better than AdamW on various language models sized from 39M to 13B for pre-training, supervised fine-tuning, and RLHF. The reduced memory footprint of Adam-mini also alleviates communication overheads among GPUs, thereby increasing throughput. For instance, Adam-mini achieves 49.6% higher throughput than AdamW when pre-training Llama 2-7B on 2× A800-80GB GPUs, which saves 33% wall-clock time for pre-training.

Cited by

• Adaptive Regularized Newton Method with Inexact Hessian
• FOAM: Blocked State Folding for Memory-Efficient LLM Training
• SpectralTrain: A Universal Framework for Hyperspectral Image Classification
• nanoTabPFN: A Lightweight and Educational Reimplementation of TabPFN
• FedAdamW: A Communication-Efficient Optimizer with Convergence and Generalization Guarantees for Federated Large Models
• FedMuon: Accelerating Federated Learning with Matrix Orthogonalization
• BSFA: Leveraging the Subspace Dichotomy to Accelerate Neural Network Training
• MISA: Memory-Efficient LLMs Optimization with Module-wise Importance Sampling
• Accelerated Distance-adaptive Methods for Hölder Smooth and Convex Optimization
• Noise-Adaptive Layerwise Learning Rates: Accelerating Geometry-Aware Optimization for Deep Neural Network Training
• Cautious Weight Decay
• Randomized Gradient Subspaces for Efficient Large Language Model Training
• AdaPM: a Partial Momentum Algorithm for LLM Training
• NorMuon: Making Muon more efficient and scalable
• Learning a Zeroth-Order Optimizer for Fine-Tuning LLMs
• DeMuon: A Decentralized Muon for Matrix Optimization over Graphs
• Conda: Column-Normalized Adam for Training Large Language Models Faster
• Partial Parameter Updates for Efficient Distributed Training
• FORGE: Fused On-Register Gradient Elimination for Memory-Efficient LLM Training
• Scalable Mean-Field Variational Inference via Preconditioned Primal-Dual Optimization
• Development of Deep Learning Optimizers: Approaches, Concepts, and Update Rules
• Matrix-Driven Identification and Reconstruction of LLM Weight Homology

Related

• When Can You Get Away with Low Memory Adam?
• APOLLO: SGD-like Memory, AdamW-level Performance
• Adam Accumulation to Reduce Memory Footprints of both Activations and Gradients for Large-scale DNN Training
• AdamZ: An Enhanced Optimisation Method for Neural Network Training
• Amos: An Adam-style Optimizer with Adaptive Weight Decay towards Model-Oriented Scale