Making Model Training a Bit More Scientific (I): Convergence of the Average Loss for SGD

让炼丹更科学一些（一）：SGD的平均损失收敛 苏剑林 · 2023-12-19 · Machine Learning inequality optimizer SGD model tuning

We often jokingly call the training process of deep learning models "alchemy," because the whole procedure resembles ancient alchemy: it looks like it has some scientific basis, yet overall it gives off a rather "mysterious and elusive" feeling. Although this blog has previously looked at some work related to optimizers, and even written a series called A Dynamical-Systems View of Optimization Algorithms, those were fairly superficial introductions that didn't touch on deeper theory. To make future training a bit more scientific, I've decided to catch up on some optimization-related theoretical results, in the hope of giving the path of model training a bit more theoretical support.

In this post, we'll study a very basic convergence result for stochastic gradient descent (SGD). Although it looks crude and impractical by today's standards, it represents a very important early attempt at proving optimizer convergence — in particular, it explicitly takes into account the fact that we actually use stochastic gradient descent (SGD) rather than full-batch gradient descent (GD), which makes the conclusion more relevant as a reference.

Problem Setup

Let the loss function be $L(\boldsymbol{x},\boldsymbol{\theta})$, where $\boldsymbol{x}$ is the training set and $\boldsymbol{\theta}\in\mathbb{R}^N$ are the trainable parameters. Constrained by compute, we can typically only perform stochastic gradient descent (SGD), i.e., at each step we sample a subset of the training set to compute the loss function and update the parameters. Assuming the sampling is i.i.d., and that the subset sampled at step $t$ is $\boldsymbol{x}_t$, then we can reasonably regard the actual optimization objective as

\begin{equation}L(\boldsymbol{\theta}) = \lim_{T\to\infty}\frac{1}{T}\sum_{t=1}^T L(\boldsymbol{x}_t,\boldsymbol{\theta})\label{eq:loss}\end{equation}more

In practice, we can also only train for a finite number of steps, so we assume $T$ is a sufficiently large positive integer constant. Our goal is to find the minimizer of $L(\boldsymbol{\theta})$, i.e., we want to find $\boldsymbol{\theta}^*$:

\begin{equation}\boldsymbol{\theta}^* = \mathop{\text{argmin}}_{\boldsymbol{\theta}\in\mathbb{R}^N} L(\boldsymbol{\theta})\label{eq:argmin}\end{equation}

Now, consider the following SGD iteration:

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

where $\eta_t > 0$ is the learning rate, and $\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta})\triangleq \nabla_{\boldsymbol{\theta}}L(\boldsymbol{x}_t,\boldsymbol{\theta})$ is the gradient of $L(\boldsymbol{x}_t,\boldsymbol{\theta})$ with respect to $\boldsymbol{\theta}$. Our task is to analyze whether, by iterating in this way, $\boldsymbol{\theta}_t$ can converge to the target point $\boldsymbol{\theta}^*$.

A First Look at the Result

First, let's state the inequality we ultimately want to prove: under suitable assumptions, we have

\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{neq:core}\end{equation}

where $R,G$ is a constant unrelated to the optimization process. We'll introduce the "suitable assumptions" one by one later; before that, let's first look at what the inequality $\eqref{neq:core}$ actually says:

1. The first term on the left is the average, over the optimization process, of the loss function at each step;
2. The second term on the left, according to equation $\eqref{eq:loss}$, can be regarded as essentially equal to $L(\boldsymbol{\theta}^*)$ when $T$ is sufficiently large;
3. Combining these, the left-hand side is the gap between the average loss during optimization and the theoretical minimum of the loss function;
4. The right-hand side is an expression that depends only on the learning rate schedule $\{\eta_t\}$.

Putting points 1–4 together, the inequality $\eqref{neq:core}$ says: under suitable assumptions, the gap between the SGD average loss and the ideal target we're seeking can be bounded by an expression that depends only on the learning rate schedule. If we can choose an appropriate learning rate schedule that drives this expression to zero, it means the average loss of SGD is guaranteed to converge to the theoretical optimum. (Of course, theoretically, this result can only guarantee finding the minimum value $L(\boldsymbol{\theta}^*)$ of the loss function, but cannot guarantee finding the actual minimizer $\boldsymbol{\theta}^*$.)

In plain terms, this is a theoretical result about the conditions under which SGD converges. By the way, if we multiply the left-hand side of inequality $\eqref{neq:core}$ by $T$, we get a quantity with a special name: "Regret" (sometimes also translated as "悔" in some tutorials).

