Making Model Training More Scientific (Part 3): Convergence of the Final-Point Loss in SGD

让炼丹更科学一些（三）：SGD的终点损失收敛 苏剑林 · 2025-12-16 · Mathematics inequality optimizer SGD model tuning

So far we have had two posts discussing the convergence properties of SGD, but both only gave convergence results for the loss value, so they only guarantee that we can find the optimal loss value, without guaranteeing that we find the location of the optimum $\boldsymbol{\theta}^*$ — this is a notable gap between the current theoretical conclusions and practice. Intuitively, the weights $\boldsymbol{\theta}_T$ at the end of training should be close to the theoretically optimal $\boldsymbol{\theta}^*$, and we would like to know whether this is theoretically supported.

So in this post we will convert the convergence result for the average loss into a convergence result for the final-point loss, giving us a preliminary theoretical sense of how far apart $\boldsymbol{\theta}_T$ and $\boldsymbol{\theta}^*$ actually are.

Locating the Optimum

We start from the post Making Model Training More Scientific (Part 2): Extending the Conclusion to Unbounded Domains, whose core result is the inequality

\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}more

Then, assuming the monotonic decrease of $\eta_t$, we replace the left-hand $\eta_t$ with $\eta_T$, and substituting $\boldsymbol{\varphi}=\boldsymbol{\theta}^*$ gives us one of the conclusions from that previous post:

\begin{equation}\frac{1}{T}\sum_{t=1}^T \mathbb{E}[L(\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}\label{leq:avg-2}\end{equation}

As mentioned at the outset, this is only a convergence result for the loss value; what we really want is to locate where it converges. To this end, a fairly simple idea is to exploit the convexity of $L$ together with Jensen's inequality, which gives

\begin{equation}\frac{1}{T}\sum_{t=1}^T \mathbb{E}[L(\boldsymbol{\theta}_t)] = \mathbb{E}\left[\frac{1}{T}\sum_{t=1}^T L(\boldsymbol{\theta}_t)\right] \geq \mathbb{E}\left[L\left(\frac{1}{T}\sum_{t=1}^T \boldsymbol{\theta}_t\right)\right]\end{equation}

Defining $\bar{\boldsymbol{\theta}}_T = \frac{1}{T}\sum_{t=1}^T \boldsymbol{\theta}_t$, we then have

\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}

That is, the loss value $L(\bar{\boldsymbol{\theta}}_T)$ corresponding to the centroid $\bar{\boldsymbol{\theta}}_T$ of the training trajectory $\boldsymbol{\theta}_1,\boldsymbol{\theta}_2,\cdots,\boldsymbol{\theta}_T$ converges on average to $L(\boldsymbol{\theta}^*)$, which in turn means that $\bar{\boldsymbol{\theta}}_T$ converges on average to $\boldsymbol{\theta}^*$ (since a strictly convex function has a unique minimizer). This to some extent explains the practice of taking a moving average over the training trajectory to obtain better weights, including the rationale behind the merge operation in WSM (Warmup-Stable and Merge).

Preliminaries

Computing $\bar{\boldsymbol{\theta}}_T$ gives us one way to locate $\boldsymbol{\theta}^*$, but it doesn't fully answer the question posed at the start of this post — what we really want is the conclusion that $\boldsymbol{\theta}_T$ converges to $\boldsymbol{\theta}^*$. Next we follow the approach of Last Iterate of SGD Converges (Even in Unbounded Domains) to convert convergence of the average loss into convergence of the final-point loss.

Before giving the formal proof, we need some preliminary work, one part of which is generalizing equation $\eqref{leq:avg-2-mid3}$. From its proof, we know that the lower limit of the summation can in principle be arbitrary — that is, we can replace the starting point $\boldsymbol{\theta}_1$ with an arbitrary $\boldsymbol{\theta}_{T-k}$, and the inequality still holds, though now $\boldsymbol{\theta}_{T-k}$ may also depend on $\boldsymbol{x}_t$, so we need to add $\mathbb{E}$ to the right-hand side as well, giving

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

Again assuming the monotonic decrease of $\eta_t$, we replace the left-hand $\eta_t$ with $\eta_T$, and divide both sides by $\eta_T$, obtaining

