Making Alchemy a Bit More Scientific (Part 2): Extending the Conclusion to Unbounded Domains

Two years ago, I planned to start a "Scientific Alchemy" series, intending to systematically organize the classic theoretical results on optimizers. But after writing the first post, Making Alchemy a Bit More Scientific (Part 1): Convergence of the Average Loss of SGD, I let the series sit untouched until now. The main reason is that I've always felt the conditions these classical optimization results depend on are too restrictive, quite far removed from practical applications—especially now that we're in the LLM era, where the reference value of these results seems even more limited. So I never had much motivation to keep writing.

Recently, however, while thinking through some questions related to Scaling Laws, I found that these results aren't as "useless" as I'd imagined—they can actually provide useful theoretical insight into some empirical findings. So I've decided to restart this series and continue writing this line of posts, "paying off" the debt I owe from before.

Recap of the Conclusion

We'll keep using the same notation as in the first post, so there's no need to reintroduce it. The main conclusion of the first post was: under suitable assumptions, SGD satisfies

\begin{equation}\frac{1}{T}\sum_{t=1}^T L(\boldsymbol{x}_t,\boldsymbol{\theta}_t) - \frac{1}{T}\sum_{t=1}^T L(\boldsymbol{x}_t,\boldsymbol{\theta}^*)\leq \frac{R^2}{2T\eta_T} + \frac{G^2}{2T}\sum_{t=1}^T\eta_t\label{leq:avg-1}\end{equation}

where $R, G$ is a constant independent of the optimization trajectory, and the "suitable assumptions" include:

1. $\boldsymbol{\Theta}$ is a bounded convex set, $R=\max\limits_{\boldsymbol{\theta}_1,\boldsymbol{\theta}_2\in \boldsymbol{\Theta}}\Vert\boldsymbol{\theta}_1-\boldsymbol{\theta}_2\Vert < \infty$;
2. for any $\boldsymbol{\theta}\in \boldsymbol{\Theta}$ and any $\boldsymbol{x}$, $L(\boldsymbol{x},\boldsymbol{\theta})$ is a convex function of $\boldsymbol{\theta}$;
3. for any $\boldsymbol{\theta}\in \boldsymbol{\Theta}$ and any $\boldsymbol{x}$, we have $\Vert\nabla_{\boldsymbol{\theta}}L(\boldsymbol{x},\boldsymbol{\theta})\Vert\leq G < \infty$;
4. the learning rate $\eta_t$ is a monotonically decreasing function of $t$ (i.e., $\eta_t\geq \eta_{t+1}$);

The most "awkward" part here is probably the "bounded" in "bounded convex set" in the first condition, since the optimization trajectory of SGD generally cannot be guaranteed to be bounded. To ensure boundedness, we had to modify SGD into projected SGD:

\begin{equation}\boldsymbol{\theta}_{t+1} = \Pi_{\boldsymbol{\Theta}}\big(\boldsymbol{\theta}_t - \eta_t \boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\big)\in \boldsymbol{\Theta}\label{eq:sgd-p}\end{equation}

Of course, boundedness poses no problem at all in practice—an appropriate degree of boundedness can even enhance the stability of the optimization algorithm. But from the standpoint of "theoretical purism," adding an extra constraint is never entirely comfortable. So in this post, we'll first remove this boundedness condition, in order to obtain a simpler, more illuminating proof.

A New Approach

By the way, the proofs in this section and the sections that follow mainly draw on the article Last Iterate of SGD Converges (Even in Unbounded Domains) from the blog Parameter-free. This is a classic blog on optimization theory; besides this particular post, it contains a great deal of theoretical material on optimizers, much of it original to the author, and is well worth reading.

Just as in the proof of the previous post, our starting point is the identity

\begin{equation}\begin{aligned} \Vert\boldsymbol{\theta}_{t+1} - \boldsymbol{\varphi}\Vert^2=&\, \Vert\boldsymbol{\theta}_t - \eta_t \boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)- \boldsymbol{\varphi}\Vert^2 \\ =&\, \Vert\boldsymbol{\theta}_t - \boldsymbol{\varphi}\Vert^2 - 2\eta_t (\boldsymbol{\theta}_t- \boldsymbol{\varphi})\cdot\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t) + \eta_t^2\Vert\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\Vert^2 \end{aligned}\end{equation}