Two Examples

For example, suppose the learning rate is a constant $\eta$; then for the right-hand side of inequality $\eqref{neq:core}$ we have

\begin{equation}\frac{R^2}{2T\eta} + \frac{G^2}{2T}\sum_{t=1}^T\eta = \frac{R^2}{2T\eta} + \frac{G^2\eta}{2}\geq \frac{RG}{\sqrt{T}}\end{equation}

Equality holds when $\eta=\frac{R}{G\sqrt{T}}$, that is, when the learning rate is set to the constant $\frac{R}{G\sqrt{T}}$, in which case we get

\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{RG}{\sqrt{T}}\label{neq:case-1}\end{equation}

As $T\to\infty$, the right-hand side tends to zero, meaning that when the number of training steps $T$ is sufficiently large, setting the learning rate to the constant $\frac{R}{G\sqrt{T}}$ allows the gap between the average of SGD's iterates and the theoretical optimum to be made arbitrarily small.

Another example is to consider a decay schedule $\eta_t = \frac{\alpha}{\sqrt{t}}$. Using

\begin{equation}\sum_{t=1}^T \frac{1}{\sqrt{t}} = 1+\sum_{t=2}^T \frac{1}{\sqrt{t}}\leq 1+\sum_{t=2}^T \frac{2}{\sqrt{t-1} + \sqrt{t}}=1+\sum_{t=2}^T 2(\sqrt{t}-\sqrt{t-1})=2\sqrt{T}-1 < 2\sqrt{T}\end{equation}

and substituting into equation $\eqref{neq:core}$, we get

\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}^*) < \frac{R^2}{2\alpha\sqrt{T}} + \frac{G^2\alpha}{\sqrt{T}}\label{neq:case-2}\end{equation}

Equations $\eqref{neq:case-2}$ and $\eqref{neq:case-1}$ are both $\mathcal{O}\left(\frac{1}{\sqrt{T}}\right)$ with respect to $T$, so in theory both converge. Compared with $\eqref{neq:case-1}$, the constant in $\eqref{neq:case-2}$ is larger, which suggests that $\eta_t\equiv\frac{R}{G\sqrt{T}}$ likely converges faster than $\eta_t = \frac{\alpha}{\sqrt{t}}$. In practice, however, we prefer to use the latter, because the former requires fixing the total number of training steps $T$ in advance — once training finishes, that's it, and the precision is also fixed — whereas the latter has no such restriction; in fact $\alpha$ doesn't even need to be tuned — we can simply use $\eta_t = \frac{1}{\sqrt{t}}$ and keep training indefinitely, and in theory the gap between the average loss and the theoretical minimum will keep shrinking.

Even so, learning rate schedules like $\eta_t = \frac{1}{\sqrt{t}}$, whether in magnitude or in how they evolve, are still quite far from what we actually use in practice, so it's not hard to guess that a number of strong assumptions must have been baked in. Without further ado, let's dive into the proof and unpack these assumptions one by one.

The Proof

At the start of the proof, we assume that for any $\boldsymbol{x}$, $L(\boldsymbol{x},\boldsymbol{\theta})$ is a convex function of $\boldsymbol{\theta}$. This is a very strong assumption that is usually quite far from the truth in actual training, but there's no way around it — theoretical analysis generally has to make some strong assumptions, and then apply the resulting conclusions heuristically to real-world scenarios.

There are many different equivalent definitions of convex functions; here we use the following:

\begin{equation}L(\boldsymbol{x}_t,\boldsymbol{\theta}_2) - L(\boldsymbol{x}_t,\boldsymbol{\theta}_1) \geq (\boldsymbol{\theta}_2-\boldsymbol{\theta}_1)\cdot\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_1),\quad \forall \boldsymbol{\theta}_1,\boldsymbol{\theta}_2\label{eq:convex}\end{equation}

where $\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_1)\triangleq \nabla_{\boldsymbol{\theta}}L(\boldsymbol{x}_t,\boldsymbol{\theta})$ is the gradient of $L(\boldsymbol{x}_t,\boldsymbol{\theta})$ with respect to $\boldsymbol{\theta}$, and $\cdot$ denotes the inner product of vectors. The geometric meaning of this definition is that the graph of a convex function always lies above its tangent line (or tangent plane).

The key idea of the proof is to consider the distance between $\boldsymbol{\theta}_{t+1}$ and $\boldsymbol{\theta}^*$:

\begin{equation}\begin{aligned} \Vert\boldsymbol{\theta}_{t+1} - \boldsymbol{\theta}^*\Vert^2=&\, \Vert\boldsymbol{\theta}_t - \eta_t \boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)- \boldsymbol{\theta}^*\Vert^2 \\ =&\, \Vert\boldsymbol{\theta}_t - \boldsymbol{\theta}^*\Vert^2 - 2\eta_t (\boldsymbol{\theta}_t- \boldsymbol{\theta}^*)\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}

Rewrite this as

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

According to equation $\eqref{eq:convex}$, we have $L(\boldsymbol{x}_t,\boldsymbol{\theta}_t) - L(\boldsymbol{x}_t,\boldsymbol{\theta}^*)\leq (\boldsymbol{\theta}_t- \boldsymbol{\theta}^*)\cdot\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)$; substituting this into the above gives

\begin{equation}L(\boldsymbol{x}_t,\boldsymbol{\theta}_t) - L(\boldsymbol{x}_t,\boldsymbol{\theta}^*)\leq \frac{\Vert\boldsymbol{\theta}_t - \boldsymbol{\theta}^*\Vert^2}{2\eta_t} - \frac{\Vert\boldsymbol{\theta}_{t+1} - \boldsymbol{\theta}^*\Vert^2}{2\eta_t} + \frac{1}{2}\eta_t\Vert\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\Vert^2\end{equation}

Summing both sides over $t=1,2,\cdots,T$:

\begin{equation}\sum_{t=1}^T L(\boldsymbol{x}_t,\boldsymbol{\theta}_t) - \sum_{t=1}^TL(\boldsymbol{x}_t,\boldsymbol{\theta}^*)\leq \sum_{t=1}^T\left(\frac{\Vert\boldsymbol{\theta}_t - \boldsymbol{\theta}^*\Vert^2}{2\eta_t} - \frac{\Vert\boldsymbol{\theta}_{t+1} - \boldsymbol{\theta}^*\Vert^2}{2\eta_t}\right) + \sum_{t=1}^T\frac{1}{2}\eta_t\Vert\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\Vert^2\label{neq:base}\end{equation}

We introduce one more new assumption — that $\eta_t$ is a monotonically decreasing function of $t$ (i.e., $\eta_t\geq \eta_{t+1}$) — and, letting $R = \max_t \Vert\boldsymbol{\theta}_t - \boldsymbol{\theta}^*\Vert$, we obtain

\begin{equation}\begin{aligned} &\,\sum_{t=1}^T\left(\frac{\Vert\boldsymbol{\theta}_t - \boldsymbol{\theta}^*\Vert^2}{2\eta_t} - \frac{\Vert\boldsymbol{\theta}_{t+1} - \boldsymbol{\theta}^*\Vert^2}{2\eta_t}\right)\\ =&\,\frac{\Vert\boldsymbol{\theta}_1 - \boldsymbol{\theta}^*\Vert^2}{2\eta_1} - \frac{\Vert\boldsymbol{\theta}_{T+1} - \boldsymbol{\theta}^*\Vert^2}{2\eta_T} + \sum_{t=2}^T\left(\frac{1}{2\eta_t} - \frac{1}{2\eta_{t-1}}\right)\Vert\boldsymbol{\theta}_t - \boldsymbol{\theta}^*\Vert^2\\ \leq&\,\frac{R^2}{2\eta_1} + \sum_{t=2}^T\left(\frac{1}{2\eta_t} - \frac{1}{2\eta_{t-1}}\right)R^2\\ =&\, \frac{R^2}{2\eta_T} \end{aligned}\end{equation}

Finally, letting $G = \max_t \Vert\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\Vert$ and substituting into equation $\eqref{neq:base}$, we get

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

Dividing both sides by $T$ finally gives us inequality $\eqref{neq:core}$.

Note that the constant $R,G$ here is currently optimization-dependent — that is, we first need to fix a learning rate schedule $\{\eta_t\}$, then complete the optimization process, before we can obtain $R,G$. For $R,G$ to become a constant that is independent of the optimization process, we need to assume that for any $\{\eta_t\}$, both $\Vert\boldsymbol{\theta}_t - \boldsymbol{\theta}^*\Vert$ and $\Vert\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\Vert$ are bounded by some constant $R,G$, so that the right-hand side of inequality $\eqref{neq:core}$ depends only on the learning rate schedule.

Projection Onto a Domain