\begin{equation}\sum_{t=T-k}^T \mathbb{E}[L(\boldsymbol{\theta}_t) - L(\boldsymbol{\varphi})] \leq \frac{\mathbb{E}[\Vert\boldsymbol{\theta}_{T-k} - \boldsymbol{\varphi}\Vert^2]}{2\eta_T} + \frac{G^2}{2} \sum_{t=T-k}^T\frac{\eta_t^2}{\eta_T}\end{equation}

Here $\boldsymbol{\varphi}$ is an arbitrary data-independent vector, though this "data independence" is relative: reviewing the proof, we find that when we choose the starting point to be $T-k$, it may depend on $\boldsymbol{x}_1,\boldsymbol{x}_2,\cdots,\boldsymbol{x}_{T-k-1}$ at most. In particular, $\boldsymbol{\theta}_{T-k}$ satisfies this condition, and substituting $\boldsymbol{\varphi}=\boldsymbol{\theta}_{T-k}$ gives

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

This is an important intermediate result that we will use later.

A Key Identity

To convert the conclusion about the average loss into a statement about the final point, we need one more crucial 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})\label{eq:qt}\end{equation}

This identity elegantly connects the final-point value to the average value. I spent several days trying to find an intuitive way to understand it, without success, so I will simply walk through its proof step by step. The idea of the proof is to consider the cumulative average of $q_t$ counted from the end back to the beginning — that is, define $S_k = \frac{1}{k}\sum_{t=T-k+1}^T q_t$ — and then we can write

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

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

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

Finally, substituting in the original definition of $S_1,S_T$ gives us equation $\eqref{eq:qt}$. The core of the whole derivation is that the operation of "cumulative averaging" serves as a natural bridge from the final-point value $q_T$ to the average value $\frac{1}{T}\sum_{t=1}^T q_t$. In the original blog post, this identity appears in a slightly different inequality form, but I believe the identity form is more fundamental, and the subsequent proof only needs the equality form anyway.

Completing the Proof

Now we can push through and complete the proof in one go. Define $q_t = \mathbb{E}[L(\boldsymbol{\theta}_t) - L(\boldsymbol{\theta}^*)]$, and substitute it into the identity $\eqref{eq:qt}$ to get:

\begin{equation}\mathbb{E}[L(\boldsymbol{\theta}_T) - L(\boldsymbol{\theta}^*)] = \underbrace{\frac{1}{T}\sum_{t=1}^T \mathbb{E}[L(\boldsymbol{\theta}_t) - L(\boldsymbol{\theta}^*)]}_{\eqref{leq:avg-2}} + \sum_{k=1}^{T-1} \frac{1}{k(k+1)}\underbrace{\sum_{t=T-k}^T \mathbb{E}[L(\boldsymbol{\theta}_t) - L(\boldsymbol{\theta}_{T-k})]}_{\eqref{leq:last-mid2}}\end{equation}

Substituting inequalities $\eqref{leq:avg-2}$ and $\eqref{leq:last-mid2}$ respectively gives

