Adam, short for Adaptive Moment Estimation, is an iterative optimization algorithm widely used for training neural networks and other Machine learning models. It was introduced by Diederik P. Kingma and Jimmy Ba in a 2014 paper titled "Adam: A Method for Stochastic Optimization." Adam combines the advantages of two other stochastic gradient descent extensions: adaptive learning rates (as in AdaGrad) and momentum (as in RMSProp). It computes individual adaptive learning rates for each parameter from estimates of the first and second moments of the gradients, making it well-suited for problems with large datasets and high-dimensional parameter spaces.
The algorithm is a variant of stochastic gradient descent (SGD), which itself is an iterative method for minimizing an objective function by replacing the true gradient with an estimate from a randomly selected subset of data. Adam has become a default optimizer in Deep learning frameworks and is used extensively in training Transformer (architecture) models, including large language models.
Algorithm
Adam maintains two moving averages per parameter: the first moment (mean) of gradients and the second moment (uncentered variance) of gradients. At each iteration \(t\), given the gradient \(g_t\) of the loss with respect to the parameters, the updates are computed as follows:
- Update biased first moment estimate: \(m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t\)
- Update biased second moment estimate: \(v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2\)
- Compute bias-corrected estimates: \(\hat{m}_t = m_t / (1 - \beta_1^t)\) and \(\hat{v}_t = v_t / (1 - \beta_2^t)\)
- Update parameters: \(\theta_t = \theta_{t-1} - \alpha \hat{m}_t / (\sqrt{\hat{v}_t} + \epsilon)\)
Here, \(\alpha\) is the learning rate (step size), \(\beta_1\) and \(\beta_2\) are exponential decay rates for the moment estimates (typically 0.9 and 0.999), and \(\epsilon\) is a small constant (e.g., \(10^{-8}\)) to prevent division by zero. The bias correction step is crucial in the early iterations when the moment estimates are initialized at zero, as it counteracts the bias toward zero.
Bias Correction
Because both \(m_t\) and \(v_t\) are initialized as zero vectors, the first few iterations produce estimates that are biased toward zero. Adam addresses this by dividing the moment estimates by \((1 - \beta_1^t)\) and \((1 - \beta_2^t)\), respectively. This correction becomes less significant as \(t\) grows, since the denominators approach 1. The bias correction is a key feature that distinguishes Adam from earlier adaptive methods and contributes to its stable convergence behavior.
Hyperparameters
Adam introduces several hyperparameters beyond the learning rate:
- Learning rate (\(\alpha\)): Controls the step size. Common default is 0.001.
- \(\beta_1\): Exponential decay rate for the first moment estimate. Default 0.9.
- \(\beta_2\): Exponential decay rate for the second moment estimate. Default 0.999.
- \(\epsilon\): Small constant for numerical stability. Default \(10^{-8}\).
In practice, the default values work well for many tasks, but tuning the learning rate is often necessary. Some implementations also support learning rate schedules, such as warmup and decay, which are common in training Transformer (architecture) models.
Variants and Extensions
Several variants of Adam have been proposed to address specific limitations:
- AdamW: Decouples weight decay from the gradient update, applying it directly to the parameters. This improves generalization and is now standard in many Deep learning libraries.
- Nadam: Combines Adam with Nesterov momentum, which can accelerate convergence.
- AMSGrad: Modifies the second moment estimate to ensure the learning rate does not increase, addressing convergence issues in some settings.
- Adamax: Uses the infinity norm for the second moment, making it more robust to large gradients.
- RAdam: Rectifies the variance of the adaptive learning rate, reducing the need for warmup.
These variants are used in specific contexts, but the original Adam remains widely used.
Applications
Adam is used across many domains of Artificial intelligence and Machine learning. It is the default optimizer in many frameworks, including TensorFlow and PyTorch. It has been applied to training neural networks for image classification, natural language processing, speech recognition, and reinforcement learning. In particular, Adam is the optimizer of choice for training Transformer (architecture)-based models, such as large language models developed by organizations like OpenAI, Anthropic, and Google DeepMind.
Adam's popularity stems from its robustness to hyperparameter settings and its ability to handle sparse gradients and noisy data. It is also memory-efficient, requiring only two additional variables per parameter.
Theoretical Properties
Adam does not always converge to a global minimum, especially for non-convex objectives, but it has been shown to converge to a critical point under certain conditions. The convergence analysis of Adam is more complex than that of plain SGD due to the adaptive learning rates. Some studies have shown that Adam can fail to converge in certain convex settings, which motivated the development of AMSGrad and other fixes.
Empirically, Adam often achieves faster convergence than SGD in the early stages of training, but may generalize slightly worse than SGD with momentum in some tasks. This has led to hybrid approaches, such as switching from Adam to SGD during training.
Comparison with Stochastic Gradient Descent
Stochastic gradient descent (SGD) updates parameters using a single sample or a mini-batch, with a fixed or decaying learning rate. Adam adapts the learning rate per parameter based on the history of gradients. This makes Adam less sensitive to the choice of learning rate and often requires less tuning. However, SGD with momentum can sometimes reach better final performance, especially with careful learning rate schedules.
Adam also differs from SGD in that it normalizes the gradient by the square root of the second moment, which can lead to more stable updates in the presence of large or small gradients.
Practical Considerations
When using Adam, it is common to set the learning rate to 0.001 and the betas to their defaults. For large-scale training, such as with large language models, learning rate schedules with warmup are often employed. Weight decay, as in AdamW, is recommended for regularization.
Memory usage is a consideration: Adam stores two additional values per parameter, which can be significant for models with billions of parameters. This has motivated research into memory-efficient optimizers, such as Adafactor, which approximates the second moment with lower-rank factors.
See Also
- stochastic-gradient-descent
- Deep learning
- Neural network
- Transformer (architecture)
References
- Kingma, D. P., & Ba, J. (2014). Adam: A Method for Stochastic Optimization. arXiv:1412.6980.
- Loshchilov, I., & Hutter, F. (2017). Decoupled Weight Decay Regularization. arXiv:1711.05101.
- Reddi, S. J., Kale, S., & Kumar, S. (2018). On the Convergence of Adam and Beyond. ICLR.