However, the final assumption — "for any $\{\eta_t\}$, both $\Vert\boldsymbol{\theta}_t - \boldsymbol{\theta}^*\Vert$ and $\Vert\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\Vert$ are bounded by some constant $R,G$" — always has a "putting the cart before the horse" feel to it: it seems to require completing the optimization first in order to determine $R,G$, whereas our goal is precisely to use the proven result to improve the optimization. To remove this odd feeling, let's replace these two assumptions with the following:

1. In equation $\eqref{eq:argmin}$, replace $\boldsymbol{\theta}\in\mathbb{R}^N$ with $\boldsymbol{\theta}\in\boldsymbol{\Theta}$, where $\boldsymbol{\Theta}\subseteq \mathbb{R}^N$ is a bounded convex set;
1.1) Bounded: $R=\max\limits_{\boldsymbol{\theta}_1,\boldsymbol{\theta}_2\in \boldsymbol{\Theta}}\Vert\boldsymbol{\theta}_1-\boldsymbol{\theta}_2\Vert < \infty$;
1.2) Convex set: for any $\forall \boldsymbol{\theta}_1,\boldsymbol{\theta}_2\in \boldsymbol{\Theta}$ and $\forall\lambda\in[0,1]$, we have $\lambda \boldsymbol{\theta}_1 + (1-\lambda)\boldsymbol{\theta}_2 \in \boldsymbol{\Theta}$.
2. For any $\boldsymbol{\theta}\in \boldsymbol{\Theta}$ and any $\boldsymbol{x}$, we have $\Vert\boldsymbol{g}(\boldsymbol{x},\boldsymbol{\theta})\Vert\leq G$

Point 2 might be easier to accept: it's really just adding one more assumption on top of the loss function $L(\boldsymbol{x},\boldsymbol{\theta})$ — a case of "in for a penny, in for a pound," since we've already made the very strong convexity assumption, one more assumption doesn't hurt. But assumption 1 might not be so easy to grasp: the convex set requirement is understandable, since a convex function can only be defined on a convex set to begin with, but how do we ensure boundedness? That is, how do we guarantee that the output of the iteration $\eqref{eq:sgd}$ is always bounded?

The answer is: "add one extra projection step." We define the projection operation:

\begin{equation}\Pi_{\boldsymbol{\Theta}} (\boldsymbol{\varphi}) = \mathop{\text{argmin}}_{\boldsymbol{\theta}\in\boldsymbol{\Theta}}\Vert\boldsymbol{\varphi}-\boldsymbol{\theta}\Vert\end{equation}

That is, we find the vector in $\boldsymbol{\Theta}$ closest to $\boldsymbol{\varphi}$. This allows us to rewrite equation $\eqref{eq:sgd}$ as

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

which guarantees that the iterates always stay within the set $\boldsymbol{\Theta}$.

But after this modification, does the proof from the previous section, and the earlier conclusion (mainly inequality $\eqref{neq:core}$), still hold? Fortunately, yes. We just need to show that for the projected SGD defined by equation $\eqref{eq:sgd-p}$,

\begin{equation}\Vert\boldsymbol{\theta}_{t+1} - \boldsymbol{\theta}^*\Vert \leq \Vert\boldsymbol{\theta}_t - \eta_t \boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t) - \boldsymbol{\theta}^*\Vert\end{equation}

holds; then the derivation in the "The Proof" section can still go through, except that some of the equalities become $\leq$. To this end, we just need to show that for $\forall \boldsymbol{\varphi}\in\mathbb{R}^N, \boldsymbol{\theta}\in \boldsymbol{\Theta}$,

\begin{equation}\Vert\Pi_{\boldsymbol{\Theta}} (\boldsymbol{\varphi}) - \boldsymbol{\theta}\Vert \leq \Vert \boldsymbol{\varphi} - \boldsymbol{\theta}\Vert\end{equation}

