The Imperative for Dynamic Learning Rates

In the optimization of deep neural networks, the learning rate stands as arguably the most critical hyperparameter, directly governing the magnitude of weight updates. If the rate is set too high, the optimization process risks instability; large weight updates can cause the training loss to oscillate erratically or diverge entirely, overshooting the optimal points in the loss landscape[4, 5, 6]. Conversely, if the learning rate is set too low, training convergence becomes prohibitively slow, and the model is more susceptible to becoming permanently trapped in suboptimal local minima, failing to reach a more globally optimal solution [4, 6].

This inherent tension necessitates a more sophisticated approach: dynamic learning rate scheduling. Learning rate schedules are algorithms or predefined frameworks that systematically adjust the learning rate during the training process. The core principle is to modulate the learning rate to suit different phases of training. Typically, a higher learning rate is employed at the outset to allow the model to make rapid progress across the loss landscape when its randomly initialized weights are far from optimal. As training progresses and the model’s parameters approach a more favorable region, the learning rate is gradually decreased. This reduction facilitates a more delicate, fine-tuning process, enabling the optimizer to converge carefully to a minimum without overshooting it.

The Cosine Annealing Schedule: Mechanics and Mathematics

Among the various dynamic scheduling techniques, cosine annealing has emerged as a particularly effective and widely adopted method. Its defining characteristic is the smooth, gradual adjustment of the learning rate according to a cosine curve.[5, 7, 11] Unlike step decay schedules, which introduce abrupt, discrete changes in the learning rate, the non-linear decay of cosine annealing provides a continuous and gentle transition.[5, 11] This smooth decay pattern can significantly enhance training stability, especially for complex models that might otherwise oscillate around a solution when subjected to sharp changes in the learning rate.[5] The conceptual underpinning of this technique draws an analogy from annealing in metallurgy, where a material is heated and then slowly cooled to relieve internal stresses and reach a minimum energy state, resulting in a stronger, more refined structure[11]. Similarly, cosine annealing “cools” the learning rate to guide the optimizer gracefully toward a low-error state.

The mathematical formulation of a single cosine annealing cycle is elegant and captures this behavior precisely. As implemented in leading deep learning frameworks like PyTorch and TensorFlow, the learning rate at a given time step, ηt​, is calculated as follows [12, 13, 14, 15]:

The components of this equation are defined as:

• ​: The learning rate at the current epoch or iteration t.
• ​: The initial and maximum learning rate, which serves as the upper bound for the schedule. This is typically set to the initial learning rate configured in the optimizer[12, 13]
• ​: The minimum learning rate, which acts as a lower bound. The schedule ensures the learning rate does not decay below this value, allowing for continued, albeit slow, learning even at the end of training.[12, 13] In some frameworks like Keras, this is controlled by an alpha parameter, which specifies ηmin​ as a fraction of the initial learning rate.[16, 17]
• ​: The maximum number of iterations or epochs in the decay cycle. This parameter defines the period over which the learning rate will travel from to ​.[5, 13] It corresponds to a half-period of the cosine function.
• ​: The number of epochs or iterations that have elapsed since the beginning of the current decay cycle.[5, 12]

Visually, this formula traces a distinctive curve. At the start of the cycle (Tcur​=0), the cosine term is cos(0)=1, which sets ηt​=ηmax​. As training progresses towards the end of the cycle (Tcur​→Tmax​), the argument of the cosine function approaches π, making the cosine term approach cos(π)=−1. This results in ηt​ approaching ηmin​. The rate of decay is not linear; it is slowest at the beginning and end of the cycle and fastest in the middle, creating a concave decay curve that contrasts sharply with the sudden drops of a StepLR schedule or the steady decay of an ExponentialLR schedule.[7, 9, 10]

SGDR and the Concept of Warm Restarts

The application of cosine annealing in modern deep learning was popularized by the seminal paper “SGDR: Stochastic Gradient Descent with Warm Restarts” by Loshchilov & Hutter[15, 18, 19, 20]. It gave rise to two distinct but related learning rate schedulers, which are now standard components in libraries like PyTorch.

First, CosineAnnealingLR implements only the cosine decay portion of the SGDR strategy[5, 13]. It executes a single, continuous decay of the learning rate from ηmax​ to ηmin​ over a specified duration of T_max epochs or iterations. This scheduler is best suited for scenarios where the primary goal is a smooth, uninterrupted convergence over the entire training run, such as during the fine-tuning of a pre-trained model on a new task where extensive exploration is not the main objective.[5]