Here $\boldsymbol{\varphi}$ is an arbitrary vector; the previous post directly set $\boldsymbol{\varphi}=\boldsymbol{\theta}^*$, but here we keep open the possibility that it takes any value. After establishing this identity, the earlier post divided both sides by $2\eta_t$ and then summed over $t$, which meant that the later bounding steps required the boundedness condition to be tractable. In this post, instead of dividing by $2\eta_t$, we sum both sides directly over $t$, so that the $\Vert\boldsymbol{\theta}_t - \boldsymbol{\varphi}\Vert^2$ term in the middle cancels out naturally:

\begin{equation}\begin{aligned} 2\sum_{t=1}^T \eta_t (\boldsymbol{\theta}_t- \boldsymbol{\varphi})\cdot\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t) =&\, \Vert\boldsymbol{\theta}_1 - \boldsymbol{\varphi}\Vert^2 - \Vert\boldsymbol{\theta}_{T+1} - \boldsymbol{\varphi}\Vert^2 + \sum_{t=1}^T \eta_t^2 \Vert\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\Vert^2 \\ \leq&\, \Vert\boldsymbol{\theta}_1 - \boldsymbol{\varphi}\Vert^2 + G^2\sum_{t=1}^T \eta_t^2 \end{aligned}\end{equation}

Using the convexity of $L$, we know $(\boldsymbol{\theta}_t- \boldsymbol{\varphi})\cdot\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t) \geq L(\boldsymbol{x}_t,\boldsymbol{\theta}_t) - L(\boldsymbol{x}_t,\boldsymbol{\varphi})$; substituting this into the left-hand side above and rearranging gives

