Making Model Training More Scientific (IV): A New Identity, A New Learning Rate

让炼丹更科学一些（四）：新恒等式，新学习率 苏剑林 · 2025-12-26 · Mathematics learning rate optimizer SGD model tuning

In the previous post, Making Model Training More Scientific (III): Endpoint Loss Convergence of SGD], we successfully converted the convergence conclusion from average loss to endpoint loss, obtaining a convergence rate of $\mathcal{O}(\sqrt{\ln T/T})$. However, on closer reflection this result seems somewhat counterintuitive: intuitively, the endpoint loss should be closer to the optimum, and yet the average loss can already achieve a rate of $\mathcal{O}(1/\sqrt{T})$ — so why would the endpoint convergence rate be slower instead?

The latest progress on this question comes from Optimal Linear Decay Learning Rate Schedules and Further Refinements]. The paper first generalizes the key identity we proved earlier, then points out the importance of learning rate scheduling for endpoint convergence, thereby accelerating the endpoint loss convergence rate to $\mathcal{O}(1/\sqrt{T})$.

A New Identity

The original paper's results are rich, and we will cover them across several posts. This post mainly continues along the line of thought from the previous one, giving an initial introduction. To convert the average loss convergence conclusion into one about endpoint loss, the previous post introduced the following key identity:

\begin{equation}q_T = \frac{1}{T}\sum_{t=1}^T q_t + \sum_{k=1}^{T-1} \frac{1}{k(k+1)}\sum_{t=T-k}^T (q_t - q_{T-k})\end{equation}more

In this post we generalize it to a weighted-average version: defining $w_{k:T}\triangleq\sum_{t=k}^T w_t$, we have

\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}^T w_t (q_t - q_k)\label{eq:qt-g}\end{equation}

The proof strategy is basically the same. Let $\lambda_k = \sum_{t=T-k+1}^T w_t, S_k = \frac{1}{\lambda_k}\sum_{t=T-k+1}^T w_t q_t$, then

\begin{equation}\begin{aligned} \lambda_k S_k =&\, \lambda_{k+1} S_{k+1} - w_{T-k} q_{T-k} \\[5pt] =&\, \lambda_k S_{k+1} + w_{T-k} (S_{k+1} - q_{T-k}) \\ =&\, \lambda_k S_{k+1} + \frac{w_{T-k}}{\lambda_{k+1}}\sum_{t=T-k}^T w_t (q_t - q_{T-k}) \end{aligned}\end{equation}

Dividing both sides by $\lambda_k$ and summing over $k=1\sim T-1$ gives

\begin{equation}S_1 = S_T + \sum_{k=1}^{T-1}\frac{w_{T-k}}{\lambda_k\lambda_{k+1}}\sum_{t=T-k}^T w_t (q_t - q_{T-k})\end{equation}

Noting that $\frac{w_{T-k}}{\lambda_k\lambda_{k+1}} = \frac{1}{\lambda_k} - \frac{1}{\lambda_{k+1}}$, and substituting in the definition of $S_1, S_T$, we get

\begin{equation}q_T = \frac{1}{\sum_{t=1}^T w_t}\sum_{t=1}^T w_t q_t + \sum_{k=1}^{T-1}\left(\frac{1}{\sum_{t=T-k+1}^T w_t} - \frac{1}{\sum_{t=T-k}^T w_t}\right)\sum_{t=T-k}^T w_t (q_t - q_{T-k})\end{equation}

Finally, substituting $T-k$ with $k$ yields identity $\eqref{eq:qt-g}$.

A General Conclusion

Next, we again start from the core inequality of the second post, Making Model Training More Scientific (II): Extending the Conclusion to Unbounded Domains]:

\begin{equation}\sum_{t=1}^T \eta_t \mathbb{E}[L(\boldsymbol{\theta}_t) - L(\boldsymbol{\varphi})]\leq \frac{\Vert\boldsymbol{\theta}_1 - \boldsymbol{\varphi}\Vert^2}{2} + \frac{G^2}{2}\sum_{t=1}^T \eta_t^2\label{leq:avg-2-mid3}\end{equation}

