PaperlayerView exact revision on arXiv

Source-linked AI summary

Adam-mini: Use Fewer Learning Rates To Gain More

Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Diederik P. Kingma, Yinyu Ye, Zhi-Quan Luo, Ruoyu Sun

arXiv:2406.16793v7cs.LGcs.AI

TL;DR

Adam-mini addresses Adam’s substantial optimizer-state memory cost by replacing most per-parameter learning rates with one learning rate per Hessian-informed parameter block. It achieves on-par or better performance than AdamW with 50% less memory, including 49.6% higher throughput and 33.1% lower wall-clock time in one Llama 2-7B setting.

  • Problem

    Adam requires optimizer states consuming at least 2× model memory, creating a major burden for large-language-model training and motivating more memory-efficient optimization.

  • Method

    Adam-mini partitions parameters using dense Hessian sub-blocks and assigns each block a single learning rate based on the average of Adam’s v.

  • Results

    Adam-mini cuts at least 99.9% of Adam’s v and saves 50% memory while matching or exceeding AdamW; it also delivers 49.6% higher throughput and 33.1% lower wall-clock time in a Llama 2-7B pre-training setting.

  • Takeaways & Limitations

    The reduced memory footprint enables larger per-GPU batch sizes and eases GPU communication overhead, improving training efficiency.

  • Takeaways & Limitations

    Adam-mini’s learning-rate design uses blockwise averages of Adam’s v and may not be optimal, leaving room for stronger designs.

Abstract

