Source-linked AI summary
On the Variance of the Adaptive Learning Rate and Beyond
Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, Jiawei Han
TL;DR
The paper investigates why adaptive optimizers can converge poorly early in training, identifying excessive adaptive-learning-rate variance from limited samples as the central issue. It theoretically and empirically analyzes this mechanism, proposes RAdam to rectify the variance, and reports consistent improvement over vanilla Adam across tasks and architectures.
Problem
Adaptive learning rates have undesirably large variance early in training because only limited samples are available, while warmup lacks theoretical underpinnings and consistent guidance.
Method
The paper combines empirical and theoretical analysis with a variance-based rectification of Adam's adaptive learning rate, producing RAdam.
Results
RAdam brings consistent improvement over vanilla Adam across language modeling, image classification, and neural machine translation on different network architectures.
Takeaways & Limitations
The findings support warmup as a variance-reduction technique and show that RAdam provides a theoretically grounded alternative to heuristic warmup.
Takeaways & Limitations
The analysis assumes normally distributed gradients and a scaled inverse chi-square model for the second-moment estimate.
Abstract
from arXiv · showhide
The learning rate warmup heuristic achieves remarkable success in stabilizing training, accelerating convergence and improving generalization for adaptive stochastic optimization algorithms like RMSprop and Adam. Here, we study its mechanism in details. Pursuing the theory behind warmup, we identify a problem of the adaptive learning rate (i.e., it has problematically large variance in the early stage), suggest warmup works as a variance reduction technique, and provide both empirical and theoretical evidence to verify our hypothesis. We further propose RAdam, a new variant of Adam, by introducing a term to rectify the variance of the adaptive learning rate. Extensive experimental results on image classification, language modeling, and neural machine translation verify our intuition and demonstrate the effectiveness and robustness of our proposed method. All implementations are available at: https://github.com/LiyuanLucasLiu/RAdam.
1 INTRODUCTION
Adaptive optimizers can become unstable early in training because their adaptive learning rates have large variance, motivating warmup and the variance-rectifying optimizer RAdam.
- Removing warmup raises Transformer training loss from 3 to around 10 on De-En IWSLT’14.
- The paper attributes early convergence problems to undesirably large adaptive-learning-rate variance caused by limited training samples.
- Using smaller learning rates during the first few epochs reduces this variance and theoretically justifies warmup.
- RAdam explicitly rectifies adaptive-learning-rate variance based on theoretical derivations and is evaluated on language modeling, image classification, and neural machine translation.
- RAdam consistently improves over vanilla Adam across tasks and network architectures, supporting the general presence of the variance issue.
2 PRELIMINARIES AND MOTIVATIONS
Adaptive optimization combines momentum with an element-wise adaptive learning rate, while warmup uses smaller early learning rates to reduce instability observed in Transformer training.
- Generic adaptive methods: Generic adaptive methods compute momentum through φ(.) and adaptive learning rates through ψ(.), with Adam as one specification.
- Generic adaptive methods: Numerical implementations calculate ψ(.) with a small stabilizing ϵ, such as 1 × 10−8.
- Learning rate warmup: Linear warmup sets α_t = t α_0 when t < T_w, using smaller learning rates during the first steps rather than a constant or decreasing schedule.
- Learning rate warmup: Without warmup, the gradient distribution develops a mass center at relatively small absolute values within 10 updates.
- Learning rate warmup: Warmup reduces the impact of problematic early updates, and the analysis extends to adaptive-rate designs such as RMSprop and Nadam.
3 VARIANCE OF THE ADAPTIVE LEARNING RATE
The analysis attributes early optimization instability to high variance in the adaptive learning rate caused by limited samples. Controlled experiments and variance analysis show that increasing effective samples or reducing variance alleviates convergence problems, while excessive ϵ can introduce bias.
- 3 VARIANCE OF THE ADAPTIVE LEARNING RATE: For t = 1, the adaptive ratio has divergent variance, while scaling it by a small early learning rate reduces variance and alleviates problematic updates.This special-case result motivates interpreting warmup as variance reduction.
- 3.1 WARMUP AS VARIANCE REDUCTION: Adam-2k avoids vanilla Adam’s convergence problem by collecting 2,000 additional samples for adaptive-rate estimation while freezing momentum and parameters.The resulting gradient distribution is also not distorted, supporting insufficient early samples as the root cause.
- 3.1 WARMUP AS VARIANCE REDUCTION: A larger ϵ reduces adaptive learning-rate variance, and Adam-eps with ϵ = 10^-4 avoids vanilla Adam’s serious convergence problem.Its performance is worse than Adam-2k and Adam-warmup, which the authors conjecture results from bias introduced by large ϵ.
- 3.2 ANALYSIS OF ADAPTIVE LEARNING RATE VARIANCE: The analysis approximates Adam’s exponential moving average with a simple average because early exponential weights differ relatively little.The paper notes that the exponential moving average has larger variance than the simple average, so the approximation is used for analytical convenience.
- Early adaptive learning-rate variance is larger because few training samples are available, and it decreases as the effective sample count increases.The analysis approximates the adaptive-rate distribution and derives a variance that monotonically decreases with ρ, the effective degrees of freedom.
4 RECTIFIED ADAPTIVE LEARNING RATE
The paper estimates the effective sample size and adaptive-learning-rate variance, then uses these estimates to rectify Adam during early training. RAdam disables unreliable adaptation when variance is intractable and otherwise applies a variance correction without an extra warmup hyperparameter.
- 4.1 ESTIMATION OF ρ: The effective sample size ρ_t is estimated from the EMA decay rate and timestep by matching the EMA’s center of mass to an equivalent SMA.The resulting f(t, β2) is treated as ρ_t, while ρ∞ denotes the maximum length of the approximated SMA.
- 4.2 VARIANCE ESTIMATION AND RECTIFICATION: Adaptive-learning-rate variance is substantially larger early in training because few samples contribute to the estimate.The variance decays roughly as O(1/ρ_t); at ρ_t = 5 it is over 100 times larger than at ρ_t = 500.
- 4.2 VARIANCE ESTIMATION AND RECTIFICATION: The rectification term is designed so the variance of the adjusted adaptive learning rate matches the minimum late-stage variance C_var.The paper uses a first-order approximation because the analytic variance expression is not numerically stable.
- 4.2 VARIANCE ESTIMATION AND RECTIFICATION: RAdam deactivates adaptive learning when the approximated SMA length is at most 4, then applies variance rectification once the estimate becomes tractable.If β2 ≤ 0.6, ρ∞ ≤ 4 and RAdam degenerates to SGD with momentum.
- 4.3 IN COMPARISON WITH WARMUP AND OTHER STABILIZATION TECHNIQUES: RAdam outperforms Adam on One Billion Word, CIFAR10, and ImageNet, although rectification makes it slower during the first few epochs before faster convergence.The method is presented as orthogonal to other stabilization techniques and does not require the warmup horizon T_w.
5 EXPERIMENTS
Experiments across language modeling, image classification, and neural machine translation show that RAdam addresses early adaptive-learning-rate variance while improving robustness and matching heuristic warmup.
- Performance Comparison: RAdam outperforms Adam on One Billion Word, CIFAR10, and ImageNet, although it does not surpass SGD in test accuracy on CIFAR10 and ImageNet.On ImageNet, training accuracy is 69.57 for SGD, 69.12 for Adam, and 70.30 for RAdam.
- Robustness to Learning Rate Change: RAdam achieves consistent model performance across learning rates 0.1, 0.03, 0.01, and 0.003, whereas Adam and SGD are more learning-rate-sensitive.The CIFAR10 test-accuracy curves for RAdam highly overlap across this range.
- Comparing to Heuristic Warmup: RAdam achieves performance similar to heuristic warmup on neural machine translation while maintaining consistent adaptive-learning-rate variance.The comparison is reported for IWSLT’14 De-En, IWSLT’14 En-De, and WMT’16 En-De.
- Comparing to Heuristic Warmup: On CIFAR10, Adam with warmup reaches 91.29 test accuracy and RAdam reaches 91.38, but RAdam requires less hyperparameter tuning.At learning rate 0.1, 100 warmup steps yield 90.13 for Adam with warmup versus 91.06 for RAdam.
- Simulated Verification: The analytic variance expression and first-order approximation highly overlap across ρ = 5, …, 500, with a much smaller absolute difference.The result is presented as verification that the first-order approximation is very accurate.
- Simulated Verification: Simulations across six settings show larger early-stage variance for the adaptive learning rate and relatively consistent variance for the rectified rate.The simulations sample gradients from N(µ, 1), use β2 = 0.999, 5000 trajectories, and 6000 iterations.
6 CONCLUSION
The paper attributes warmup’s effectiveness to reducing the large early-stage variance of adaptive learning rates. It supports this hypothesis with empirical and theoretical evidence and proposes a new Adam variant.
- Limited early-stage samples make the adaptive learning rate undesirably high-variance.The paper links this variance to the early stage of model training.
- The large adaptive-learning-rate variance can cause convergence to suspicious or bad local optima.
- The paper investigates warmup’s underlying principle using both empirical and theoretical evidence.
A PROOF OF THEOREM 1
The proof derives variance expressions under a scaled inverse chi-square model and establishes the stated monotonic property through algebraic inequalities and special-function identities.
- A PROOF OF THEOREM 1: The proof models x as Scale-inv-X^2(ρ, τ^2) and uses its moments to derive Var[ψ(.)].The derivation includes E[x] and E[√x] through gamma-function expressions.
- A PROOF OF THEOREM 1: The monotonicity argument rewrites the target inequality into equivalent forms and applies bounds involving the digamma function.The proof also invokes the Legendre duplication formula.
B.1 LANGUAGE MODELING
The language-modeling experiments use a two-layer LSTM with adaptive softmax on the One Billion Word dataset, with specified embedding, cutoff, and preprocessing settings.
- B.1 LANGUAGE MODELING: The language-modeling setup uses two-layer LSTMs with 2048 hidden states and adaptive softmax on One Billion Word.The adaptive-softmax cutoffs are [4000, 40000, 200000].
- B.1 LANGUAGE MODELING: Inputs use randomly initialized 300-dimensional word embeddings, and tokens occurring three times or fewer are replaced with UNK.These are the stated embedding and preprocessing choices.
B.2 IMAGEINE CLASSIFICATION
The experiments use standard ResNet and Transformer implementations across image-classification and machine-translation settings, with task-specific architectures, batch sizes, and training schedules.
- Image classification: CIFAR-10 uses a 20-layer ResNet with 9 Basic Blocks and batch size 128.
- Image classification: ImageNet uses an 18-layer ResNet with 8 Basic Blocks and batch size 256.
- Image classification: CIFAR-10 training lasts 186 epochs, with learning-rate decays at epochs 81 and 122 by 0.1.
- Image classification: ImageNet training lasts 90 epochs, with learning-rate decays at epochs 31 and 61 by 0.1.
- Machine translation: Machine-translation experiments use fairseq Transformers with six-layer encoders and decoders, 512-dimensional embeddings, and dataset-specific attention and feedforward dimensions.IWSLT14 uses 4 heads and 1024 feedforward dimensions, while WMT14 uses 8 heads and 2048 feedforward dimensions.
C DOWNGRADING TO SGDM
RAdam degenerates to SGD with momentum during its first several updates, but these updates can distort gradients and make adaptive-variance instability damaging to early training reliability.
- C. DOWNGRADING TO SGDM: RAdam degenerates to SGD with momentum during the first several updates as a consequence of its mathematical derivation.
- C. DOWNGRADING TO SGDM: Gradient distributions can become distorted within 10 gradient updates during the early stage.The passage presents this distortion as an example from Figure 2.
- C. DOWNGRADING TO SGDM: Divergent adaptive learning-rate variance is described as potentially more damaging than converged variance because it implies greater instability.
- C. DOWNGRADING TO SGDM: A CIFAR-10 case study reports five-run average results in Table 3 while examining whether changing the first four updates affects model reliability.