\begin{equation}\mathbb{E}[L(\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} + \frac{G^2}{2}\sum_{k=1}^{T-1} \frac{1}{k(k+1)}\sum_{t=T-k}^T \frac{\eta_t^2}{\eta_T}\label{leq:last-mid3}\end{equation}

For the last term, we have

\begin{equation}\sum_{k=1}^{T-1}\frac{1}{k(k+1)}\sum_{t=T-k}^{T}\frac{\eta_t^2}{\eta_T} =\sum_{t=1}^{T}\frac{\eta_t^2}{\eta_T}\sum_{k=\max(1,\,T-t)}^{T-1} \frac{1}{k(k+1)} =\sum_{t=1}^{T}\frac{\eta_t^2}{\eta_T}\left(\frac{1}{\max(1,\,T-t)}-\frac{1}{T}\right)\label{eq:last-mid4}\end{equation}

Substituting this into equation $\eqref{leq:last-mid3}$ gives

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

This is our final conclusion. Because we used the identity transformation $\eqref{eq:last-mid4}$ to swap the order of summation and simplify things in advance, this result ends up being more concise and more general than the one in the original post Last Iterate of SGD Converges (Even in Unbounded Domains).

Two Examples

It is not hard to see that the shape of the right-hand side of conclusion $\eqref{leq:last-1}$ is quite similar to that of $\eqref{leq:avg-2}$, which suggests that the final-point loss and the average loss should converge at similar rates. As before, we look at two examples — static learning rate and dynamic learning rate — to see how the final conclusion behaves. First, with a static learning rate $\eta_t = \eta$, we get

\begin{equation}\begin{aligned} \mathbb{E}[L(\boldsymbol{\theta}_T) - L(\boldsymbol{\theta}^*)] \leq&\, \frac{\Vert\boldsymbol{\theta}_1 - \boldsymbol{\theta}^*\Vert^2}{2T\eta} + \frac{G^2\eta}{2} + \frac{G^2\eta}{2}\sum_{t=1}^{T-1}\frac{1}{T-t} \\ \leq&\, \frac{\Vert\boldsymbol{\theta}_1 - \boldsymbol{\theta}^*\Vert^2}{2T\eta} + \frac{G^2\eta}{2} (2 + \ln T) \end{aligned}\end{equation}

Taking $\eta = \frac{\Vert\boldsymbol{\theta}_1 - \boldsymbol{\theta}^*\Vert/G}{\sqrt{T(2+\ln T)}}$ minimizes the rightmost expression, giving a convergence rate of $\mathcal{O}(\sqrt{\ln T/T})$, which is slightly slower than that of the average loss. In the previous post we showed that, under a constant learning rate, the average loss can achieve a convergence rate of $\mathcal{O}(1/\sqrt{T})$. Of course, this difference only shows up in the limiting case; in practice the difference between $\sqrt{\ln T}$ may be entirely negligible.

Next, consider the dynamic learning rate $\eta_t = \frac{\alpha}{\sqrt{t}}$. Substituting into equation $\eqref{leq:last-1}$ gives

\begin{equation}\begin{aligned} \mathbb{E}[L(\boldsymbol{\theta}_T) - L(\boldsymbol{\theta}^*)] \leq&\, \frac{\Vert\boldsymbol{\theta}_1 - \boldsymbol{\theta}^*\Vert^2}{2\alpha\sqrt{T}} + \frac{G^2\alpha}{2\sqrt{T}} + \frac{G^2\alpha\sqrt{T}}{2}\sum_{t=1}^{T-1}\frac{1}{t(T-t)} \\ =&\, \frac{\Vert\boldsymbol{\theta}_1 - \boldsymbol{\theta}^*\Vert^2}{2\alpha\sqrt{T}} + \frac{G^2\alpha}{2\sqrt{T}} + \frac{G^2\alpha}{2\sqrt{T}}\sum_{t=1}^{T-1}\left(\frac{1}{t} + \frac{1}{T-t}\right) \\ =&\, \frac{\Vert\boldsymbol{\theta}_1 - \boldsymbol{\theta}^*\Vert^2}{2\alpha\sqrt{T}} + \frac{G^2\alpha}{2\sqrt{T}} + \frac{G^2\alpha}{\sqrt{T}}\sum_{t=1}^{T-1}\frac{1}{t} \\ \leq&\, \frac{\Vert\boldsymbol{\theta}_1 - \boldsymbol{\theta}^*\Vert^2}{2\alpha\sqrt{T}} + \frac{G^2\alpha}{2\sqrt{T}} + \frac{G^2\alpha}{\sqrt{T}}(1 + \ln T) \\ \sim&\, \mathcal{O}\left(\frac{\ln T}{\sqrt{T}}\right) \end{aligned}\end{equation}

Just as with the convergence of the average loss over unbounded domains in the previous post, under the dynamic learning rate $\eta_t = \frac{\alpha}{\sqrt{t}}$ the convergence rate is again $\mathcal{O}(\ln T / \sqrt{T})$, only with a somewhat larger constant here.

Summary

In this post, we extended the convergence result for SGD from the average loss to the final-point loss — that is, we considered how close the loss value at the end of training is to the theoretical optimum. This setting is a closer match to actual training practice.