Hop, skip, jump to Convergence: Dynamics of Learning Rate Transitions for Improved Training of Large Language Models (2024)

Abstract: Various types of learning rate (LR) schedulers are being used for training or fine tuning of Large Language Models today. In practice, several mid-flight changes are required in the LR schedule either manually, or with careful choices around warmup steps, peak LR, type of decay and restarts. To study this further, we consider the effect of switching the learning rate at a predetermined time during training, which we refer to as “SkipLR”. We model SGD as a stochastic gradient flow and show that when starting from the same initial parameters, switching the learning rate causes the loss curves to contract towards each other. We demonstrate this theoretically for some simple cases, and empirically on large language models. Our analysis provides insight into how learning rate schedules affect the training dynamics, and could inform the design of new schedules to accelerate convergence.

## introduction modern deep neural networks have achieved stateof-the-art performance across a wide range of machine learning tasks. one critical hyperparameter that significantly influences the training dynamics is the learning rate (lr). aside from using adaptive optimizers such as adam @xcite , adamw @xcite , and adafactor @xcite , lr schedulers have become a necessity for training large language models (llms) @xcite . despite being used pervasively, the choice of lr schedules is often made based on empirical observations or best practices, rather than any known theoretical basis. why are such abrupt changes in lr required to achieve peak performance, and how do they impact training? in this paper, we investigate these questions by studying the impact of drastic lr changes on the training of language models. our theoretical and empirical analyses elucidate the dynamics of optimizers under varying lr schedules, focusing on abrupt lr transitions at predetermined epochs, which we term the "skiplr" experiment. first, we aim to understand the impact of lr schedules on training by utilizing the skiplr framework to study the relationship between multiple lr transitions and changes in loss. while traditional schedulers emphasize the long-term effects of continuous changes in the lr, we argue that even instantaneous changes in lr during a single step can have long-term impacts on loss. we clarify here that skiplr is a framework to study this specific phenomenon, and not a scheduler to be used in place of a scheduler like cosine or exponential decay. through theoretical analysis and hundreds model optimizer initial lr scheduler citation falcon 40b adamw 1.85𝑒 -4 cosine decay to 10% initial lr (almazrouei et al., 2023) meta llama 2 adamw 1.5𝑒 -4 cosine decay to 10% initial lr (touvron et al., 2023) flan t5, palm adafactor 5𝑒 -4 constant (chung et al., 2022) gpt-j adamw 1.2𝑒 -4 polynomial decay to 10% initial lr (wang and komatsuzaki, 2021) opt 1.3b adamw 2𝑒 -4                warmup from 0 to initial lr followed by decaying down to 10% of the maximum lr over 300b tokens along with a number of mid-flight changes (zhang et al., 2022) opt 6.7b adamw 1.2𝑒 -4 opt 30b adamw 1𝑒 -4 opt 66b adamw 0.8𝑒 -4 opt 175b adamw 1.2𝑒 -4 instructgpt adam 9𝑒 -6 cosine decay to 10% initial lr (ouyang et al., 2022) llama 3.1 70b adamw 1.5𝑒 -4 warmup 2000 steps with cosine decay (dubey et al., 2024) llama 3.1 405b adamw 8𝑒 -5 warmup 8000 steps with cosine decay to 1% of peak lr over 1.2m steps (dubey et al., 2024) to demonstrate this, we first investigate how skiplr transitions affect the descent on the loss surface at any given iteration. we model sgd dynamics as simple yet functionally correct continuoustime gradient flow, with 𝑑 𝜃 𝑑𝑡 = -𝜂∇φ(𝜃 (𝑡)), where φ(𝜃) denotes the loss function. our theoretical results show that, under common assumptions, loss function trajectories contract towards each other after a skiplr transition as illustrated in figure 1 . this suggests that skiplr techniques could help escape sharp local minima or transition between high and low loss regions during training in a principled manner. next, to support our theoretical findings, we present extensive numerical simulations visualizing loss curve contractions on synthetic functions. specifically, for quadratic objectives, these simulations validate our analysis by demonstrating the contraction of loss curves following skips in strongly convex functions. we also demonstrate successful skiplr transitions on large transformer networks for causal language modeling, translation, and named entity recognition (ner) tasks, empirically showing convergence in these models. in this work, through our comprehensive analysis and experimentation with skiplr, we aim to provide valuable insights that can guide informed decisions on learning rate schedules and adjustments during training. by quantifying the impact of abrupt lr changes and exploring their interaction with the loss landscape and gradient dynamics, we hope to enhance the design of adaptive learning rate schedules tailored to specific models and training conditions. ## related work learning rate schedulers, such as linear or cosine scheduling @xcite , have been extensively studied and shown to accelerate convergence and improve training performance in neural networks. for instance, @xcite discusses the effect of learning rate annealing on generalization performance in convex learning problems, demonstrating that annealing leads to better generalization compared to using a constant small learning rate. this improvement is attributed to the mismatch between the test and train loss landscapes and the benefits of early stopping. additionally, they have studied the impact of learning rate schedules on sgd and their influence on the local geometry of the parameter space. in linear scheduling @xcite , where the learning rate decreases linearly over time. this approach has been found effective for achieving fast convergence and improving training performance by allowing the model to finetune its parameters gradually. another technique, warmup steps @xcite , involves gradually increasing the learning rate at the beginning of training before applying the scheduled learning rate. warmup steps help the model explore the parameter space more effectively and avoid local minima, thus improving convergence.uickly adapt to the training data and improve convergence. cosine scheduling @xcite ) is another widely used approach, where the learning rate decreases following a cosine function over time. this method has been shown to improve convergence and prevent models from getting stuck in sub-optimal solutions. research has also been conducted on restarts in learning rate scheduling and their impact on convergence and training. restarting the learning rate schedule involves periodically resetting the learning rate to its initial value during training. this technique has been found to improve convergence and prevent the model from getting trapped in poor local minima @xcite . by periodically resetting the learning rate, restarts allow the model to explore different regions of the parameter space and potentially find better solutions. other schedulers that control the training have also been shown to accelerate convergence and improve the performance of models. for example, in the paper by @xcite , the authors propose a method called "decaying loss" where a fixed learning rate is used, but the magnitude of the update is controlled by gradually reducing the impact of noise on the network. another paper by @xcite introduces eigencurve, a family of learning rate schedules that achieve minimax optimal convergence rates for stochastic gradient descent (sgd) on quadratic objectives with skewed hessian spectrums. the authors show that eigencurve outperforms step decay in image classification tasks, especially when the number of epochs is small. the proposed schedulers are designed to approximate eigencurve and show superior performance compared to cosine decay in certain situations. in federated learning, @xcite propose a device scheduling policy to achieve fast convergence by considering the trade-off between the number of rounds required to attain a certain model accuracy and the latency per round. their greedy policy selects devices with the least time consumption in model updating, resulting in a good trade-off between learning efficiency and latency per round. our previous work also explored scheduling such as using reinforcement learning (rl) to learn lr schedules as a policy @xcite . overall, these papers provide evidence that learning rate schedulers can accelerate convergence and improve the performance of models. this has been well known in research in deep learning. the proposed methods, such as decaying loss, eigencurve, learning rate annealing, and device scheduling policies, offer different strategies for controlling the learning rate and achieving better convergence rates. several of these schedulers impose drastic changes to the lr schedule which we hypothesize can cause long term changes to the loss trajectory. our intention is to show a specific transition phenomenon via theoretical and experimental analysis. in the next section, we model sgd dynamics and prove that after a transition, a loss curve can converge or contract towards another loss curve that is predetermined by the learning rates involved in the transition. ## proof sketch sgd dynamics have been modeled in various ways in the past for the purpose of studying theoretical characteristics of the algorithm, and for informing the design of new optimizers and schedulers. in this section we model sgd dynamics in continuous time, and use this to prove how sudden changes in learning rate (lr) can cause one loss trajectory corresponding to the original lr to converge or contract towards the loss trajectory corresponding to the the changed lr. we then use these results and demonstrate the same phenomenon using theoretical and real-world experiments. first we state the standard sgd algorithm. here, 𝜃 𝑘 ∈ r 𝐾 is the model parameters at iteration 𝑘, φ 𝑖 : r 𝐾 → r is the loss function for data subset 𝑖, 𝜂 𝑘 > 0 is the learning rate at iteration 𝑘, 𝑁 is the number of data subsets or batches of data. this matches the sgd algorithm description in section 1.1 of @xcite . at a small enough discretization of steps, we model sgd dynamics by replacing 𝜃 ← 𝜃 -𝜂δφ 𝑖 (𝜃), by 𝑑 𝜃 𝑑𝑡 = -𝜂δφ 𝑖 (𝜃 (𝑡)). in practice, we also commonly see fixing 𝜂 𝑘 = 𝜂 0 , or using a scheduler (for example, cosine scheduler) to control the learning rate. while several choices of schedulers that are available continuously change the learning rate, our intuition is that even a single step change to a different learning rate would introduce a drastic change to the dynamics even if starting from the same initial point. to study the effect of changes in lr, we define the algorithm to be studied slightly differently by adding steps related to the skiplr transition to the original algorithm 1, as shown algorithm 2. here, we study the dynamics of a process where the learning rate switches at a predetermined time step or epoch. ## main result we can now state and prove the main result: theorem 1. let 𝜃 1 (𝑡), 𝜃 2 (𝑡), 𝜃 3 (𝑡) be trajectories evolving according to the sgp dynamics: what we intend to show is that when starting from the same initial point 𝜃 0 , that the conditions presented around 𝜂 using skip lr as shown in algorithm 2 cause the solutions φ 𝑖 to contract towards each other. ## limitations while our findings are promising, there are several important considerations. first, the theoretical analysis relies on strong convexity and smoothness assumptions about the loss landscape. while this is a standard setup in many theoretical settings, our work doesn't explore more complex settings, especially a setting that includes modern-day llms. second, our experiments are limited to fixed learning rates. we believe dynamic learning schedulers add an additional layer of complexity that can affect the skip phenomenon, which we do not explore in this work. ## conclusion in this work, we have analyzed the dynamics of sgd training under different learning rate schedules. we modeled sgd as a stochastic gradient flow and considered the effect of switching the learning rate at a predetermined time, which we call "skiplr". under standard assumptions on the loss landscape, we proved that the loss curves resulting from different constant learning rates contract towards each other after a skiplr switch. we verified this behavior theoretically for simple cases and empirically for large language models. our analysis gives insight into how learning rate schedules impact the training dynamics. this could help design new schedules that lead to faster convergence. an interesting direction for future work is to extend the analysis to stochastic processes with non-uniform sampling, and establish quantitative bounds on the convergence rates. c synthetic data experiments we run some theoretical experiments approximating the dynamics of φ using φ 0 (𝜃) = 𝜃, i.e. only a single potential function. when using skip lr scheduler, we see that the skips cause the curves to exponentially converge (fig. 1 a) and b). when 𝑁 = 2, we uniformly randomly switch between φ 0 (𝜃) = 𝜃 + 1 and φ 1 (𝜃) = 𝜃 -1, with two sets of learning rates. we use the runge kutta 5 solver from python3's scipy library for solving the initial value problem with a fixed 𝜃 (0). the stationary point 𝜃 * = 0 for all problems, and the loss function used is rmse; note that here since 𝜃 * = 0, the plot also reveals the parameter 𝜃 (𝑡). as we can see in all cases of figure 8, transition within some 𝜖 happens well before 𝑡 → ∞. note here that with this theoretical set up, the ranking of the loss curves corresponds to the values of the lrs used. that is, a higher learning rate corresponds to a steeper descent as expected; in the main paper we demonstrate how this is different for real-world transformer based training experiments. in these experiments modeling learning dynamics using potential functions, when transitioning from a higher average loss value to a lower one, the convergence rate effectively accelerates to eventually match the lower loss curve. this matches our proof (φ 2 (𝑡) → φ 3 (𝑡). we also note that the loss curves as 𝑁 increases beyond 1 with random switching between potential closely resemble what one might see when training standard deep neural networks.