Second, CosineAnnealingWarmRestarts implements the full SGDR paradigm[5, 21, 22]. It employs the same cosine decay mechanism but introduces periodic “warm restarts.” At the end of each decay cycle, the learning rate is abruptly reset to its initial maximum value, ηmax​, and a new decay cycle begins.[5, 14, 21] These restarts are considered “warm” because the model’s weights are not re-initialized; the optimizer continues from its current state, preserving the learned knowledge.[14] The theoretical motivation for these restarts is to help the optimizer escape from sharp, narrow local minima in the loss landscape. By periodically increasing the learning rate, the optimizer is given a renewed burst of momentum, enabling it to “jump” out of a suboptimal basin and explore other, potentially wider and more generalizable regions of the solution space[7, 23, 24].

The CosineAnnealingWarmRestarts scheduler introduces two additional key hyperparameters to manage these cycles:

• T_0: An integer specifying the number of iterations or epochs for the first restart cycle.[22, 25]
• T_mult: An integer factor by which the cycle length, Ti​, is multiplied after each restart. For example, if T_mult is 2, each subsequent cycle will be twice as long as the previous one.[22, 25] This allows for broad, rapid exploration in the early cycles and progressively longer, more focused fine-tuning periods in later cycles.

When a practitioner selects CosineAnnealingLR for a single, long decay, they are operating under the assumption that the optimization journey is primarily one of convergence. The model, after some initial training, is presumed to have found a promising “basin of attraction” in the loss landscape. The main challenge is no longer exploration but exploitation: carefully navigating to the bottom of this basin to find the best possible solution within it. The smooth, decelerating nature of the cosine curve is ideally suited for this final “zeroing-in” phase, as it reduces the step size precisely when the optimizer is close to a minimum, thereby preventing overshooting[5]. This strategy is often highly effective for transfer learning, where a pre-trained model already possesses a strong feature representation and only needs to be fine-tuned.

In contrast, choosing CosineAnnealingWarmRestarts implies a belief that the loss landscape is highly complex and multimodal, riddled with numerous local minima that could trap the optimizer.[15] Here, the warm restarts are not just a feature but a core part of the strategy. They are a deliberate mechanism to combat getting stuck. By periodically resetting the learning rate to a high value, the optimizer is endowed with renewed “energy” to surmount the walls of a given basin and explore entirely different regions of the parameter space[23, 24]. This exploratory power can be crucial for training large, complex models from scratch on novel datasets, where finding a wide, flat minimum (which is often associated with better generalization) is paramount.

This distinction reveals a critical trade-off. While the exploratory power of warm restarts is valuable, it is not without cost. Recent research, particularly in the context of continual learning, has suggested that the “re-warming” phase of repeated schedules can induce catastrophic forgetting, where the model loses previously acquired knowledge upon being exposed to new data or a reset learning rate[26]. This introduces a new dimension to the decision: the potential for enhanced exploration must be weighed against the risk of destabilizing and erasing learned representations.

The Role of Learning Rate Warmup: Stabilizing the Ascent

Warmup as a Stabilization Mechanism

Applying a high learning rate during the initial training phase is a recipe for instability. The large gradients, when multiplied by a large learning rate, result in massive, erratic updates to the model’s weights. Learning rate warmup is a simple yet profoundly effective technique designed to navigate this treacherous initial phase of training. The core idea is to begin the optimization process with a very small learning rate, often a value close to zero, and then gradually increase it over a specified number of initial training steps. This ramp-up is typically linear, though other functions can be used[28, 31, 32]. After this “warmup period,” which lasts for a pre-defined number of warmup_steps, the learning rate reaches its target base value (the η_max of the subsequent decay schedule), and the main training phase begins[33, 34].

The benefits of this approach are multifaceted and address the initial instability problem directly. Work by Kalra & Barkeshli (2024) demonstrates that the primary benefit of warmup is its ability to force the network into more well-conditioned, or “flatter,” areas of the loss landscape.[30, 31, 35] The “sharpness” of the loss landscape at a given point can be characterized by the largest eigenvalue of the Hessian matrix. A high sharpness value indicates a steep, narrow valley, where large learning steps can easily cause the optimizer to overshoot and become unstable. The research suggests that warmup acts as a process of sharpness annealing: the small, careful steps taken during the warmup phase gradually guide the model parameters towards regions of lower sharpness. By the time the warmup period ends and the learning rate reaches its high peak value, the model is in a much flatter region of the loss landscape and can tolerate these large steps without diverging.[30]