Following the approach in the "Preliminaries]" section of Making Model Training More Scientific (III): Endpoint Loss Convergence of SGD], we change the starting point to $k$ and substitute in $\boldsymbol{\varphi}=\boldsymbol{\theta}_k$. But this time, instead of assuming monotonicity of $\eta_t$ in order to divide both sides by $\eta_T$, we directly obtain

\begin{equation}\sum_{t=k}^T \eta_t\mathbb{E}[L(\boldsymbol{\theta}_t) - L(\boldsymbol{\theta}_k)] \leq \frac{G^2}{2} \sum_{t=k}^T \eta_t^2\label{leq:avg-2-mid4}\end{equation}

Substituting $w_t = \eta_t, q_t = \mathbb{E}[L(\boldsymbol{\theta}_t)] - L(\boldsymbol{\theta}^*)$ into identity $\eqref{eq:qt-g}$, we get

\begin{equation}\begin{aligned} \mathbb{E}[L(\boldsymbol{\theta}_T) - L(\boldsymbol{\theta}^*)] =&\, \frac{1}{\eta_{1:T}}\underbrace{\sum_{t=1}^T \eta_t \mathbb{E}[L(\boldsymbol{\theta}_t) - L(\boldsymbol{\theta}^*)]}_{\eqref{leq:avg-2-mid3}} + \sum_{k=1}^{T-1}\left(\frac{1}{\eta_{k+1:T}} - \frac{1}{\eta_{k:T}}\right)\underbrace{\sum_{t=k}^T \eta_t \mathbb{E}[L(\boldsymbol{\theta}_t) - L(\boldsymbol{\theta}_k)]}_{\eqref{leq:avg-2-mid4}} \\ \leq&\, \frac{\Vert\boldsymbol{\theta}_1 - \boldsymbol{\theta}^*\Vert^2}{2\eta_{1:T}} + \frac{G^2}{2\eta_{1:T}}\sum_{t=1}^T \eta_t^2 + \frac{G^2}{2}\sum_{k=1}^{T-1}\left(\frac{1}{\eta_{k+1:T}} - \frac{1}{\eta_{k:T}}\right)\sum_{t=k}^T \eta_t^2 \\ \end{aligned}\label{leq:avg-2-mid5}\end{equation}

For the second term, applying $\sum_{k=1}^{T-1}\sum_{t=k}^T = \sum_{t=1}^{T} \sum_{k=1}^{\min(t, T-1)}$:

\begin{equation}\begin{aligned} \sum_{k=1}^{T-1}\left(\frac{1}{\eta_{k+1:T}} - \frac{1}{\eta_{k:T}}\right)\sum_{t=k}^T \eta_t^2 =&\, \sum_{t=1}^T\eta_t^2\sum_{k=1}^{\min(t, T-1)}\left(\frac{1}{\eta_{k+1:T}} - \frac{1}{\eta_{k:T}}\right) = \sum_{t=1}^T\eta_t^2\left(\frac{1}{\eta_{\min(t+1, T):T}} - \frac{1}{\eta_{1:T}}\right) \end{aligned}\end{equation}

Substituting this into equation $\eqref{leq:avg-2-mid5}$ gives

\begin{equation}\mathbb{E}[L(\boldsymbol{\theta}_T) - L(\boldsymbol{\theta}^*)] \leq \frac{\Vert\boldsymbol{\theta}_1 - \boldsymbol{\theta}^*\Vert^2}{2\eta_{1:T}} + \frac{G^2}{2}\sum_{t=1}^T\frac{\eta_t^2}{\eta_{\min(t+1, T):T}}\label{leq:last-2}\end{equation}