\begin{equation}\sum_{t=1}^T \eta_t [L(\boldsymbol{x}_t,\boldsymbol{\theta}_t) - L(\boldsymbol{x}_t,\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-mid1}\end{equation}

This is already very close to the result we want, and crucially the bounding steps did not require boundedness at any point—so we only need an assumption on the distance between the initialization $\boldsymbol{\theta}_1$ and the endpoint $\boldsymbol{\varphi}$.

Adding an Expectation

Next, if we could further assume the non-negativity of $L(\boldsymbol{x}_t,\boldsymbol{\theta}_t) - L(\boldsymbol{x}_t,\boldsymbol{\varphi})$, then combined with the monotonic decrease of $\eta_t$ we would get $\eta_t [L(\boldsymbol{x}_t,\boldsymbol{\theta}_t) - L(\boldsymbol{x}_t,\boldsymbol{\varphi})] \geq \eta_T [L(\boldsymbol{x}_t,\boldsymbol{\theta}_t) - L(\boldsymbol{x}_t,\boldsymbol{\varphi})]$, and substituting this into the equation above would give

\begin{equation}\frac{1}{T}\sum_{t=1}^T L(\boldsymbol{x}_t,\boldsymbol{\theta}_t) - \frac{1}{T}\sum_{t=1}^T L(\boldsymbol{x}_t,\boldsymbol{\varphi})\leq \frac{\Vert\boldsymbol{\theta}_1 - \boldsymbol{\varphi}\Vert^2}{2 T \eta_T} + \frac{G^2}{2T}\sum_{t=1}^T \frac{\eta_t^2}{\eta_T}\end{equation}

The problem is that even if we set $\boldsymbol{\varphi}$ to be the global optimum $\boldsymbol{\theta}^*$, there's no guarantee that $L(\boldsymbol{x}_t,\boldsymbol{\theta}_t) - L(\boldsymbol{x}_t,\boldsymbol{\varphi})$ is non-negative, because the optimal point is optimal for the whole dataset but not necessarily for any particular subset of samples. To sidestep this issue, notice that the right-hand side of $\eqref{leq:avg-2-mid1}$ is independent of $\boldsymbol{x}_t$, while the left-hand side depends on $\boldsymbol{x}_t$. So if we take the expectation of both sides, the inequality still holds:

\begin{equation}\sum_{t=1}^T \eta_t \mathbb{E}[L(\boldsymbol{x}_t,\boldsymbol{\theta}_t) - L(\boldsymbol{x}_t,\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-mid2}\end{equation}

Here $\mathbb{E}$ denotes $\mathbb{E}_{\boldsymbol{x}_1,\boldsymbol{x}_2,\cdots,\boldsymbol{x}_T}$, i.e., "repeating SGD infinitely many times with different sampling orders, and averaging over each optimization trajectory." Now, since $\boldsymbol{\varphi}$ is an arbitrary vector independent of the data, we have $\mathbb{E}[L(\boldsymbol{x}_t,\boldsymbol{\varphi})] = \mathbb{E}_{\boldsymbol{x}_t}[L(\boldsymbol{x}_t,\boldsymbol{\varphi})] = L(\boldsymbol{\varphi})$, which is exactly our optimization objective. As for $L(\boldsymbol{x}_t,\boldsymbol{\theta}_t)$, note that $\boldsymbol{\theta}_t$ depends on $\boldsymbol{x}_1,\cdots,\boldsymbol{x}_{t-1}$, so at best we have

\begin{equation}\mathbb{E}[L(\boldsymbol{x}_t,\boldsymbol{\theta}_t)] = \mathbb{E}_{\boldsymbol{x}_1,\cdots,\boldsymbol{x}_t}[L(\boldsymbol{x}_t,\boldsymbol{\theta}_t)] = \mathbb{E}_{\boldsymbol{x}_1,\cdots,\boldsymbol{x}_{t-1}}[\mathbb{E}_{\boldsymbol{x}_t}[L(\boldsymbol{x}_t,\boldsymbol{\theta}_t)]] = \mathbb{E}_{\boldsymbol{x}_1,\cdots,\boldsymbol{x}_{t-1}}[L(\boldsymbol{\theta}_t)] = \mathbb{E}[L(\boldsymbol{\theta}_t)]\end{equation}

Given that $\boldsymbol{\varphi}$ is independent of the data, $L(\boldsymbol{\varphi})$ can also be written as $\mathbb{E}[L(\boldsymbol{\varphi})]$. Substituting these into $\eqref{leq:avg-2-mid2}$ gives

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

A Digression on the Reasoning

Before pushing the derivation further, I want to say a bit more: why do we need to, and why are we allowed to, introduce this kind of expectation?

The "why do we need to" part has a simple answer—so that the derivation can go through. For instance, we want the non-negativity of $L(\boldsymbol{x}_t,\boldsymbol{\theta}_t) - L(\boldsymbol{x}_t,\boldsymbol{\varphi})$, but this can't be guaranteed; once we add the expectation, we can instead obtain the non-negativity of $\mathbb{E}[L(\boldsymbol{\theta}_t) - L(\boldsymbol{\varphi})]$, as long as $\boldsymbol{\varphi}$ is taken to be the theoretical optimum. Broadly speaking, taking the expectation removes the dependence on the data variable $\boldsymbol{x}_t$, simplifying the subsequent derivation—that's the main point of this step.

As for "why are we allowed to"—some readers might object: we're not actually going to retrain many times and average the results, so why should we care about such an expected outcome? To answer this, we first need to ask ourselves: what do we fundamentally want here? In fact, all we want is some degree of convergence guarantee for SGD; and since the objective function is restricted to being convex, we'd like this guarantee to be mathematically rigorous.

For a complicated stochastic process, the expectation is often the most basic quantity we're able to compute—sometimes the only one we're able to compute at all. So although taking the expectation on both sides pulls the conclusion somewhat away from practice, it still remains a convergence result with genuine reference value. As we all know, practice is generally more complicated than theory, so as long as a theoretical result maintains some connection to practice, can describe certain phenomena, or offer some insight, it still has value.

A Monotonic Bound

Back to the main thread. Now, given $\eqref{leq:avg-2-mid3}$, if we set $\boldsymbol{\varphi}$ to be the theoretical optimum $\boldsymbol{\theta}^*$, we obtain the non-negativity of $\mathbb{E}[L(\boldsymbol{x}_t,\boldsymbol{\theta}_t) - L(\boldsymbol{x}_t,\boldsymbol{\theta}^*)]$. Continuing to assume the monotonic decrease of $\eta_t$, we can then replace $\eta_t$ with $\eta_T$, and finally, dividing both sides by $T\eta_T$, we get

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

Comparing this with $\eqref{leq:avg-1}$, the dependence on $T$ and $\eta_t$ is similar in both: the first term on the right-hand side, $R^2$, has been replaced with $\Vert\boldsymbol{\theta}_1 - \boldsymbol{\theta}^*\Vert^2$, which is relatively speaking a shrinking; while in the summation in the second term on the right, $\eta_t$ has been replaced with $\eta_t^2/\eta_T$, which typically enlarges it. With one term shrinking and the other growing, how big is the actual difference? Let's look at two examples. The first is a constant learning rate $\frac{\alpha}{\sqrt{T}}$, in which case the conclusions $\eqref{leq:avg-1}$ and $\eqref{leq:avg-2}$ are essentially the same; taking $\eqref{leq:avg-2}$ as an example and substituting in, we get

\begin{equation}\frac{1}{T}\sum_{t=1}^T \mathbb{E}[L(\boldsymbol{\theta}_t) - L(\boldsymbol{\theta}^*)] \leq \frac{1}{2\sqrt{T}}\left(\frac{\Vert\boldsymbol{\theta}_1 - \boldsymbol{\theta}^*\Vert^2}{\alpha} + G^2\alpha\right)\end{equation}

which means both achieve a convergence rate of $\mathcal{O}(1/\sqrt{T})$. The second example is the dynamic learning rate $\eta_t = \frac{\alpha}{\sqrt{t}}$, which we already proved in the previous post; substituting it into $\eqref{leq:avg-1}$ also yields a convergence rate of $\mathcal{O}(1/\sqrt{T})$, but if we substitute it into $\eqref{leq:avg-2}$, the result is

\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}{2 \alpha \sqrt{T}} + \frac{G^2\alpha}{2\sqrt{T}}\sum_{t=1}^T \frac{1}{t} \sim \mathcal{O}\left(\frac{\ln T}{\sqrt{T}}\right)\end{equation}

So the price of extending to the unbounded domain is that the convergence rate is slightly enlarged from $\mathcal{O}(1/\sqrt{T})$ to $\mathcal{O}(\ln T/\sqrt{T})$—a degree of enlargement that makes no practical difference. Moreover, this result only concerns $\eta_t = \frac{\alpha}{\sqrt{t}}$; by setting $\eta_t = \frac{\alpha}{\sqrt{t\ln t}}$, one could push it arbitrarily close to $\mathcal{O}(\sqrt{\ln T/T})$, but this doesn't represent any essential difference—it's more of a derivation exercise.

Weighted Averaging

In fact, for $\eqref{leq:avg-2-mid3}$, there's an even simpler way to handle things: divide both sides directly by $\sum_{t=1}^T \eta_t$, then substitute in $\boldsymbol{\varphi}=\boldsymbol{\theta}^*$ to get

\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{\Vert\boldsymbol{\theta}_1 - \boldsymbol{\theta}^*\Vert^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}\label{leq:avg-3}\end{equation}

This way, there's no need to assume the learning rate is monotonically decreasing—it's a rather loose convergence conclusion, one might say, describing the convergence behavior of a weighted average along the optimization trajectory. Some readers might wonder why we're weighting things at all—the answer is, again, the same as before: "all we want is some degree of convergence guarantee for SGD." So as far as approaching the optimal loss value goes, does it really matter whether we weight things or not?

In a sense, $\eqref{leq:avg-3}$ offers us more insight than $\eqref{leq:avg-2}$ does. For instance, for the right-hand side to converge as fast as possible, the sum of learning rates $\sum_{t=1}^T \eta_t$ should be as large as possible, while the sum of squares $\sum_{t=1}^T \eta_t^2$ should be as small as possible. If we want convergence to the optimum as $T\to\infty$, the learning rate should satisfy

\begin{equation}\sum_{t=1}^{\infty} \eta_t = \infty \qquad\text{and}\qquad \sum_{t=1}^{\infty} \eta_t^2 \bigg/ \sum_{t=1}^{\infty} \eta_t = 0\end{equation}

These are exactly the two classical conditions on learning rate schedules. Furthermore, $\eqref{leq:avg-3}$ has no explicit dependence on the terminal learning rate $\eta_T$, and so it also allows for $\eta_T\to 0$, giving us more flexibility in scheduling the learning rate. All in all, $\eqref{leq:avg-3}$ gives us a more "relaxed" feeling, and is also more aesthetically pleasing.

Summary

In this post, we've restarted the "Scientific Alchemy" series and extended the previous post's conclusion about SGD convergence on a bounded domain to the unbounded case, arriving at a richer set of results.

English translation of a post from 科学空间 | Scientific Spaces by 苏剑林. Original: https://kexue.fm/archives/11469
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.