This work also introduces the “loss catapult” mechanism, a counterintuitive phenomenon where a carefully chosen learning rate can cause a temporary increase in training loss, which paradoxically leads to a significant decrease in sharpness[31, 35]. This insight opens up possibilities for more principled and efficient warmup strategies, where one could potentially search for this critical learning rate to accelerate the sharpness annealing process and reduce the required number of warmup steps.

The performance benefits are well-documented in ablation studies. One study reported that applying a learning rate warmup allowed a model to achieve 95% of its final accuracy using 28% fewer training epochs compared to a fixed learning rate, demonstrating a significant improvement in convergence speed[36]. Another experiment on a channel state information (CSI) feedback task found that a cosine annealing schedule with warmup dramatically outperformed a constant learning rate, even when using the adaptive Adam optimizer[37]. The impact is particularly stark for very large models; practitioners have noted that while smaller “base” models might see little benefit, large-scale models often fail to train altogether without a warmup phase, with the loss collapsing at the very beginning of the run[27].

A deeper analysis reveals that learning rate warmup is more than a simple stabilization trick; it functions as a sophisticated pre-conditioning phase for the entire optimization process. It fundamentally alters the trajectory of the optimizer in the initial stages of training, preparing the model for the more aggressive exploration and convergence phases that follow. However, the very fact that this technique is so critical suggests that it may be a highly effective compensation for underlying limitations in our current generation of optimization algorithms and initialization schemes.

This perspective reframes warmup from a fundamental principle of deep learning to a powerful and necessary heuristic patch. The initial instability that warmup is designed to solve points directly to specific weaknesses in popular optimizers. For instance, research has shown that the bias-correction mechanism for the first-moment estimate in the Adam optimizer can cause an artificial spike in the magnitude of the initial weight updates[38]. Furthermore, at the start of training, the gradients calculated from different mini-batches are often highly correlated, which limits the effective batch size and can contribute to instability [29, 38].

This understanding has profound implications for the future of deep learning optimization. It suggests that the reliance on warmup is not an immutable law but a symptom of the current state of our tools. The success of warmup highlights a clear frontier for research: the development of next-generation optimizers and initialization methods that are inherently stable from the very first step. Innovations like GI-Adam, which proposes a different initialization for the second moment in Adam to mimic the effects of warmup, or LionAR, which aims to control update sizes more directly, are steps in this direction[30, 31, 39]. For the practitioner today, this means that while mastering warmup is an essential skill for achieving state-of-the-art results, it is equally important to remain aware of emerging optimization techniques that may one day render it obsolete by providing a simpler, more robust, and less hyperparameter-dependent training process out-of-the-box.

The Combined Strategy: Cosine Annealing with Linear Warmup

Play with the Interactive Version of the Graph here.

The combination of a linear warmup phase followed by a cosine annealing decay phase has become a de facto standard for training high-performance deep learning models, particularly in domains like natural language processing and computer vision[4, 34, 40, 41]. The learning rate trajectory unfolds in two distinct phases:

• Phase 1: Linear Warmup. The schedule begins with the learning rate at a very low value, denoted as warmup_start_lr, which is often set to 0[33]. Over a predefined number of warmup_steps, the learning rate increases linearly with each training iteration. At the conclusion of this phase, the learning rate reaches its peak value, η_max (also referred to as base_lr or target_lr), which will serve as the starting point for the subsequent decay phase[33, 34]. Mathematically, for a step t within the warmup period (t≤twarmup​), the learning rate ηt​ can be expressed as:ηt​=ηstart​+(ηmax​−ηstart​)×twarmup​t​
• Phase 2: Cosine Annealing Decay. Immediately upon completion of the warmup, the schedule transitions seamlessly into a cosine decay. The learning rate then begins its smooth, non-linear descent from the peak value η_max down towards a specified minimum value η_min. This decay occurs over the remaining duration of the training run, typically defined by a parameter like T_max or decay_steps[34, 40]. The shape of this decay follows the half-period cosine curve described previously, ensuring a gradual reduction that is ideal for fine-tuning as the model converges.

A visualization of this combined schedule reveals a characteristic shape: a straight, upward-sloping line representing the linear ramp-up, followed immediately by the familiar concave curve of the cosine decay, ending at or near zero at the conclusion of training.[14, 37, 40, 42] This two-part structure provides a complete, end-to-end strategy for managing the learning rate throughout the entire training process.

Performance Comparison: With vs. Without Warmup

To appreciate the value of the combined schedule, it is instructive to compare its performance against a standard cosine annealing schedule that lacks a warmup phase.

