Making Model Training a Bit More Scientific (Part 7): Step Size Scheduling and Weight Averaging
Anyone experienced in model training knows that step size scheduling, or learning rate scheduling (LR Schedule), is crucial to a model's final performance. As we've derived in earlier posts in this series, even for plain SGD, the learning rate function that is optimal for terminal convergence naturally takes on the Warmup-Decay shape so commonly used in practice. That said, some recent work — such as Schedule-Free — has tried to replace learning rate scheduling with some form of weight averaging, and has made real progress along these lines.
How can we analyze, theoretically, the connection between learning rate scheduling and weight averaging? To what extent can one substitute for the other? This post tries to answer these questions.
Weight Averaging
As for weight averaging, we already gave a basic result on this in Making Model Training a Bit More Scientific (Part 3): Terminal Loss Convergence for SGD:
\begin{equation}\mathbb{E}[L(\bar{\boldsymbol{\theta}}_T) - L(\boldsymbol{\theta}^*)] \leq \frac{\Vert\boldsymbol{\theta}_1 - \boldsymbol{\theta}^*\Vert^2}{2T\eta_T} + \frac{G^2}{2T}\sum_{t=1}^T \frac{\eta_t^2}{\eta_T}\end{equation}more
where $\bar{\boldsymbol{\theta}}_T = \frac{1}{T}\sum_{t=1}^T \boldsymbol{\theta}_t$. This result means that if we choose a constant learning rate $\eta\propto 1/\sqrt{T}$, the right-hand side achieves the ideal convergence rate $\mathcal{O}(1/\sqrt{T})$, while the left-hand side implies that $\bar{\boldsymbol{\theta}}_T$ converges to $\boldsymbol{\theta}^*$. This motivates the following optimizer design
\begin{equation}\begin{aligned} \boldsymbol{\theta}_{t+1} =&\, \boldsymbol{\theta}_t - \eta\, \boldsymbol{g}(\boldsymbol{x}_t, \boldsymbol{\theta}_t) \\[4pt] \boldsymbol{\mu}_{t+1} =&\, (1 - c_{t+1}) \boldsymbol{\mu}_t + c_{t+1} \boldsymbol{\theta}_{t+1} \end{aligned}\label{eq:pr}\end{equation}
where $c_t = 1/t$. In other words, we add one extra step that has no effect on the training trajectory itself: we simply average the trajectory with equal weights, and the model weights we eventually use for prediction are not the terminal value $\boldsymbol{\theta}_T$ but the average $\boldsymbol{\mu}_T$. In practice, however, this approach performs poorly — one of the classic examples of theory failing to match practice.
That said, if we replace $c_t$ with a suitable constant that doesn't depend on $t$ — say, $10^{-3}$ — then in practice this often performs quite well. This is exactly the exponential moving average (EMA) trick the author often used in competitions. Compared with $c_t=1/t$, this moving-average approach tends to average over the most recent part of the optimization trajectory rather than the entire history, and the fact that it works well suggests we don't need to pay much attention to distant history.
Explaining this moving-average trick theoretically — or, more generally, explaining the viability of weight averaging in general — is the main task ahead.
Generalizing Again
In the post Making Model Training a Bit More Scientific (Part 4): A New Identity, A New Learning Rate, in order to extend the average-loss result to a terminal-loss result, we introduced the identity
\begin{equation}q_T = \frac{1}{w_{1:T}}\sum_{t=1}^T w_t q_t + \sum_{k=1}^{T-1}\left(\frac{1}{w_{k+1:T}} - \frac{1}{w_{k:T}}\right)\sum_{t=k+1}^T w_t (q_t - q_k)\end{equation}
where $w_{k:T}\triangleq\sum_{t=k}^T w_t$. In this section, we generalize this once more to
\begin{equation}\frac{1}{v_{1:T}}\sum_{t=1}^T v_t q_t = \frac{1}{w_{1:T}}\sum_{t=1}^T w_t q_t + \frac{1}{v_{1:T}}\sum_{k=1}^{T-1}\left(\frac{v_{k+1:T}}{w_{k+1:T}} - \frac{v_{k:T}}{w_{k:T}}\right)\sum_{t=k+1}^T w_t (q_t - q_k)\label{eq:qt-gg}\end{equation}
Intuitively, it provides a way to convert one weighted sum into another weighted sum: when $v_T=1$ and the remaining $v_t$ are all zero, it degenerates to the terminal form. The proof is in fact a generalization of the earlier one. Writing $S_k = \frac{1}{w_{k:T}}\sum_{t=k}^{T}w_t q_t$, we can write
\begin{equation}\begin{aligned} v_{k:T} S_k - v_{k+1:T} S_{k+1} =&\, v_{k:T} \left(\frac{w_{k+1:T}S_{k+1} + w_k q_k}{w_{k:T}}\right) - v_{k+1:T} S_{k+1} \\ =&\, v_k q_k + \left(\frac{v_{k:T}}{w_{k:T}} w_k - v_k\right) q_k + \left(\frac{v_{k:T}}{w_{k:T}}w_{k+1:T} - v_{k+1:T}\right) S_{k+1} \\ =&\, v_k q_k + \left(\frac{v_{k:T}}{w_{k:T}}w_{k+1:T} - v_{k+1:T}\right) (S_{k+1} - q_k) \\ =&\, v_k q_k + \left(\frac{v_{k:T}}{w_{k:T}} - \frac{v_{k+1:T}}{w_{k+1:T}}\right) (w_{k+1:T} S_{k+1} - w_{k+1:T} q_k) \\ =&\, v_k q_k + \left(\frac{v_{k:T}}{w_{k:T}} - \frac{v_{k+1:T}}{w_{k+1:T}}\right) \sum_{t=k+1}^T w_t (q_t - q_k) \end{aligned}\end{equation}
The third equality holds because adding the two parenthesized terms from the second equality gives zero, so the two are negatives of each other. Now summing both sides over $k=1\sim T-1$, the left-hand side equals $v_{1:T} S_1 - v_T q_T$; moving $v_T q_T$ to the right-hand side, dividing both sides by $v_{1:T}$, and rearranging slightly gives the identity we wanted to prove, $\eqref{eq:qt-gg}$.
Rescaling Transformation
Next, following the approach of the previous post Making Model Training a Bit More Scientific (Part 6): An Elegant Top-Down Construction, we derive a more general convergence result. Setting $q_t=\mathbb{E}[L(\boldsymbol{\theta}_t)-L(\boldsymbol{\theta}^*)]$ and substituting into equation $\eqref{eq:qt-gg}$ gives
\begin{equation}\frac{1}{v_{1:T}}\sum_{t=1}^T v_t \mathbb{E}[L(\boldsymbol{\theta}_t) - L(\boldsymbol{\theta}^*)] = \frac{1}{w_{1:T}}\sum_{t=1}^T w_t \mathbb{E}[L(\boldsymbol{\theta}_t) - L(\boldsymbol{\theta}^*)] + \frac{1}{v_{1:T}}\sum_{k=1}^{T-1}\left(\frac{v_{k+1:T}}{w_{k+1:T}} - \frac{v_{k:T}}{w_{k:T}}\right)\sum_{t=k+1}^T w_t \mathbb{E}[L(\boldsymbol{\theta}_t) - L(\boldsymbol{\theta}_k)]\end{equation}
Using the convexity of $L$ to bound $\mathbb{E}[L(\boldsymbol{\theta}_t)-L(\boldsymbol{\theta}^*)]$ and $\mathbb{E}[L(\boldsymbol{\theta}_t)-L(\boldsymbol{\theta}_k)]$, and then carefully rearranging as in the "Identity Transformation" section of the previous post, we obtain
\begin{equation}\frac{1}{v_{1:T}}\sum_{t=1}^T v_t\,\mathbb{E}[L(\boldsymbol{\theta}_t)-L(\boldsymbol{\theta}^*)]\leq \frac{1}{w_{1:T}}\sum_{t=1}^T w_t\,\mathbb{E}[\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\cdot(\boldsymbol{\psi}_t-\boldsymbol{\theta}^*)] \label{leq:mid}\end{equation}
where
\begin{equation}\boldsymbol{\psi}_t=\frac{w_{1:T}}{v_{1:T}}\left[\frac{v_{t:T}}{w_{t:T}}\boldsymbol{\theta}_t-\sum_{k=1}^{t-1}\left(\frac{v_{k+1:T}}{w_{k+1:T}}-\frac{v_{k:T}}{w_{k:T}}\right)\boldsymbol{\theta}_k\right]\end{equation}
and direct verification gives
\begin{equation}\boldsymbol{\psi}_{t+1}-\boldsymbol{\psi}_t=\frac{v_{t+1:T}}{v_{1:T}}\frac{w_{1:T}}{w_{t+1:T}}(\boldsymbol{\theta}_{t+1}-\boldsymbol{\theta}_t)\end{equation}
Therefore, if we update $\boldsymbol{\psi}_t$ according to $\boldsymbol{\psi}_{t+1}=\boldsymbol{\psi}_t-w_t\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)$, then the corresponding update rule for $\boldsymbol{\theta}_t$ is
\begin{equation}\boldsymbol{\theta}_{t+1}=\boldsymbol{\theta}_t-\underbrace{\frac{v_{1:T}}{v_{t+1:T}}\frac{w_t w_{t+1:T}}{w_{1:T}}}_{\eta_t}\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\end{equation}
The General Result
The remaining steps are fairly straightforward. Applying the conclusions of posts such as Making Model Training a Bit More Scientific (Part 5): Fine-Tuning the Learning Rate Based on Gradients to the right-hand side of equation $\eqref{leq:mid}$, we get
\begin{equation}\frac{1}{w_{1:T}}\sum_{t=1}^T w_t\,\mathbb{E}[\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\cdot(\boldsymbol{\psi}_t-\boldsymbol{\theta}^*)] \leq \frac{1}{2 w_{1:T}}\left(R^2 + \sum_{t=1}^T w_t^2 G_t^2\right) \end{equation}
where $R = \Vert\boldsymbol{\theta}_1 - \boldsymbol{\theta}^*\Vert, G_t^2 = \mathbb{E}[\Vert\boldsymbol{g}(\boldsymbol{x}_t, \boldsymbol{\theta}_t)\Vert^2]$, with the convention $\boldsymbol{\psi}_1=\boldsymbol{\theta}_1$. As for the left-hand side, applying convexity gives
\begin{equation}\mathbb{E}\left[L\left(\frac{1}{v_{1:T}}\sum_{t=1}^T v_t \boldsymbol{\theta}_t \right) - L(\boldsymbol{\theta}^*)\right] \leq \frac{1}{v_{1:T}}\sum_{t=1}^T v_t\,\mathbb{E}[L(\boldsymbol{\theta}_t)-L(\boldsymbol{\theta}^*)]\end{equation}
Putting these together we obtain
\begin{equation}\mathbb{E}\left[L\left(\frac{1}{v_{1:T}}\sum_{t=1}^T v_t \boldsymbol{\theta}_t \right) - L(\boldsymbol{\theta}^*)\right] \leq \frac{1}{2 w_{1:T}}\left(R^2 + \sum_{t=1}^T w_t^2 G_t^2\right)\label{eq:avg-final}\end{equation}
The right-hand side attains its minimum when
\begin{equation}w_t = \frac{R G_t^{-2}}{\sqrt{Q_T}},\qquad Q_T=\sum_{k=1}^T G_k^{-2}\end{equation}
which means the optimal learning rate takes the form
\begin{equation}\eta_t = \frac{v_{1:T}}{v_{t+1:T}}\frac{w_t w_{t+1:T}}{w_{1:T}} = \frac{v_{1:T}}{v_{t+1:T}}\frac{R G_t^{-2}}{\sqrt{Q_T}} \left(1 - \frac{Q_t}{Q_T}\right) \label{eq:opt-lr}\end{equation}
Interpreting the Result
The previous post Making Model Training a Bit More Scientific (Part 6): An Elegant Top-Down Construction gave the following terminal convergence result:
\begin{equation}\mathbb{E}\left[L(\boldsymbol{\theta}_T) - L(\boldsymbol{\theta}^*)\right] \leq \frac{1}{2 w_{1:T}}\left(R^2 + \sum_{t=1}^T w_t^2 G_t^2\right)\end{equation}
The right-hand side agrees with equation $\eqref{eq:avg-final}$, but the learning rate becomes $\eta_t = \frac{w_t w_{t+1:T}}{w_{1:T}}$. This shows that if we want weight averaging to match the performance of the terminal weights, the learning rate needs to be multiplied by an extra factor of $\frac{v_{1:T}}{v_{t+1:T}}$ — or, put the other way around, under the same learning rate, the effect of weight averaging is equivalent to the terminal-point effect obtained after multiplying the learning rate by $\frac{v_{t+1:T}}{v_{1:T}}$. Clearly $\frac{v_{t+1:T}}{v_{1:T}}$ is at most 1 and monotonically decreasing, so multiplying by this factor plays a role similar to LR decay.
Let's look at some concrete examples. First, the simplest case, $v_t\equiv 1$, corresponding to equation $\eqref{eq:pr}$ at the beginning of this post; here $\frac{v_{t+1:T}}{v_{1:T}} = 1 - t/T$, which is exactly linear decay. This means that "constant learning rate + equal-weight averaging" should, in theory, have an effect similar to a linearly decaying learning rate. In practice, however, what we usually observe is that directly applying linear decay to the learning rate works quite well at the terminal point, whereas the averaged weights from equation $\eqref{eq:pr}$ underperform.
Looking back over the whole derivation, one can guess that the problem lies in the "convex function" assumption. This entire series of posts relies on the assumption that the objective function is convex, but in practice it is usually non-convex, so the conclusions of convex optimization cannot be reproduced exactly. A compromise is to assume local convexity, in which case we can do equal-weight averaging over a finite interval — or, alternatively, consider its smoothed version, the exponential moving average (EMA). In that case, $v_t = \gamma^{-t}$ (where $0 < \gamma < 1$), and
\begin{equation}\frac{v_{t+1:T}}{v_{1:T}} = \frac{1 - \gamma^{T-t}}{1 - \gamma^T}\end{equation}
which maintains an averaging window of length $\mathcal{O}((1-\gamma)^{-1})$. Because of the memorylessness of the exponential function, if we run EMA alongside constant-learning-rate training, it's as if at every step we simultaneously obtain the weights we would get after $\mathcal{O}((1-\gamma)^{-1})$ steps of linear decay — letting us peek ahead at what the effect of linear decay would be.
Related Work
Regarding the underperformance of equation $\eqref{eq:pr}$, the paper The Road Less Scheduled proposed "Schedule-Free," which changes the point at which the gradient is evaluated to an interpolation between $\boldsymbol{\theta}_t$ and $\boldsymbol{\mu}_t$:
\begin{equation}\begin{aligned} \hat{\boldsymbol{\theta}}_t =&\, (1-\beta)\boldsymbol{\theta}_t + \beta \boldsymbol{\mu}_t \\[4pt] \boldsymbol{\theta}_{t+1} =&\, \boldsymbol{\theta}_t - \eta\, \boldsymbol{g}(\boldsymbol{x}_t, \hat{\boldsymbol{\theta}}_t) \\[4pt] \boldsymbol{\mu}_{t+1} =&\, (1 - c_{t+1}) \boldsymbol{\mu}_t + c_{t+1} \boldsymbol{\theta}_{t+1} \end{aligned}\end{equation}
Intuitively, the original averaging operation was completely decoupled from the training process — the model never "sensed" whether $\boldsymbol{\mu}_t$ was actually good or not. Now, by evaluating the gradient at an interpolation between $\boldsymbol{\theta}_t$ and $\boldsymbol{\mu}_t$, the model can sense the quality of $\boldsymbol{\mu}_t$ and adjust accordingly. This modification comes with theoretical guarantees, and it does indeed noticeably improve the performance of the averaged weights $\boldsymbol{\mu}_T$.
However, as training scale increased, Schedule-Free gradually began to show weaknesses, and the authors recently proposed ScheduleFree+, adding quite a few "patches" to the original Schedule-Free to bring it back to SOTA status. That said, the Plus version introduces a lot of extra machinery, making the whole algorithm considerably more complex and harder to understand, so we won't go into detail here — we'll come back to discuss it later.
It's worth noting that the authors of Schedule-Free and its Plus version are the same authors of the classic paper Optimal Linear Decay Learning Rate Schedules and Further Refinements that we introduced in earlier posts in this series. This body of work has successfully applied classical convex optimization theory to real practical training, rather than remaining confined to "armchair theorizing" on toy models — and that deserves our respect!
Further Thoughts
What's really worth pondering is whether the premise of Schedule-Free actually holds up. "Schedule-Free" has two goals: first, to train from start to finish with a constant learning rate; second, to have all hyperparameters be independent of the total number of training steps $T$. If both conditions hold, training could be stopped at any point and resumed at any point, and no matter where you stop, you'd have the optimal solution given the steps taken so far — which would indeed be ideal.
However, both of these goals are debatable. First, for actual training, there's no fundamental difference between using a constant learning rate versus a dynamic one like cosine decay or linear decay — either way, you tune hyperparameters and pick the best-performing option, so pursuing a constant learning rate per se has little practical significance. Moreover, even Schedule-Free and its Plus version haven't managed to escape the dependence on Warmup, so in practice they still, to some degree, rely on a manually designed LR schedule.
Second, wanting the hyperparameters to be independent of the total number of training steps $T$ is indeed a worthy goal, but it's hard to justify purely theoretically. Take equation $\eqref{eq:pr}$ as an example again: even under all the stated assumptions, while equation $\eqref{eq:pr}$ performs well, its optimal learning rate $\eta$ still depends on $T$. Likewise, Schedule-Free cannot escape this fact either — it achieves the result that $\eta$ is independent of $T$ only by introducing new assumptions.
We can also see this from equation $\eqref{eq:opt-lr}$: for equal-weight averaging, $\frac{v_{1:T}}{v_{t+1:T}}=1/(1-t/T)$; if we further assume $G_t$ is a constant $G$, then $1 - Q_t/Q_T=1-t/T$, which cancels out exactly, so the optimal learning rate is $\frac{R G^{-2}}{\sqrt{Q_T}}$, and $Q_T$ does depend on $T$. On the other hand, $G_t$ is usually large early on and small later, so given that, according to equation $\eqref{eq:opt-lr}$, the learning rate is inversely proportional to $G_t^2$, the learning rate should be smaller early in training — which actually explains why Schedule-Free still can't completely shed the need for a schedule, and still needs Warmup.
For this reason, ScheduleFree+ and some contemporaneous work have gradually shifted their focus from "Schedule-Free" toward "LR-Free." But overall, this area is still in its early stages, and a more fundamental exploration remains to be done.
Summary
This post introduced a new identity, further generalizing the results of the previous post, which lets us quantitatively reveal, from a theoretical standpoint, the connection between weight averaging and learning rate decay. We also discussed why equal-weight averaging underperforms moving averages, and why Schedule-Free still cannot fully escape the need for a schedule.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.