Making Model Training a Bit More Scientific (VIII): Learning Rates for Multi-Stage Training

让炼丹更科学一些（八）：多阶段训练的学习率 苏剑林 · 2026-08-31 · Mathematics inequality learning rate optimizer SGD model tuning

In the previous post Making Model Training a Bit More Scientific (VII): Step-Size Schedules and Weight Averaging we gave a brief introduction to Schedule-Free learning rates, which try to replace learning-rate schedules with a certain kind of weight averaging, so that training with a constant learning rate can still achieve the effect of an optimally scheduled run. However, as we also noted there, without introducing additional assumptions, even this "optimal" constant learning rate still depends on the total number of training steps, so true schedule-freeness cannot really be achieved.

In this post we revisit the problem from the perspective of multi-stage training. The core idea is to relax the scheduling objective to "being close to optimal at the end of each stage," which makes the problem much simpler and more practical to handle.

A Classic Result

Let's start from a classic convergence result, which will help us understand why the optimal learning rate depends on the number of training steps. We first introduced this convergence result in Making Model Training a Bit More Scientific (II): Extending the Result to Unbounded Domains:

\begin{equation}\frac{\sum_{t=1}^T \eta_t \mathbb{E}[L(\boldsymbol{\theta}_t) - L(\boldsymbol{\theta}^*)]}{\sum_{t=1}^T \eta_t}\leq \frac{R^2}{2\sum_{t=1}^T \eta_t} + \frac{G^2}{2}\frac{\sum_{t=1}^T \eta_t^2}{\sum_{t=1}^T \eta_t}\triangleq f_T(\eta)\end{equation}

We won't repeat the meaning of the notation here. This convergence result is representative, since the later posts Making Model Training a Bit More Scientific (VI): An Elegant Top-Down Construction and the previous Making Model Training a Bit More Scientific (VII): Step-Size Schedules and Weight Averaging are both built upon it.

What we want to do now is: given $R,G,T > 0$, find $\eta_1,\cdots,\eta_T$ that minimizes $f_T(\eta)$, which can be solved with a simple inequality. First it's easy to show that $\sum_{t=1}^T \eta_t^2\geq (\sum_{t=1}^T \eta_t)^2/T$, so we get

\begin{equation}f_T(\eta)\geq \frac{R^2}{2\sum_{t=1}^T \eta_t} + \frac{G^2}{2T}\sum_{t=1}^T \eta_t \geq \frac{RG}{\sqrt{T}}\triangleq f^*_T\end{equation}

Both equalities hold simultaneously when $\eta_1=\cdots=\eta_T=\frac{R}{G\sqrt{T}}$, which depends on $T$. If instead we use a learning rate schedule that is genuinely independent of $T$, such as $\eta_t \propto 1/\sqrt{t}$, then at best we can achieve $f_T(\eta)=\mathcal{O}(\log T / \sqrt{T})$, and one can further prove that no matter how we improve it, a learning rate that does not depend on $T$ can never achieve a $f$ better than $\mathcal{\Theta}(\sqrt{\log T / T})$.

Reframing the Problem

Now let's go back and think about the motivation behind Schedule-Free. Ideally, Schedule-Free hopes to keep training indefinitely, without needing to know the total number of training steps $T$ in advance, in such a way that no matter at which step we stop, we get model weights that are close to optimal up to that point.

However, this perfect goal looks nearly impossible. Even under the convex optimization assumptions behind Schedule-Free, where we can make LR decay equivalent to weight averaging and thus train with a constant learning rate, as analyzed in the previous section, the optimal value of this constant learning rate still depends on $T$. So we cannot obtain the up-to-date optimal solution no matter when we stop, unless we introduce further additional assumptions.

Given this, let's reconsider what we actually want from Schedule-Free. First, a constant learning rate is not really a hard requirement — if some time-varying learning rate can also train a good model, that's acceptable too. Second, "optimal no matter when you stop" is certainly a nice property, but it's not really a strict necessity either; what we mainly want, in multi-stage training, is for the endpoint of each stage to be as close to optimal as possible.