This is the strengthened version of the endpoint loss convergence result. It does not depend on the learning rate being monotonically decreasing, nor does it depend on the operation of dividing both sides by $\eta_T$, which gives the learning rate schedule more flexibility. The previous post's result essentially just replaced $\eta_{1:T},\eta_{\min(t+1, T):T}$ with $T\eta_T,\max(1, T-t)\eta_T$ throughout, giving a noticeably cruder conclusion. This result is new, but it is essentially equivalent to Theorem 10 in Appendix F of the original paper.

Accelerating Convergence

In this section we will see that, under an appropriate setup, equation $\eqref{leq:last-2}$ can achieve a convergence rate of $\mathcal{O}(1/\sqrt{T})$. Here the "appropriate setup" mainly refers to the learning rate scheduling strategy. Unlike the constant learning rates or the "endpoint-free" schedules such as $\alpha/\sqrt{t}$ and $\alpha/t$ used previously, this time we choose "linear decay":

\begin{equation}\eta_t = \alpha\left(1 - \frac{t}{T+1}\right)\label{eq:linear-decay}\end{equation}

This learning rate function deserves to be highlighted on its own line, because it is one of the best practices for learning rate schedules — for instance, Straight to Zero: Why Linearly Decaying the Learning Rate to Zero Works Best for LLMs] claims that it even outperforms cosine decay. This shows that our discussion is getting closer and closer to practical scenarios.

Computing term by term directly, we get:

\begin{gather}\eta_{1:T} = \sum_{\tau=1}^T \alpha\left(1 - \frac{\tau}{T+1}\right) = \frac{\alpha T}{2} \\[6pt] \eta_{t+1:T} = \sum_{\tau=t+1}^T \alpha\left(1 - \frac{\tau}{T + 1}\right) = \frac{\alpha(T-t)(T+1-t)}{2(T+1)} \\[6pt] \frac{\eta_t^2}{\eta_{t+1:T}} = \frac{2\alpha (T+1-t)}{(T-t)(T+1)}\leq \frac{4\alpha}{T+1} \\[6pt] \sum_{t=1}^T\frac{\eta_t^2}{\eta_{\min(t+1, T):T}} = \eta_T + \sum_{t=1}^{T-1}\frac{\eta_t^2}{\eta_{t+1:T}} \leq \frac{\alpha}{T+1} + \sum_{t=1}^{T-1} \frac{4\alpha}{T+1} \leq 4\alpha\end{gather}

Substituting these results into equation $\eqref{leq:last-2}$ gives

\begin{equation}\mathbb{E}[L(\boldsymbol{\theta}_T) - L(\boldsymbol{\theta}^*)] \leq \frac{\Vert\boldsymbol{\theta}_1 - \boldsymbol{\theta}^*\Vert^2}{\alpha T} + 2 G^2 \alpha\end{equation}

Taking $\alpha=\frac{\Vert\boldsymbol{\theta}_1 - \boldsymbol{\theta}^*\Vert}{G\sqrt{2T}}$ minimizes the right-hand side, achieving an endpoint convergence rate of $\mathcal{O}(1/\sqrt{T})$. It should be pointed out that, without introducing stronger assumptions, $\mathcal{O}(1/\sqrt{T})$ cannot be improved further — this is guaranteed by information theory (see Information-theoretic lower bounds on the oracle complexity of stochastic convex optimization]). So this is already the theoretically optimal convergence rate.

The Variational Approach

Readers may wonder: how can we tell that linear decay achieves the optimal convergence rate? Or more generally: how do we find $\eta_1,\eta_2,\cdots,\eta_T\geq 0$ so that the right-hand side of inequality $\eqref{leq:last-2}$ is minimized? This is a multivariate minimization problem. The standard approach is to take partial derivatives, but the exact solution is complicated and hard to interpret intuitively.

Here we instead consider a continuous relaxation and give a quick heuristic derivation using the calculus of variations] (quick, of course, assuming the reader is already familiar with it). Let $S_t = \eta_{\min(t+1, T):T}$, so that for $t < T - 1$ we have $\eta_t = S_{t-1} - S_t\approx -\dot{S}_t$. We uniformly approximate $\eta_t$ with $-\dot{S}_t$ and approximate the sum with an integral, so that the original problem can be approximated by minimizing the following integral functional:

\begin{equation}\frac{c}{S_0} + \int_0^T \frac{(\dot{S}_t)^2}{S_t} dt \label{eq:func-min}\end{equation}

According to the definition $S_T=0$, let us first fix $S_0$; this then becomes a variational problem with fixed boundary conditions. Substituting into the Euler-Lagrange equation] gives

\begin{equation}\frac{d}{dt}\frac{2\dot{S}_t}{S_t} = -\frac{(\dot{S}_t)^2}{S_t^2}\end{equation}

which is easily solved to give $S_t = S_0(1-t/T)^2$. Substituting this into $\eqref{eq:func-min}$ gives $c/S_0 + 4S_0/T$, whose minimum is attained at $S_0=\sqrt{c T/4}$. So the optimal solution is $S_t = \sqrt{c T/4}(1-t/T)^2$; differentiating gives $\eta_t = -\dot{S}_t = \sqrt{c/T}(1-t/T)$, which recovers linear decay — including the constant factor $\propto 1/\sqrt{T}$ derived earlier.

This result is still only approximate; to get a rigorous proof we need to substitute it back into the original discrete formulation. There, $\eta_t\propto 1 - t/T$ runs into a division-by-zero issue, so a slight adjustment is needed, which ultimately yields equation $\eqref{eq:linear-decay}$.

Further Reflections

There are a few key points in the derivation and conclusions above that deserve special attention. In a sense, this represents a milestone advance in the theory of stochastic optimization convergence.

First, if the learning rate is set to a constant, the conclusion $\eqref{leq:last-2}$ agrees with that of the previous post: we already showed that it can achieve at most a convergence rate of $\mathcal{O}(\sqrt{\ln T/T})$, which is not optimal. By contrast, the linear-decay learning rate in equation $\eqref{eq:linear-decay}$ achieves $\mathcal{O}(1/\sqrt{T})$. This demonstrates, on one hand, the necessity of learning rate decay for endpoint convergence, and on the other hand, provides theoretical support for the linear decay strategy.

It is not hard to show that, in the first three posts, the best convergence rate was always achieved with a constant learning rate — but this constant depends on the total number of training steps $T$, e.g., $\alpha/\sqrt{T}$. Many works consider this a drawback, preferring schedules like $\alpha/\sqrt{t}$ or $\alpha/t$ instead, since these do not require knowing the training step count $T$ in advance — they belong to the category of "endpoint-free" learning rate strategies that let you "stop anytime, resume anytime, and train for as many steps as you like."

However, such strategies tend to perform poorly in practice. The new results in this post show that, once we switch to endpoint loss, some new characteristics emerge: the fastest-converging strategy is neither a constant learning rate tied to $T$, nor a dynamic learning rate independent of $T$, but rather "a bit of both" — as exemplified by linear decay. In addition, cosine-type decay is also commonly used in practice. What these strategies have in common is: better endpoint performance, dependence on $T$, and dynamic variation over training.

In other words, there is no "one-size-fits-all" learning rate strategy. Only by fine-tuning the learning rate schedule according to the number of training steps can we obtain the best endpoint convergence result. This actually aligns quite well with current Scaling Law practice — for example, Step Law] finds that both the optimal learning rate and the optimal batch size should be finely tuned according to the amount of data. Note that once the amount of data and the batch size are fixed, the number of training steps $T$ is also determined, so these quantities can equally be regarded as functions of the training step count $T$.

In a later post, we will also briefly discuss the connection between this series of results and today's Scaling Laws — stay tuned.

Summary

In this post, we generalized the core identity from the previous post, and thereby obtained the theoretically optimal convergence rate for endpoint loss. Interestingly, the learning rate strategy that achieves this result is neither a constant learning rate nor the traditional inverse-step or inverse-square-root-of-step schedules, but rather linear decay — something much closer to everyday practice. In the next post, we will continue exploring the deeper implications behind this result.