from arXiv · show

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/\sqrt{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 $\geq$ 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\times$ A800-80GB GPUs, which saves 33% wall-clock time for pre-training.

1 INTRODUCTION

Adam-mini targets Adam’s substantial optimizer-state memory cost by replacing nearly all per-parameter learning rates with blockwise rates derived from Hessian structure. It retains comparable language-model performance while reducing memory and improving throughput.

  • 1 INTRODUCTION: Adam requires at least twice the model memory for momentum states, reaching about 56 GB for a 7B model before gradients are included.This memory burden can require CPU offload or optimizer-state sharding for training.
  • 1 INTRODUCTION: Adam-mini partitions parameters using dense Hessian sub-blocks and assigns one learning rate to each block, computed from the block’s averaged Adam v.This design uses fewer learning rates than Adam while retaining more structure than SGD.
  • 1 INTRODUCTION: Adam-mini cuts at least 99.9% of Adam’s learning-rate resources and saves 50% of Adam’s memory cost.The reduction targets the per-parameter rates represented by Adam’s v.
  • 1 INTRODUCTION: Adam-mini performs on par or better than AdamW across 39M–13B language models for pre-training, SFT, and RLHF.It also performs similarly to Adam on diffusion, vision, and graph neural-network training.
  • 1 INTRODUCTION: 49.6% higher throughput than AdamW on Llama 2-7B pre-training with 2× A800-80GB GPUs saves 33.1% wall-clock time.The efficiency is attributed to reduced memory enabling larger per-GPU batches and easing GPU communication.

2 METHOD

Adam-mini replaces Adam’s parameter-wise learning rates with one learning rate per carefully chosen parameter block, guided by neural-network Hessian structure. This design can preserve or improve optimization performance while sharply reducing optimizer-state memory and increasing throughput.

  • 2.1 MOTIVATIONS AND OBSERVATIONS: Neural-network Hessians are near-block-diagonal, with blocks corresponding to neurons, and this structure can persist throughout training.The paper uses this structure to motivate parameter blocks aligned with dense Hessian sub-blocks.
  • 2.1 MOTIVATIONS AND OBSERVATIONS: For block-diagonal quadratic problems, one well-chosen learning rate per dense sub-block can outperform Adam despite using tens or hundreds fewer learning rates.The result motivates replacing parameter-wise rates with blockwise rates rather than removing adaptivity entirely.
  • 2.1 MOTIVATIONS AND OBSERVATIONS: Adam’s diagonal preconditioner is effective when a Hessian sub-block is close to diagonal but not as effective when the sub-block is dense.This observation supports using Hessian-informed blocks instead of applying independent rates within dense sub-blocks.
  • 2.3 PRINCIPLE FOR THE PARTITION STRATEGY: The proposed partition follows the smallest dense Hessian sub-block, and the resulting Transformer strategy stabilizes training where PyTorch’s default partition can produce loss spikes.For Transformers, the paper further partitions parameters by heads, output neurons, or tokens according to their Hessian structure.
  • 2.4 SOME CHARACTERISTICS OF ADAM-MINI AND DISCUSSIONS: Adam-mini cuts more than 99.9% of Adam’s v and saves 50% of Adam’s memory by reducing learning rates from parameters to parameter blocks.The effective number of retained learning-rate values equals the number of blocks rather than the number of parameters.
  • 2.4 SOME CHARACTERISTICS OF ADAM-MINI AND DISCUSSIONS: 49.6% higher throughput than AdamW is achieved when pre-training Llama 2-7B on 2× A800-80GB GPUs, reducing wall-clock time by 33.1%.The memory reduction also enables larger per-GPU batch sizes and reduces inter-GPU communication.
  • 2.4 SOME CHARACTERISTICS OF ADAM-MINI AND DISCUSSIONS: Adam-mini’s blockwise average of Adam’s v is computationally cheap but may not be the best learning-rate design.The authors identify this design as an area with room for improvement.

3 EXPERIMENTS

Across language-model experiments, Adam-mini generally matches or exceeds AdamW while using 50% less memory, with benefits across pre-training, SFT, and RLHF. It also preserves similar training behavior and can improve throughput, while outperforming tested memory-efficient alternatives.

  • Pre-training: Adam-mini matches AdamW on GPT-2 and Llama pre-training, while the PyTorch-default partition becomes unstable on larger GPT-2 models.
  • Pre-training: Adam-mini’s loss curves and optimization trajectory closely resemble AdamW’s on the reported Transformer experiments.
  • Scaling laws: Adam-mini scales from 39M to 1B models with similar loss curves, lower final loss than AdamW in Figure 11, and substantially lower memory use.
  • SFT and RLHF: Adam-mini performs on par or better than AdamW for Llama 2-7B SFT and RLHF, including better alignment performance on MT-Bench.
  • Comparison with Adafactor: Adam-mini consistently outperforms Adafactor variants and achieves 40% higher throughput than Adafactor on Llama 2-1B.
  • Hyperparameter sensitivity: Adam-mini works with AdamW’s hyperparameters, whereas Adafactor has many tunable hyperparameters and remains difficult to tune in the reported comparisons.
  • Adam-mini performs on par with AdamW with 50% less memory across the tested experiments.

4 CONCLUDING REMARKS

The paper concludes that Adam-mini reduces Adam’s memory use by 50% while retaining comparable performance, but identifies room to improve its blockwise learning-rate design.

  • Adam-mini saves 50% of Adam’s memory while retaining comparable performance.
  • The current learning-rate design for dense Hessian sub-blocks is simple and cost-effective but may not be optimal.

BROADER IMPACT

The work presents a method for reducing LLM training costs, while noting that AI systems enabled by it could pose risks if used illegally.

  • Adam-mini is designed to reduce the training cost of large language models.
  • The paper notes a potential threat if resulting AI models are used for illegal purposes.

A RELATED WORKS

Adam-mini is positioned against memory-efficient and adaptive optimizers by using Hessian-informed blockwise learning rates rather than coordinate-wise rates, with reported gains in memory and throughput.

  • Unlike several lightweight optimizers that reduce Adam’s v through factorization or candidate-set rules, Adam-mini reportedly maintains AdamW-level effectiveness while those alternatives often degrade.The comparison includes Adafactor, CAME, and SM3.
  • The paper contributes a Hessian-based understanding of why fewer learning rates can suffice on near-block-diagonal optimization problems.This connects the optimizer design to a generic principle for problems with block-diagonal Hessians.
  • Adam-mini can combine with GaLore or Sophia for further memory reduction and throughput improvements.The paper reports about 40% additional memory reduction on GaLore and up to 50% savings in Sophia.
  • Adam-mini’s implementation is compatible with orthogonal memory-saving techniques such as checkpointing, offloading, sharding, quantization, and fused updates.
  • Adam-mini differs from LAMB because it replaces coordinate-wise 1/√v rates with blockwise rates, whereas LAMB retains coordinate-wise rates and adds layer-wise scaling.

C MORE DISCUSSIONS

The discussion motivates Adam-mini through near-block-diagonal Hessian structure and explains why blockwise averages of Adam’s v provide practical learning rates.

  • Near-block-diagonal Hessians arise because a shrinking loss-dependent term suppresses off-diagonal interactions, with the structure observed early and maintained during training.The cited analysis links blocks to output neurons and reports the effect after one training step.
  • The Hessian analysis is explicitly described as informal, with a more rigorous theoretical study left for future work.
  • Mean(v) also represents a block’s parameters because backpropagation gives entries in a weight row a shared error factor, making their gradients more similar.

D.1 MORE RESULTS FOR MOTIVATION

Motivation experiments extend the redundancy argument from quadratic problems to Transformers, showing that dense Hessian blocks need not use many coordinate-wise learning rates.

  • Adam’s preconditioner does not effectively reduce the condition number of dense Transformer Hessian blocks, leaving many learning rates redundant.The comparison is between κ(H) and κ(DAdamH).
  • Together, the experiments show that a single learning rate per block can work well and that the redundancy conclusion extends from random quadratics to Transformers.
  • Blockwise gradient descent outperforms AdamW in the 1-layer Transformer experiment with grid-searched learning rates for each block.
  • Mean(v) outperforms alternative quantities, including 1-norm(v), 2-norm(v), max(v), and min(v), in the Adam-mini ablation.
  • Adam-mini’s loss closely resembles AdamW during complete Llama 2-1B pre-training, while final validation perplexity is lower for models from 39M to 1B.

D.4 GPT-4 EVALUATION SCORE OF SFT AND RLHF

The GPT-4 evaluation compares Adam-mini and AdamW for SFT and RLHF on Llama 2-7B, alongside broader pre-training results and the optimizer’s memory and efficiency claims.

  • Adam-mini outperforms AdamW on MT-Bench chat evaluations for both supervised fine-tuning and reinforcement learning from human feedback.MT-Bench uses GPT-4 scores from 0 to 10, with higher scores better.
  • The complete Llama 2-1B loss curves show Adam-mini performing similarly to AdamW while using 50% less memory.
  • Adam-mini’s reported pre-training results include lower final validation perplexity than AdamW across models from 39M to 1B.

D.5 NON-LLM TASKS

Adam-mini matches or exceeds AdamW across the reported non-LLM and language-model comparisons while using substantially fewer learning-rate resources. The experiments also compare its behavior with Adafactor and Lion, including instability cases.

  • Adam-mini performs on par or better than AdamW on popular non-LLM tasks, including vision, diffusion, and graph neural-network training.
  • Adafactor-Zhai-version consistently underperforms Adam-mini across three hyperparameter setups, and tuning changes provide little improvement.
  • Lion underperforms Adam-mini on Llama 2-20M and GPT-2-125M, while also encountering loss spikes on GPT-2-125M.
  • On GPT-2-330M, AdamW encounters loss spikes whereas Adam-mini does not, although changing ϵ from 1e-8 to 1e-6 mitigates AdamW’s spike.
  • Replacing Adam steps with Adam-mini improves LoRA SFT performance after learning-rate grid search.

E.2 PRELIMINARY RESULTS IN (ZHANG ET AL., 2024)

Prior Hessian analyses find similar eigenvalue distributions across CNN parameter blocks but substantial differences across Transformer blocks. This motivates using different learning rates across Transformer blocks.

  • The observed heterogeneity suggests that Transformers need different learning rates for different blocks.
  • Transformer parameter blocks have different Hessian eigenvalue distributions, whereas CNN block distributions are similar.

F.1 TRAINING CONFIGURATIONS FOR SECTION 3

The experiments use standard or explicitly specified training protocols across GPT-2, Llama, SFT, RLHF, vision, diffusion, graph, and scaling-law studies. Configurations include model-specific learning rates, schedules, datasets, and evaluation setups.

  • GPT-2 experiments train 125M, 330M, and 1.5B models on OpenWebText with cosine-decay schedules, warm-up, weight decay, and specified Adam hyperparameters.
  • Llama experiments span 20M to 13B models on C4 using 1% warm-up and linear decay after warm-up, with learning rates specified for larger models.
  • SFT and RLHF use a Llama 2-7B pretrained model and UltraFeedback, splitting data between SFT and reward optimization with ReMax.
  • SFT uses validation-loss learning-rate selection, while ReMax selects peak learning rates using validation reward for AdamW and Adam-mini.
  • Vision, diffusion, and graph experiments use official implementations, standard configurations, and task-specific learning rates for AdamW and Adam-mini.
  • Trajectory comparisons save weights every 250 iterations and measure Euclidean distance to AdamW’s trajectory under matched seed and learning rate.

F.2 DETAILED SETUP FOR OTHER EXPERIMENTS

Additional experiments construct and analyze Hessian-related matrices and small neural networks, while other setups examine partition strategies and throughput. These procedures use exact or explicitly defined computational configurations.

  • The Figure 3 setup computes exact Hessians for a one-hidden-layer, 8-neuron MLP trained on CIFAR-100 with automatic differentiation.
  • Synthetic positive-definite Hessian blocks use prescribed eigenvalue ranges, and the blockwise learning-rate baseline uses the optimal constant rate 2/(L + µ).
  • Matrices H_b are generated as QΛQ^T, with Q formed through sampled rotations and Λ set to diag(κ, 1, ..., 1).
  • Figure 7 analyzes a one-layer Transformer with vocabulary size 8, embedding size 16, four heads, and 32 output neurons, using the same Hessian procedure as Figure 3.
  • Hessian rows and columns are reordered when codebases represent weight matrices with transposes or flattened vectors.
  • Throughput measurements use two A800-80GB GPUs without CPU offload and report values from the Weights & Biases summary log.
Loading 2406.16793v7…