In other words, when our training is split into multiple stages, we want the model at the end of each stage to be as close to optimal as possible. The number of stages is usually not very large, so for simplicity let's first consider two-stage training, where the first stage trains for $T_1$ steps and the second stage trains for $T_2 - T_1$ steps, for a total of $T_2$ steps.

This can further be split into two scenarios — "unplanned" and "planned." "Unplanned" means there was originally no plan to continue training, so the first stage is tuned to be as optimal as possible on its own; once we decide to continue training, the second stage then tries to become optimal given what's already been done. This gives a greedy solution, whose downside is that no matter how we tune the second stage, the result is relatively suboptimal. As for the "planned" case, the multi-stage training schedule is planned in advance, and in theory we can better balance the outcomes of the two stages.

Greedy Decision-Making

Let's first look at the greedy version. Since the first stage is taken to be optimal, it naturally uses a constant learning rate $\eta_{(1)}^* = \frac{R}{G\sqrt{T_1}}$, and the $f$ of the first stage reaches its theoretical optimum $f^*_{T_1}$. In the second stage, building on this, we need to find $\eta_{T_1+1},\cdots,\eta_{T_2}$ that minimizes

\begin{equation}f_{T_2}(\eta) = \frac{R^2 + G^2(T_1 (\eta_{(1)}^*)^2 + \sum_{t=T_1+1}^{T_2} \eta_t^2)}{2(T_1 \eta_{(1)}^* + \sum_{t=T_1+1}^{T_2} \eta_t)}\end{equation}

Based on the same inequality $\sum_{t=T_1+1}^{T_2} \eta_t^2 \geq (\sum_{t=T_1+1}^{T_2} \eta_t)^2/(T_2-T_1)$, we know the minimum is still attained at a constant learning rate, so this simplifies to

\begin{equation}f_{T_2}(\eta) = \frac{R^2 + G^2(T_1 (\eta_{(1)}^*)^2 + (T_2 - T_1) (\eta_{(2)})^2)}{2(T_1 \eta_{(1)}^* + (T_2 - T_1) \eta_{(2)})}\end{equation}

which can be solved either by completing the inequality or by direct differentiation, giving the minimizer

