RAdam, short for Rectified Adam, is an optimization algorithm for training neural networks. It was introduced in 2019 by Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han in the paper "On the Variance of the Adaptive Learning Rate and Beyond." RAdam addresses a known issue with the Adam optimizer, where the adaptive learning rate can have high variance early in training, leading to poor convergence or suboptimal final performance. By rectifying this variance, RAdam aims to combine the benefits of Adam's fast convergence with the stability of stochastic gradient descent (SGD).
RAdam is widely used in deep learning, particularly in training transformers and large language models, where it has shown competitive performance with AdamW and other variants. It is implemented in major deep learning frameworks, including PyTorch and TensorFlow, and is a default choice in some training pipelines.
Background: The Adam Optimizer and Its Limitations
The Adam optimizer, introduced by Diederik Kingma and Jimmy Ba in 2014, is a popular adaptive learning rate method that computes individual learning rates for each parameter based on estimates of the first and second moments of the gradients. Adam has been highly successful in training deep neural networks, but it has known issues. One issue is that the adaptive learning rate can be too large in the early stages of training, especially when the second moment estimate is small. This can lead to large updates that cause the model to converge to sharp minima, which often generalize poorly. Additionally, Adam may require careful tuning of the learning rate and other hyperparameters.
Several variants have been proposed to address these issues, such as AdamW, which decouples weight decay, and AMSGrad, which uses a maximum of past squared gradients. However, these do not directly address the variance in the adaptive learning rate.
The Variance Problem in Adam
In Adam, the adaptive learning rate is computed as the ratio of the first moment estimate to the square root of the second moment estimate. Early in training, the second moment estimate is initialized to zero and is updated with a bias correction. However, the bias correction can still leave the estimate with high variance, especially when the decay rate (beta2) is close to 1. This variance causes the learning rate to fluctuate wildly, which can lead to unstable training and poor convergence.
RAdam's key insight is to quantify this variance and apply a rectification term that reduces the learning rate when the variance is high, and gradually increases it as the variance decreases. This is analogous to the rectified linear unit (ReLU) activation, which clips negative values to zero.
How RAdam Works
RAdam computes the adaptive learning rate similarly to Adam, but with an additional rectification term. The algorithm maintains the first moment (mean) and second moment (uncentered variance) of the gradients, denoted as m_t and v_t, respectively. It also tracks the timestep t and the decay rate beta2.
At each step, RAdam computes the bias-corrected estimates: m_t_hat = m_t / (1 - beta1^t) and v_t_hat = v_t / (1 - beta2^t). It then calculates a parameter rho_t = (1 - beta2^t) (2 / (1 - beta2) - 1) - t, which measures the effective length of the moving average. If rho_t is greater than a threshold (typically 4), RAdam uses the rectified term: the learning rate is scaled by sqrt((rho_t - 4) (rho_t - 2) rho_t / ((rho_t - 4) (rho_t - 2) rho_t - 4 (rho_t - 2) * (rho_t - 2))). If rho_t is less than or equal to 4, the update is simplified to use the first moment directly, similar to SGD with momentum.
This rectification ensures that the learning rate is not too large in the early stages, preventing the model from taking overly large steps that could lead to poor minima.
Advantages of RAdam
RAdam offers several advantages over Adam and other optimizers. First, it reduces the need for learning rate warmup, which is often required for Adam to avoid early instability. This simplifies hyperparameter tuning and can save training time. Second, RAdam has been shown to improve generalization performance on various tasks, including image classification and language modeling, compared to Adam. Third, it is computationally efficient, adding only a small overhead to the Adam update.
Empirical studies have shown that RAdam performs well on a range of models, including convolutional neural networks and transformers. It is particularly useful when training with small batch sizes or when the data is noisy, as the variance in gradients is higher in these scenarios.
RAdam in Practice
RAdam is implemented in popular deep learning libraries. In PyTorch, it is available as torch.optim.RAdam. In TensorFlow, it is available via the Keras API as tf.keras.optimizers.RAdam. It can be used as a drop-in replacement for Adam, with the same hyperparameters (learning rate, beta1, beta2, epsilon). The default values are typically learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8.
In practice, RAdam has been used in training large-scale models, including some large language models and transformers. For example, it has been adopted in certain training pipelines for models like GPT and BERT variants, where it has shown comparable or better performance than AdamW. However, AdamW remains a popular choice due to its decoupled weight decay, which can be beneficial for regularization. Some practitioners combine RAdam with weight decay or use it with a cosine learning rate schedule.
Comparison with Other Optimizers
RAdam is often compared with Adam, AdamW, and SGD with momentum. Compared to Adam, RAdam provides more stable training without the need for warmup, and often leads to better generalization. Compared to AdamW, RAdam does not decouple weight decay, but it can be combined with L2 regularization. In some benchmarks, RAdam outperforms AdamW on tasks like image classification, while in others AdamW is superior. The choice depends on the specific problem and hyperparameters.
RAdam also relates to other variance-reduction techniques, such as gradient clipping, which can be used in conjunction with RAdam to further stabilize training. Additionally, RAdam can be used with learning rate schedules like cosine annealing or step decay.
Extensions and Variants
Several extensions of RAdam have been proposed. For instance, RAdam has been combined with lookahead, a technique that maintains a set of slow weights, to create the Ranger optimizer, which has shown strong performance on various tasks. Another variant is AdaBelief, which modifies the second moment to use the deviation from the mean, but RAdam remains a baseline for comparison.
Research has also explored the theoretical properties of RAdam, providing convergence guarantees under certain conditions. The rectification term is derived from the analysis of the variance of the adaptive learning rate, and it has been shown to reduce the generalization gap.
Conclusion
RAdam is a robust optimizer that addresses the variance issue in Adam, leading to more stable and better generalizing training. It is easy to use, widely available, and has been proven effective in many deep learning applications. While it is not always the best choice, it is a valuable tool in the optimizer toolbox, especially for training large models where stability is critical.
For further reading, see Adam, SGD variants, and deep learning optimization techniques.