Cosine annealing is a learning rate schedule used in training Machine learning models, particularly Deep learning models. It adjusts the learning rate during optimization by following a cosine curve, starting from an initial high value and decreasing smoothly to a minimum over a specified number of epochs or iterations. This approach is designed to balance the trade-off between convergence speed and stability, helping models settle into better minima of the loss function. The schedule is widely adopted in modern training pipelines for Neural network architectures, including Transformer (architecture)-based models such as Large language models, due to its simplicity and effectiveness.
In optimization, the learning rate determines the step size taken toward a minimum of the loss function. A constant learning rate can lead to slow convergence or oscillations, as too high a rate may overshoot minima and too low a rate may stall. Cosine annealing addresses this by providing a gradual decay that reduces the learning rate smoothly, allowing for finer adjustments as training progresses. Unlike step-based or exponential schedules, which drop the rate abruptly, cosine annealing offers a continuous, differentiable decay that often yields more stable training dynamics.
Mathematical Formulation
The cosine annealing schedule is defined by the formula:
\eta_t = \eta_{min} + \frac{1}{2}(\eta_{max} - \eta_{min}) \left(1 + \cos\left(\frac{t \pi}{T}\right)\right)
where \eta_t is the learning rate at iteration or epoch t, \eta_{max} is the initial learning rate, \eta_{min} is the minimum learning rate (often set to 0), and T is the total number of training steps in a cycle. The cosine function decreases from 1 to -1 as t goes from 0 to T, causing the learning rate to drop from \eta_{max} to \eta_{min} in a smooth, curved manner. This formula can be adapted for warm restarts, where the schedule is reset periodically, as described below.
The schedule is typically applied per epoch or per batch, depending on the implementation. For example, in frameworks like PyTorch and TensorFlow, built-in functions such as torch.optim.lr_scheduler.CosineAnnealingLR implement this schedule, allowing practitioners to specify the initial learning rate and the number of epochs. The smooth decay helps avoid abrupt changes that can destabilize training.
Historical Context
Cosine annealing was introduced in the context of Deep learning research as an alternative to traditional schedules like time-based, step-based, and exponential decay. These older methods, while effective, often require manual tuning of decay parameters and can be sensitive to the choice of hyperparameters. The cosine schedule was popularized by Ilya Loshchilov and Frank Hutter in their 2017 paper "SGDR: Stochastic Gradient Descent with Warm Restarts," which proposed using cosine annealing with periodic restarts to improve convergence. This work built on earlier research on learning rate schedules, such as the time-based decay formula \eta_{n+1} = \eta_0 / (1 + dn), where \eta_0 is the initial rate and d is a decay parameter.
The idea of varying the learning rate during training has roots in adaptive-control literature, where the learning rate is sometimes referred to as gain. Over the years, researchers have explored many schedules, including exponential decay \eta_n = \eta_0 e^{-dn} and step-based decay, but cosine annealing gained traction for its ability to achieve faster convergence and better generalization in many tasks.
Relationship to Other Schedules
Cosine annealing differs from common schedules in several ways. Time-based schedules reduce the learning rate inversely with the number of iterations, which can lead to very small rates late in training. Step-based schedules drop the rate by a factor at predefined intervals, which can cause sudden changes. Exponential schedules decay continuously but often too quickly. Cosine annealing, in contrast, provides a gradual decline that is neither too fast nor too slow, and it naturally allows for a "warm-up" phase if the initial learning rate is set appropriately.
Momentum, another key component in optimization, is often used alongside cosine annealing. Momentum helps accelerate learning when gradients point in a consistent direction and helps avoid local minima, similar to a ball rolling over small bumps. The combination of momentum and a cosine-decaying learning rate is common in training Transformer (architecture) models, where the schedule helps stabilize the optimization of large parameter spaces.
Warm Restarts and Variants
A notable variant of cosine annealing is the warm restart technique, where the schedule is reset after a cycle, allowing the learning rate to jump back to a higher value. This is implemented in the SGDR method, where the learning rate follows a cosine curve over a cycle of length T, then restarts with a new cycle, often with a longer duration. This approach can help the model escape local minima and explore different regions of the loss landscape, leading to improved performance. The warm restart variant is particularly useful in Deep learning tasks where the loss surface is non-convex.
Other variants include cosine annealing with a linear warm-up phase, where the learning rate increases from a small value to the initial maximum over the first few epochs, then decays cosinely. This is often used in training Large language models to avoid instability at the start of training.
Applications in Modern AI
Cosine annealing is extensively used in training state-of-the-art models across various domains. In Artificial intelligence research, it is a default choice for many Deep learning frameworks. For instance, OpenAI and Google DeepMind have employed cosine schedules in training large-scale models, including Transformer (architecture)-based architectures. The schedule is also integrated into popular libraries such as PyTorch and TensorFlow, making it accessible to practitioners.
In computer vision, cosine annealing has been shown to improve accuracy on benchmark datasets like CIFAR-10 and ImageNet. In natural language processing, it is used to train models like BERT and GPT, where the schedule helps manage the learning rate across millions of parameters. The technique is also applied in Generative AI models, including diffusion models and GANs, to stabilize training.
Advantages and Limitations
The primary advantage of cosine annealing is its smooth decay, which reduces the risk of oscillations and helps the model converge to a good minimum. It also requires fewer hyperparameters than some other schedules, as the main parameters are the initial and minimum learning rates and the cycle length. However, the schedule is not adaptive; it does not respond to the loss landscape or gradient information. In contrast, adaptive methods like Adam adjust the learning rate per parameter based on gradient statistics, which can be more robust in some settings. Cosine annealing is often used in conjunction with such optimizers, with the schedule applied to the global learning rate.
A limitation is that the total number of epochs must be known in advance to set the cycle length T. If training is extended, the schedule may need to be restarted or adjusted. Additionally, the choice of initial learning rate remains critical, as a too-high rate can cause divergence even with cosine decay.
Implementation in Practice
In practice, implementing cosine annealing is straightforward. In PyTorch, one can use torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max, eta_min=0), where T_max is the number of epochs. In TensorFlow, the tf.keras.optimizers.schedules.CosineDecay class provides similar functionality. These implementations automatically update the learning rate at each epoch or batch, allowing researchers to focus on model architecture and data.
For example, a typical training loop might set an initial learning rate of 0.1 and decay it to 0 over 100 epochs using cosine annealing. This schedule is often combined with momentum-based optimizers like SGD with momentum, which is common in computer vision tasks. For Large language models, a warm-up phase is often added, where the learning rate increases linearly for the first few thousand steps, then follows a cosine decay.
See Also
- learning-rate (if available)
- stochastic-gradient-descent (if available)
- optimization-algorithms (if available)
- Deep learning
References
- Loshchilov, Ilya; Hutter, Frank (2017). "SGDR: Stochastic Gradient Descent with Warm Restarts." International Conference on Learning Representations.
- Géron, Aurélien (2017). "Gradient Descent." Hands-On Machine Learning with Scikit-Learn and TensorFlow. O'Reilly. pp. 113–124. ISBN 978-1-4919-6229-9.
- Plagianakos, V. P.; Magoulas, G. D.; Vrahatis, M. N. (2001). "Learning Rate Adaptation in Stochastic Gradient Descent." Advances in Convex Analysis and Global Optimization. Kluwer. pp. 433–444. ISBN 0-7923-6942-4.