• Without Warmup: A standalone CosineAnnealingLR schedule commences directly at its peak learning rate, η_max. This presents a significant challenge for the practitioner. If η_max is set to a high, aggressive value to encourage rapid exploration, there is a substantial risk that the model will diverge in the chaotic initial steps of training. To avoid this, the practitioner is often forced to select a conservatively low, suboptimal peak learning rate, thereby limiting the model’s potential for both speed of convergence and final performance.[27]
• With Warmup: The combined schedule fundamentally resolves this dilemma. The initial warmup phase acts as a stabilization period, gently guiding the model out of its random initial state. By the time the warmup is complete, the model’s weights have settled into a more stable configuration, and the optimizer is prepared to handle much larger updates without diverging. This allows the practitioner to set a significantly higher, more optimal peak learning rate for the cosine decay phase.[30, 31] The result is a training process that is both more stable at the start and more effective in its exploration and convergence, frequently leading to superior final model accuracy.

Empirical evidence consistently validates the superiority of the combined approach. Ablation studies conducted on the training of Generative Pre-trained Transformer (GPT) models show that even a short warmup period creates a significant and lasting performance gap compared to runs without warmup, even when using the same peak learning rate. The no-warmup runs, while showing faster initial loss reduction, ultimately plateau at a worse performance level, indicating that the initial instability causes irreparable harm to the optimization process.[38] Similarly, in a study on a CSI feedback task, the combination of “cosine annealing lr with warm up” was demonstrated to be a critical component for achieving high performance, substantially outperforming a constant learning rate schedule.[37]

Advantages and Disadvantages

The widespread adoption of the linear warmup with cosine annealing schedule is due to a compelling set of advantages, though it is not without its own complexities.

Advantages:

• Enhanced Stability and Performance: The primary benefit is the ability to achieve the stability of a low initial learning rate while simultaneously reaping the rewards of a high peak learning rate, namely faster convergence and more thorough exploration of the loss landscape. This combination often leads to better final model performance than either strategy could achieve in isolation[43, 44].
• Increased Robustness to Hyperparameter Choice: The warmup phase makes the entire training process more robust. Because the model is stabilized before the peak learning rate is applied, a wider range of high learning rates becomes viable options, simplifying the tuning process and reducing the risk of a poor choice leading to complete training failure[30, 31].
• A Proven State-of-the-Art Standard: This schedule has become the go-to strategy for training large, complex, and state-of-the-art models, especially Transformers. Its effectiveness has been demonstrated across a vast array of tasks and datasets, making it a reliable choice for challenging deep learning problems[4, 45].

Disadvantages:

• Increased Hyperparameter Complexity: The primary drawback is the introduction of additional hyperparameters that require careful tuning. The practitioner must now determine not only the peak and minimum learning rates but also the warmup_steps (or warmup_ratio) and the warmup_start_lr. A misconfiguration of these parameters can easily negate the potential benefits[29, 46].
• Potential for Wasted Computation: The warmup phase, while crucial, does represent a period of intentionally slowed-down learning. If the warmup duration is set to be excessively long, it can waste valuable computational resources and delay the point at which the model begins to make significant progress with a higher learning rate[31, 35]. Research into more efficient warmup methods, such as those leveraging the “loss catapult” phenomenon, aims to minimize this inefficiency[31].
• Heuristic Nature and Masking of Problems: As discussed previously, the profound need for warmup may be symptomatic of other underlying issues in the training setup, such as a suboptimal optimizer choice or a poor weight initialization strategy. Relying on the warmup-and-decay schedule as a black-box solution without understanding why it is necessary can mask these deeper problems, preventing a more fundamental resolution[39, 47].

Framework-Specific Implementation and Best Practices

Implementation in PyTorch

Implementing a cosine annealing schedule with warmup in native PyTorch requires a compositional approach, as the standard torch.optim.lr_scheduler.CosineAnnealingLR class does not include built-in parameters for a warmup phase[13, 42]. Practitioners must therefore combine multiple scheduler components to construct the desired learning rate trajectory.

Solution 1: Chaining with SequentialLR

The canonical and most explicit method in PyTorch is to use torch.optim.lr_scheduler.SequentialLR. This scheduler takes a list of other schedulers and a list of milestones, activating them in sequence. To implement the warmup-and-decay schedule, one would create:

1. A torch.optim.lr_scheduler.LinearLR instance for the warmup phase. This scheduler is configured with a start_factor (e.g., 1e-4) and total_iters corresponding to the number of warmup steps.
2. A torch.optim.lr_scheduler.CosineAnnealingLR instance for the decay phase, configured with its T_max and eta_min parameters.
3. These two schedulers are then passed to SequentialLR, along with a milestones list containing a single integer: the number of warmup steps. This tells SequentialLR to use the LinearLR scheduler up to that step, and then switch to the CosineAnnealingLR scheduler for the remainder of training.[48] While this method is somewhat verbose, it offers clear, granular control over each phase of the schedule.

Solution 2: Third-Party Libraries

Several popular third-party libraries offer more streamlined, integrated solutions, abstracting away the manual composition.

• pytorch-warmup: This library provides a warmup scheduler object that acts as a wrapper around a primary scheduler (like CosineAnnealingLR). The implementation involves using a Python context manager, with warmup_scheduler.dampening():, within the training loop. This manager applies the appropriate warmup factor to the learning rate before the main scheduler’s step() method is called.[48, 49, 50] This approach can lead to cleaner-looking training loops, but requires careful logic to ensure the main scheduler’s step() is only invoked after the warmup period is complete, to prevent the decay from starting prematurely.[48]
• pytorch-accelerated: This library offers a stateful CosineLrScheduler class that comes with built-in support for warmup. It can be initialized directly with parameters like num_warmup_epochs and warmup_starting_lr, providing an integrated experience that is more akin to the native implementation in TensorFlow/Keras.[51]
• katsura-jp/pytorch-cosine-annealing-with-warmup: This is another dedicated package that provides a single CosineAnnealingWarmupRestarts class, which combines warmup, cosine decay, and optional restarts into one convenient object.[52]

Regardless of the chosen method, it is critical to adhere to PyTorch’s conventions for scheduler usage. The scheduler.step() method should always be called after the optimizer.step() method within the training loop[53, 54, 55]. Furthermore, the frequency of the step() call must match the units used to define the scheduler’s parameters (e.g., per-iteration for step-based schedulers, or per-epoch for epoch-based ones)[56].

Implementation in TensorFlow & Keras: An Integrated Approach

In contrast to PyTorch’s compositional philosophy, TensorFlow and its high-level API, Keras, provide a more direct and integrated solution for this specific schedule.

The Native Solution: tf.keras.optimizers.schedules.CosineDecay

The most straightforward and recommended method is to use the tf.keras.optimizers.schedules.CosineDecay class. This learning rate schedule object has native support for a warmup phase, which can be enabled by providing values for the warmup_target and warmup_steps arguments during initialization.[16, 57, 58]

The mechanics are simple: an instance of the CosineDecay schedule is created with all the necessary parameters (for both warmup and decay). This schedule object is then passed directly as the learning_rate argument to a Keras optimizer, such as tf.keras.optimizers.Adam. The optimizer will then internally call the schedule at each training step to get the correct learning rate value[16, 17]. This approach encapsulates the entire logic within the schedule object, leading to minimal boilerplate code in the training script.

Alternative: Custom Callbacks

For scenarios requiring more complex or non-standard scheduling behavior that cannot be captured by the built-in classes, TensorFlow/Keras allows for the implementation of custom callbacks. A user can create a class that inherits from tf.keras.callbacks.Callback and implement the on_batch_end or on_epoch_end methods. Within these methods, one can calculate the desired learning rate based on the current batch or epoch number and then manually set the optimizer’s learning rate using the backend function tf.keras.backend.set_value(self.model.optimizer.lr, new_lr)[59, 60, 61]. This method offers maximum flexibility but comes at the cost of increased implementation complexity and verbosity compared to the native schedule object.

Clear code examples for both the native CosineDecay schedule and a custom callback implementation would highlight the elegance and simplicity of the former for this common use case[1, 59, 62]

Implementation in the Hugging Face Trainer API

The Hugging Face Trainer API provides the highest level of abstraction, designed to simplify the training process for Transformer models and minimize boilerplate code. Implementing a cosine annealing schedule with warmup is achieved by setting a few string-based arguments in the TrainingArguments class.

The key arguments are:

• lr_scheduler_type: This string argument specifies the desired schedule. Setting it to 'cosine' will automatically select the library’s internal implementation of a cosine annealing schedule with a linear warmup phase.[63, 64]
• learning_rate: This float value sets the peak learning rate (η_max) that will be reached at the end of the warmup.
• warmup_steps or warmup_ratio: These arguments control the duration of the warmup phase. warmup_steps defines an absolute number of steps, while warmup_ratio defines the duration as a fraction of the total training steps.[64]