Proof:
The key to the proof is to combine the definition of a convex set with the definition of $\Pi_{\boldsymbol{\Theta}}$. First, by the definition of a convex set, we know that for $\forall \lambda\in(0,1)$ we always have $\lambda\boldsymbol{\theta} + (1-\lambda)\Pi_{\boldsymbol{\Theta}} (\boldsymbol{\varphi})\in \boldsymbol{\Theta}$, so by the definition of $\Pi_{\boldsymbol{\Theta}}$, the following always holds:
\begin{equation}\Vert\boldsymbol{\varphi} - \Pi_{\boldsymbol{\Theta}} (\boldsymbol{\varphi})\Vert\leq \Vert\boldsymbol{\varphi} - \lambda\boldsymbol{\theta} - (1-\lambda)\Pi_{\boldsymbol{\Theta}} (\boldsymbol{\varphi})\Vert = \Vert(\boldsymbol{\varphi} - \Pi_{\boldsymbol{\Theta}} (\boldsymbol{\varphi})) + \lambda(\Pi_{\boldsymbol{\Theta}} (\boldsymbol{\varphi})-\boldsymbol{\theta})\Vert\end{equation}
Squaring both sides and subtracting, we get
\begin{equation}\lambda^2\Vert\Pi_{\boldsymbol{\Theta}} (\boldsymbol{\varphi})-\boldsymbol{\theta}\Vert^2 + 2\lambda(\Pi_{\boldsymbol{\Theta}} (\boldsymbol{\varphi})-\boldsymbol{\theta})\cdot(\boldsymbol{\varphi} - \Pi_{\boldsymbol{\Theta}} (\boldsymbol{\varphi}))\geq 0\end{equation}
Note that we've just restricted $\lambda\in(0,1)$, so we can divide both sides by $\lambda$:
\begin{equation}\lambda\Vert\Pi_{\boldsymbol{\Theta}} (\boldsymbol{\varphi})-\boldsymbol{\theta}\Vert^2 + 2(\Pi_{\boldsymbol{\Theta}} (\boldsymbol{\varphi})-\boldsymbol{\theta})\cdot(\boldsymbol{\varphi} - \Pi_{\boldsymbol{\Theta}} (\boldsymbol{\varphi}))\geq 0\end{equation}
This is an identity that holds unconditionally, so it also holds for $\lambda\to 0^+$, giving
\begin{equation}(\Pi_{\boldsymbol{\Theta}} (\boldsymbol{\varphi})-\boldsymbol{\theta})\cdot(\boldsymbol{\varphi} - \Pi_{\boldsymbol{\Theta}} (\boldsymbol{\varphi}))\geq 0\end{equation}
Adding $\Vert\Pi_{\boldsymbol{\Theta}} (\boldsymbol{\varphi})-\boldsymbol{\theta}\Vert^2 + \Vert\boldsymbol{\varphi} - \Pi_{\boldsymbol{\Theta}} (\boldsymbol{\varphi})\Vert^2$ to both sides, the left-hand side becomes exactly $\Vert\boldsymbol{\varphi}-\boldsymbol{\theta}\Vert^2$, so we have
\begin{equation}\Vert\boldsymbol{\varphi}-\boldsymbol{\theta}\Vert^2\geq \Vert\Pi_{\boldsymbol{\Theta}} (\boldsymbol{\varphi})-\boldsymbol{\theta}\Vert^2 + \Vert\boldsymbol{\varphi} - \Pi_{\boldsymbol{\Theta}} (\boldsymbol{\varphi})\Vert^2 \geq \Vert\Pi_{\boldsymbol{\Theta}} (\boldsymbol{\varphi})-\boldsymbol{\theta}\Vert^2\end{equation}

Analysis of the Assumptions

At this point, we've completed the proof. The result above comes from Online Convex Programming and Generalized Infinitesimal Gradient Ascent, a 2003 paper. In particular, since the optimization literature is vast, and I myself am a beginner in this area, there is likely to be some inaccuracy or omission in tracing the sources further below — readers who are more familiar with this material are welcome to point out any errors.

Now let's "take stock" of all the assumptions used across the whole proof:

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}$);

Under these assumptions, projected SGD, i.e., equation $\eqref{eq:sgd-p}$, satisfies inequality $\eqref{neq:core}$.

Among these, assumptions 1 and 4 are entirely reasonable — indeed one could say quite sensible. For instance, for point 1: for practical computation, a sufficiently large ball is essentially indistinguishable from $\mathbb{R}^N$; and point 4, that the learning rate decreases, aligns well with existing intuition. The strongest and least realistic assumption is point 2, the convexity assumption — but there's no way around it; after reading a few more papers on optimization theory, one comes to terms with it, since nearly all optimization theory is built on convexity assumptions, and we can only hope that once optimization enters a certain region, the loss function will partially exhibit convex-like behavior. Point 3 is, in essence, also a very strong assumption, but in practice, if initialization is done well and the learning rate is set appropriately, the gradient norm can generally be kept within a certain range, so this assumption is usually acceptable too.

Summary

In this post, we revisited an old paper on convex optimization and introduced a very basic proof of SGD's convergence: under suitable (in fact, very strong) assumptions, the convergence of SGD can be guaranteed. Although these assumptions — such as the convexity assumption and the bound on gradient norm — may not always hold in practice, these theoretical results can still offer us important insight into the convergence behavior of SGD.