\begin{equation}\eta_{(2)}^* = \frac{R/G}{\sqrt{T_2^{\#}}}, \qquad f(\eta_{(2)}^*) = \frac{RG}{\sqrt{T_2^{\#}}}\end{equation}

where

\begin{equation}T_2^{\#} = \frac{1}{2}T_2 + \frac{1}{2}\sqrt{T_1(2T_2 - T_1)}\end{equation}

Effective Step Count

Here we introduce the notation $T_2^{\#}$, which represents the "effective step count."

If there were no constraint from the first stage, we could use a constant learning rate $\frac{R}{G\sqrt{T_2}}$ from the very beginning, and the minimum achieved at the end of training would be $\frac{RG}{\sqrt{T_2}}$. But now the best we can do is $\frac{RG}{\sqrt{T_2^{\#}}}$, and the learning rate in the second stage also takes the form of $\frac{R}{G\sqrt{T_2^{\#}}}$. This is equivalent to saying that the effective number of steps has changed from $T_2$ to $T_2^{\#}$, and one can verify that

\begin{equation}T_2^{\#} \leq T_2\end{equation}

In other words, greedy two-stage training effectively "loses" some steps compared to the ideal case. Taking $T_2=2T_1$ as an example, $T_2^{\#}\approx 1.866T_1$, compared to $2T_1$, roughly loses $6.7\%$ of the steps. Note that if $T_2\to\infty$, then $T_2^{\#}/T_2\to 1/2$, meaning that if we keep continuing training indefinitely, the fraction of lost steps keeps increasing, up to a maximum loss of half.

Why bother abstracting out this concept? Because the notion of "number of steps" is universal, and it makes the conclusion easier to transfer to practice. The result above is derived for SGD, but real training typically uses non-SGD optimizers such as Adam or Muon, and hyperparameters like the learning rate are usually not set directly via the formulas above, but rather determined based on a fitted Scaling Law.

So the various conclusions derived for SGD cannot be directly applied. But if we can extract the relative change via the concept of "effective step count," then the step-count input used in the Scaling Law in practice can similarly be replaced with the effective step count, giving us a heuristic correction. This is the main bridge connecting the theoretical results of this post to practice.

The Planned Case

Now let's turn to the "planned" case. Here, $T_1,T_2$ are all known in advance, and the learning rates of the two stages can be jointly tuned. The core issue is that we have two delivery targets — the model at the end of stage one and the final model — and we need to balance the outcomes of these two stages. To this end, we introduce a minimax optimization objective:

\begin{equation}\min_{\eta\geq 0} \max_{T\in\{T_1,T_2\}} \frac{f_T(\eta)}{f_T^*}\end{equation}

This aims to make the ratio between the actual result of each stage and its ideal value as small as possible, for every stage. The advantage of this objective is that it introduces no additional hyperparameters, while also generalizing easily to multi-stage training. Introducing a new variable $m$, we can convert this into a nonlinear programming problem:

\begin{equation}\min\Big\{ m \,\,\Big|\,\, f_{T_1}(\eta)/f_{T_1}^*\leq m ,\,\, f_{T_2}(\eta)/f_{T_2}^*\leq m ,\,\, m\geq 1,\,\, \eta\geq 0\Big\}\end{equation}

To make the problem dimensionless, let $x_t = G\eta_t/R$; then for any $T$ we can complete the square to get

\begin{equation}\frac{f_T(\eta)}{f_T^*} = \frac{\sqrt{T}(1 + \sum_{t=1}^T x_t^2)}{2\sum_{t=1}^T x_t}\leq m \iff \sum_{t=1}^T\left(x_t - \frac{m}{\sqrt{T}}\right)^2 \leq m^2 - 1\end{equation}

That is, the feasible set of $x_t$ forms a hypersphere centered at $\frac{m}{\sqrt{T}}\boldsymbol{1}$ with radius $\sqrt{m^2-1}$, from which we see that the original problem is equivalent to a quadratic programming problem with objective $m$.

The Solution Procedure

For the two-stage case, this problem can be solved analytically. Note that the first constraint only involves $x_1,\cdots,x_{T_1}$, while the second constraint involves all $T_2$ variables. For the $x_t$ of index $t > T_1$, they only appear in the second constraint, so we can directly set $x_t = m/\sqrt{T_2}$ so that their contribution vanishes. The problem thus reduces to: the first $T_1$ $x_t$ must simultaneously lie in two spheres

\begin{equation}\sum_{t=1}^{T_1}\left(x_t - \frac{m}{\sqrt{T_1}}\right)^2 \leq m^2 - 1, \qquad \sum_{t=1}^{T_1}\left(x_t - \frac{m}{\sqrt{T_2}}\right)^2 \leq m^2 - 1\end{equation}

Both spheres have radius $\sqrt{m^2-1}$, and both centers lie along the all-ones vector direction, with the squared distance between the centers being

\begin{equation}T_1\left(\frac{m}{\sqrt{T_1}} - \frac{m}{\sqrt{T_2}}\right)^2 = m^2 (1 - \sqrt{\tau} )^2,\qquad \tau = T_1/T_2\end{equation}

The two spheres intersect if and only if the distance between their centers does not exceed twice the radius, i.e., $m^2 (1 - \sqrt{\tau} )^2 \leq 4(m^2 - 1)$, from which we can solve for the optimal $m$:

\begin{equation}m^* = \frac{2}{\sqrt{4 - (1 - \sqrt{\tau})^2}}\end{equation}

The corresponding solution can be taken as the midpoint of the line connecting the two sphere centers, i.e.,

\begin{equation}x_t = \left\{\begin{aligned} &\frac{m^*}{2}\left(\frac{1}{\sqrt{T_1}} + \frac{1}{\sqrt{T_2}}\right), & t \leq T_1 \\ &\frac{m^*}{\sqrt{T_2}}, & t > T_1 \end{aligned}\right.\end{equation}

Analyzing the Result

Restoring $\eta_t = Rx_t/G$ gives

\begin{equation}\eta_{(1)}^* = \frac{m^* R}{2G}\left(\frac{1}{\sqrt{T_1}} + \frac{1}{\sqrt{T_2}}\right),\qquad \eta_{(2)}^* = \frac{m^* R}{G\sqrt{T_2}}\end{equation}

Interestingly, $\eta_{(1)}^*$ is exactly the average (times $m^*$) of the two single-stage optimal constant learning rates for "training only $T_1$ steps" and "training $T_2$ steps from scratch." From the perspective of effective step count, $\eta_{(1)}^*$ corresponds to presupposing a larger number of training steps, while $\eta_{(2)}^*$ corresponds to presupposing a smaller number of training steps:

\begin{equation}T_1^{\#} = \left(\frac{R}{G\eta_{(1)}^*}\right)^2 = T_1\left(\frac{3-\sqrt{\tau}}{1 + \sqrt{\tau}}\right),\qquad T_2^{\#} = \left(\frac{R}{G\eta_{(2)}^*}\right)^2 = T_2\left(1-\frac{(1-\sqrt{\tau})^2}{4}\right)\end{equation}

In other words, the first stage only needs to train for $T_1$ steps, but when setting its learning rate, we assume it will train for $T_1^{\#}$ steps; the two stages together train for $T_2$ steps, but when setting the learning rate for the second stage, we assume it will train for $T_2^{\#}$ steps. As for the result, at this point $f_{T_1}/f_{T_1}^* = f_{T_2}/f_{T_2}^* = m^*$, i.e., the relative losses of the two stages are exactly equal — this is precisely the balancing property of the minimax objective: it does not favor either delivery point.

Again taking $T_2=2T_1$ as an example, here $\tau=1/2$. Substituting gives $T_2^{\#} \approx 0.979 T_2$, so in terms of performance, both stages, compared to the optimal solution, effectively lose about $2.1\%$ of the steps — and averaged over the two stages, this is better than the greedy solution. As for the learning rate settings, the first stage's learning rate should be set according to an effective step count of $T_1^{\#} \approx 1.343 T_1$, and the second stage's learning rate according to an effective step count of $T_2^{\#} \approx 0.979 T_2$.

Another limiting case is when $\tau\to 0$, giving $T_2^{\#}\to 0.75 T_2$ — that is, when the number of steps in stage $T_2$ is sufficiently large, the minimax solution loses at most $25\%$ of the steps at the endpoint, which is better than the greedy solution's $50\%$.

Generalizing to Multiple Stages

Both the greedy solution for the "unplanned" case and the minimax solution for the "planned" case can, in principle, be readily extended to multi-stage training. The greedy solution needs no further elaboration; let's focus mainly on the minimax solution. Suppose there are $K$ delivery points $T_1 < T_2 < \cdots < T_K$; the minimax objective becomes

\begin{equation}\min_{\eta\geq 0}\ \max_{T\in\{T_1,\cdots,T_K\}} \frac{f_T(\eta)}{f_T^*}\end{equation}

The same dimensionless treatment gives $K$ sphere constraints:

\begin{equation}\sum_{t=1}^{T_k}\left(x_t - \frac{m}{\sqrt{T_k}}\right)^2 \leq m^2 - 1,\qquad k = 1,2,\cdots,K\end{equation}

Although there are in principle $T_K$ variables, one can still prove that the optimal solution can always be taken to have a "piecewise constant" form, so the problem reduces to a small-scale program with only $K+1$ variables. When $K\geq 3$, there is generally no simple closed-form solution, but numerical solving is not difficult.

Summary

This post revisits "schedule-free learning rates" from the perspective of multi-stage training. By changing the scheduling objective to "being close to optimal at the end of each stage," we make the problem much simpler and more practical to work with. We further extract the notion of "effective step count," making it possible for these conclusions to transfer to the optimizers and Scaling Laws actually used in practice.