A critical nuance arises in distributed training environments. When training on multiple GPUs using a library like Hugging Face Accelerate, the learning rate scheduler may be stepped for each process on every “global” optimizer step. This means that over a single step of the training loop, the scheduler’s internal counter might advance by the number of GPUs. If the warmup_steps and total training steps are not defined with this behavior in mind, the learning rate can change much more rapidly than intended, which is a common source of confusion and unexpected results for practitioners.[65]

Framework Implementation Comparison

Feature	PyTorch (Native/Chained)	TensorFlow/Keras (Native)	Hugging Face Trainer
Base Scheduler	torch.optim.lr_scheduler.LinearLR + torch.optim.lr_scheduler.CosineAnnealingLR	tf.keras.optimizers.schedules.CosineDecay	get_cosine_schedule_with_warmup (internal)
Warmup Implementation	Chained via SequentialLR or handled by a third-party wrapper library.	Integrated via warmup_target and warmup_steps arguments.	Integrated via warmup_steps or warmup_ratio argument.
Key Parameters	start_factor, total_iters (for LinearLR); T_max, eta_min (for CosineAnnealingLR); milestones (for SequentialLR).	initial_learning_rate, decay_steps, alpha, warmup_target, warmup_steps.	learning_rate, lr_scheduler_type, warmup_steps.
Ease of Use	More verbose, requires understanding of scheduler composition. Offers high flexibility.	Very straightforward for this specific schedule. Less boilerplate.	Highest level of abstraction, minimal code required.
Snippet Refs	[48, 49, 51]	[16, 17, 58]	[63, 64]

This comparative summary highlights the fundamental design philosophies of the frameworks. PyTorch prioritizes composability and flexibility, giving the user fine-grained control at the cost of verbosity. TensorFlow/Keras offers a more integrated, user-friendly experience for common patterns. The Hugging Face Trainer takes this a step further, providing a high-level, declarative API that is extremely convenient but offers less direct control over the underlying mechanics.

A Practical Guide to Hyperparameter Tuning and Best Practices

The Core Hyperparameters: A Detailed Guide

Effectively leveraging the cosine annealing with warmup schedule requires a solid understanding of its key hyperparameters and their impact on the learning rate trajectory.

• Peak Learning Rate (η_max, learning_rate): This is the single most critical hyperparameter. It defines the maximum learning rate achieved at the end of the warmup phase and serves as the starting point for the cosine decay. This value dictates the optimizer’s maximum step size and, consequently, its exploratory potential. An improperly set peak learning rate can lead to either slow convergence or training instability. It is highly task-dependent and should be tuned with care, often using an empirical method like a learning rate range test.[2, 24]
• Warmup Duration (warmup_steps, warmup_epochs, warmup_ratio): This parameter determines the length of the initial linear ramp-up phase. A common and effective heuristic is to allocate 5-10% of the total training steps to warmup.[2] The duration represents a trade-off: if it is too short, it may not be sufficient to stabilize the model, negating the benefit of warmup. If it is too long, it can unnecessarily slow down the initial phase of training, wasting computational cycles during which more rapid progress could have been made.[31]
• Decay Duration (T_max, decay_steps): This parameter specifies the length of the cosine decay phase. In the most common configuration, this is set to be the total number of training steps minus the number of warmup steps. This ensures that the learning rate completes its decay from η_max to η_min precisely at the end of the training run, providing a complete, end-to-end schedule.[17, 59]
• Minimum Learning Rate (η_min, alpha): This value sets the floor for the learning rate. At the end of the decay phase, the learning rate will not drop below this value. Setting it to a small, non-zero number (e.g., 1/10th or 1/100th of the peak learning rate) ensures that the model can continue to make very small weight updates and fine-tune itself even in the final stages of training. In many cases, it is simply set to 0.[8, 13, 16]

A Strategic Approach to Tuning

Tuning the hyperparameters of this composite schedule should be approached systematically rather than through random guesswork. A structured, iterative process is recommended.

1. Establish a Performance Baseline: Before implementing a complex schedule, it is valuable to train the model with a simpler setup, such as a well-tuned constant learning rate or a basic step decay. This provides a baseline metric for loss and accuracy, against which the performance of the more advanced schedule can be measured.[17]
2. Find the Optimal Peak Learning Rate: The most crucial step is to determine an appropriate range for the peak learning rate, η_max. A highly effective empirical method is the learning rate range test. This involves running the training for a single epoch while exponentially increasing the learning rate from a very small value to a very large one. By plotting the training loss against the learning rate, one can identify the point at which the loss begins to diverge or explode. The optimal η_max for the full training run is typically found in the region of the steepest decline on this plot, often about one order of magnitude smaller than the rate at which divergence occurred.[24]
3. Set Initial Warmup and Decay Durations: Once a candidate η_max is chosen, the durations for the schedule’s phases can be set based on the total planned training steps. A common starting point is to allocate 10% of the total steps to the warmup phase. The remaining 90% of the steps then constitute the decay duration.
4. Iterate and Refine: With these initial parameters, begin full training runs while closely monitoring the training and validation loss curves. The shape of these curves provides valuable feedback. If the initial phase is unstable, a longer warmup period may be required. A longer warmup may, in turn, allow for a slightly higher and more aggressive peak learning rate. The goal is to find a combination that results in a stable initial ramp-up followed by a steady, consistent decrease in validation loss throughout the decay phase.

Best Practices and Common Pitfalls

Beyond tuning, several practical considerations are crucial for the successful implementation of learning rate schedulers.

• Best Practice: Update per Iteration, Not per Epoch: For schedules that include a warmup phase, it is almost always preferable to update the scheduler’s state after each optimizer step (i.e., per iteration) rather than only at the end of each epoch. An iteration-based update provides a much smoother and more fine-grained ramp-up of the learning rate, which is especially important during the very first epoch of training.[41]
• Best Practice: Correct step() Order: In frameworks like PyTorch, the order of operations within the training loop is critical. The optimizer.step() method, which updates the model’s weights, should always be called before the scheduler.step() method. The scheduler’s function is to calculate and set the learning rate that will be used for the next optimization step, based on the current state of training.[54]
• Pitfall: Mismatching Step/Epoch Units: A frequent source of error is a mismatch between the units used to define the scheduler’s parameters (e.g., T_max in epochs) and the frequency of the scheduler.step() call (e.g., per iteration). This can cause the schedule to run much faster or slower than intended, leading to unexpected and suboptimal learning rate trajectories. Ensure consistency in units throughout the implementation.
• Pitfall: Forgetting Optimizer and Scheduler State: When training is interrupted and needs to be resumed, it is essential to save and load not only the model’s weights but also the state dictionaries of both the optimizer and the learning rate scheduler. Failing to restore the scheduler’s state will cause it to restart from the beginning, disrupting the intended learning rate trajectory and likely harming the model’s performance.[22, 54]

Table 2: Hyperparameter Tuning Guide for Cosine Annealing with Warmup

The following table serves as a practical reference guide for tuning the key hyperparameters of the cosine annealing with warmup schedule. It maps the conceptual parameters to their common names in popular frameworks and provides heuristics and insights into their impact.

Hyperparameter	Framework Parameter(s)	Description	Typical Range / Heuristic	Impact of Increasing Value
Peak Learning Rate	learning_rate (HF/Keras), max_lr (some PyTorch libs)	The maximum learning rate reached after warmup.	1e-5 to 1e-3 (task-dependent). Find via LR range test.	Faster initial convergence, but higher risk of instability and divergence.
Warmup Duration	warmup_steps, warmup_ratio, num_warmup_epochs	Number of steps/epochs for the linear ramp-up.	5-15% of total training steps.	More stable start, allows higher peak LR. But, slows down initial training if too long.
Decay Duration	T_max, decay_steps	Number of steps/epochs for the cosine decay phase.	total_steps - warmup_steps.	Slower, more gradual decay. If too short, decay is too aggressive.
Minimum Learning Rate	eta_min (PyTorch), alpha (Keras)	The lowest learning rate at the end of the decay.	0 or 0.01 * peak_lr.	Allows for continued fine-tuning at the end of training.
Restart Cycle Length	T_0 (for WarmRestarts)	The length of the first cycle before a warm restart.	Problem-dependent (e.g., 10-50 epochs).	More frequent restarts lead to more exploration but can be disruptive.
Restart Cycle Multiplier	T_mult (for WarmRestarts)	Factor by which the cycle length increases after each restart.	1.0 (constant cycles) or 2.0 (doubling cycles).	Increases the duration of later fine-tuning periods.

Reference:

1. CosineAnnealingLR — PyTorch 2.7 documentation, accessed on July 7, 2025, https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.CosineAnnealingLR.html
2. Cosine Annealing Explained | Papers With Code, accessed on July 7, 2025, https://paperswithcode.com/method/cosine-annealing
3. SGDR: Stochastic Gradient Descent with Warm Restarts, accessed on July 7, 2025, https://arxiv.org/pdf/1608.03983
4. Revision History for SGDR: Stochastic Gradient Descent… – OpenReview, accessed on July 7, 2025, https://openreview.net/revisions?id=pcE3tjLzKg
5. CosineAnnealingWarmRestarts — PyTorch 2.7 documentation, accessed on July 7, 2025, https://docs.pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.CosineAnnealingWarmRestarts.html
6. CosineAnnealingWarmRestarts – BrainPy documentation – Read the Docs, accessed on July 7, 2025, https://brainpy.readthedocs.io/en/latest/apis/generated/brainpy.optim.CosineAnnealingWarmRestarts.html
7. Beyond Cosine Decay: On the effectiveness of Infinite Learning Rate Schedule for Continual Pre-training – arXiv, accessed on July 7, 2025, https://arxiv.org/html/2503.02844v1
8. Analyzing & Reducing the Need for Learning Rate Warmup in GPT Training – arXiv, accessed on July 7, 2025, https://arxiv.org/html/2410.23922v1
9. Why Warmup the Learning Rate? Underlying Mechanisms and Improvements – arXiv, accessed on July 7, 2025, https://arxiv.org/html/2406.09405v1
10. NeurIPS Poster Why Warmup the Learning Rate? Underlying Mechanisms and Improvements, accessed on July 7, 2025, https://neurips.cc/virtual/2024/poster/95431
11. [2406.09405] Why Warmup the Learning Rate? Underlying Mechanisms and Improvements, accessed on July 7, 2025, https://arxiv.org/abs/2406.09405
12. Ablation study in learning rate warm-up | Download Scientific Diagram, accessed on July 7, 2025, https://www.researchgate.net/figure/Ablation-study-in-learning-rate-warm-up_tbl2_389387589
13. Comparison between constant lr scheduler and cosine annealing lr… – ResearchGate, accessed on July 7, 2025, https://www.researchgate.net/figure/Comparison-between-constant-lr-scheduler-and-cosine-annealing-lr-scheduler-with-linear_fig1_336936339
14. Analyzing & Reducing the Need for Learning Rate … – OpenReview, accessed on July 7, 2025, https://openreview.net/pdf?id=ZgDNrpS46k
15. Analyzing & Reducing the Need for Learning Rate Warmup in GPT Training – OpenReview, accessed on July 7, 2025, https://openreview.net/forum?id=ZgDNrpS46k¬eId=Mar5dgqcSh
16. Linear Warmup With Cosine Annealing Explained – Papers With Code, accessed on July 7, 2025, https://paperswithcode.com/method/linear-warmup-with-cosine-annealing
17. LinearWarmupCosineAnnealing, accessed on July 7, 2025, https://lightning-flash.readthedocs.io/en/stable//api/generated/flash.core.optimizers.LinearWarmupCosineAnnealingLR.html
18. Why Warmup the Learning Rate? Underlying Mechanisms and …, accessed on July 7, 2025, https://openreview.net/forum?id=NVl4SAmz5c¬eId=bavOSLCuxE
19. Using both learning rate warm up and a learning rate scheduler – PyTorch Forums, accessed on July 7, 2025, https://discuss.pytorch.org/t/using-both-learning-rate-warm-up-and-a-learning-rate-scheduler/177767
20. Tony-Y/pytorch_warmup: Learning Rate Warmup in PyTorch – GitHub, accessed on July 7, 2025, https://github.com/Tony-Y/pytorch_warmup
21. Schedulers — pytorch-accelerated 0.1.3 documentation, accessed on July 7, 2025, https://pytorch-accelerated.readthedocs.io/en/latest/schedulers.html
22. katsura-jp/pytorch-cosine-annealing-with-warmup – GitHub, accessed on July 7, 2025, https://github.com/katsura-jp/pytorch-cosine-annealing-with-warmup
23. tf.keras.optimizers.schedules.CosineDecay | TensorFlow v2.16.1, accessed on July 7, 2025, https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/schedules/CosineDecay
24. tfm.optimization.CosineDecayWithOffset | TensorFlow v2.16.1, accessed on July 7, 2025, https://www.tensorflow.org/api_docs/python/tfm/optimization/CosineDecayWithOffset
25. Mr-TalhaIlyas/Learning-Rate-Schedulers-Packege-Tensorflow-PyTorch-Keras – GitHub, accessed on July 7, 2025, https://github.com/Mr-TalhaIlyas/Learning-Rate-Schedulers-Packege-Tensorflow-PyTorch-Keras