Convergence Analysis of SGD for Machine Learning Practitioners for Convex functions

$$

We introduced the basics of the stochastic gradient descent (SGD) algorithm and its implementation on Python with the MNIST dataset as a test case in our article Mastering Stochastic Gradient Descent algorithm with the MNIST Dataset. We can see that stochastic gradient descent is one of the most powerful algorithms used in the world of machine learning.

We prefer the use of the SGD algorithm over Gradient Descent (GD) solely for one main purpose the data is too large to handle for our machine, so instead of processing the entire batch of samples for learning, we process sample by sample to update our algorithm.

A key thing to remember: We use SGD/GD algorithms in machine learning to minimize the loss function associated with our machine learning problem (the problem can be classification or regression). The loss function $f(\cdot )$, however, can have certain attributes, is it convex, non-convex, smooth, non-smooth, etc. The nature of this function $f(\cdot )$ governs the performance of the SGD algorithm and its effect on the rate of convergence.

Suggested Read: Convergence analysis of Stochastic Gradient Descent (SGD) when a function is L-smooth and non-convex

In this article, we will investigate the proof of convergence when our loss function $f(\cdot )$ is convex.

Preliminaries

Let’s quickly recap a few things before jumping onto the convergence analysis. The stochastic gradient descent algorithm is given as:

 Algorithm - Stochastic Gradient Descent (SGD) 

Input: A random vector ${\mathbf{x}}_0\in {\mathbb{R}}^n$ and a decaying step size ${\gamma }_t>0$;
1. for $t=1,2,\dots ,$ do
2. $$ Evaluate the stochastic gradient mapping at ${\mathbf{x}}_t$, i.e., evaluate $\mathrm{\nabla }F({\mathbf{x}}_t,{\xi }_t)$;
3. $$ Do the following update:
   $$    $${\mathbf{x}}_{t+1}:={\mathcal{P}}_{\mathcal{X}}({\mathbf{x}}_t-{\gamma }_t\mathrm{\nabla }F({\mathbf{x}}_t,{\xi }_t))$$
4. end for

In the above algorithm, ${\mathbf{x}}_t$ is the update of the objective variable $\mathbf{x}$ at time $t$, $\mathrm{\nabla }f({\mathbf{x}}_t)$ is the gradient of the loss function $f$, and ${\mathcal{P}}_{\mathcal{X}}$ is the projection operator. We use the projection to ensure that the update of our algorithm doesn’t cause the objective variable $\mathbf{x}$ to lie outside the feasible set, and if it does, the operator projects the update ${\mathbf{x}}_{t+1}$ back to the feasible set. The details of the projection operator is given as:

Projection Operator:

We will use some properties of the projection operator in the update rule in our convergence analysis. So, we will take a brief detour towards the projection operator.

The idea of the projection operator is to project back the update iterate to the convex set by finding the point in the convex set having minimum Euclidian distance. Mathematically, it is defined as:

$${\mathcal{P}}_X\left(\hat{\mathbf{x}}\right)=\underset{\mathbf{x}\in X}{\arg \min }{\Vert \mathbf{x}-\hat{\mathbf{x}}\Vert }_2^2.$$

Properties of Projection Operator:

Let set $X\subseteq {\mathbb{R}}^n$ be a nonempty closed convex set. Consider the projection mapping ${\mathcal{P}}_X:{\mathbb{R}}^n\to X$. The projection operator has the following properties:

1. Non-expansive:
   $$\begin{array}{r}{\Vert {\mathcal{P}}_X(\mathbf{x})–{\mathcal{P}}_X(\mathbf{y})\Vert }_2\le {\Vert \mathbf{x}–\mathbf{y}\Vert }_2\text{for all }\mathbf{x},\mathbf{y}\in {\mathbb{R}}^n.\end{array}$$
2. Distance: The distance is a convex function in terms of $\mathbf{x}$ from a set $X$ is given by:
   $$dist(\mathbf{x},X)={\Vert {\mathcal{P}}_X(\mathbf{x})–\mathbf{x}\Vert }_2$$

The Stochastic Optimization Problem

The fundamental idea of the stochastic gradient descent algorithm is to find the optimal solution of the following stochastic optimization problem:

$$\begin{array}{rl}\text{minimize}& f(\mathbf{x})\triangleq \mathbb{E}[F(\mathbf{x},\mathit{\xi })]\\ \text{subject to:}& \mathbf{x}\in X.\end{array}$$

Where $f:{\mathbb{R}}^n\to \mathbb{R}$ is an unknown deterministic function, $F:{\mathbb{R}}^n\times \mathrm{\Omega }\to \mathbb{R}$ is a known stochastic function, $\mathit{\xi }:\mathrm{\Omega }\to {\mathbb{R}}^d$ is a random variable (this is a random data sample), and $\mathbb{E}[\cdot ]$ denotes the expectation operator taking an expectation due to the random variable $\mathit{\xi }$.

Concept of Filtration:

In algorithms like stochastic gradient descent, we process sample by sample. Hence, they arrive randomly, and in convergence analysis, we need to keep track of the history of these random samples. This is denoted by a set known as Filtration, which contains the random samples till iteration $t$. Mathematically, it is defined as:

$${\mathcal{F}}_t\triangleq \{{\mathbf{x}}_0,{\mathit{\xi }}_0,{\mathit{\xi }}_1,\dots ,{\mathit{\xi }}_{t-1}\}\text{for all }t\ge 1,$$

and ${\mathcal{F}}_0\triangleq x_0$. Since whenever randomness is involved, we need expectations to remove that. So, we define the conditional expectation given the filtration as $\mathbb{E}[\bullet \mid {\mathcal{F}}_t]$.

The main idea of using conditional expectation is that in our Filtration, if we know ${\mathit{\xi }}_{t-1}$, then we know what ${\mathbf{x}}_t$ i.e., by knowing the randomness of ${\mathit{\xi }}_{t-1}$ we can find the ${\mathbf{x}}_t$ by using the update rule of the algorithm (Step 3). This is how filtration and conditional expectation is useful in removing the randomness, and we can characterize the next iterate ${\mathbf{x}}_{t+1}$.

Apart from Filtration, we need the following two assumptions to prove the convergence rates for the SGD algorithm:

Assumptions:

1. $\mathbb{E}[\mathrm{\nabla }F(\mathbf{x},\mathit{\xi })\mid \mathbf{x}]=\mathrm{\nabla }f(\mathbf{x})$ for all $\mathbf{x}\in X$.
2. $\mathbb{E}[{\Vert \mathrm{\nabla }F(\mathbf{x},\mathit{\xi })-\mathrm{\nabla }f(\mathbf{x})\Vert }^2\mid \mathbf{x}]\le {\sigma }^2$ for all $\mathbf{x}\in X$ for some $\sigma >0$.

This assumption $1$ implies that the stochastic gradient is an unbiased estimator of the true gradient of the objective function, and assumption $2$

Main Theorem

Consider the stochastic optimization problem and SGD Algorithm defined above. Let Assumption 1 hold. Let $X\subset {\mathbb{R}}^n$. Suppose $X$ is nonempty, compact, and convex. Assume that $f:{\mathbb{R}}^n\to \mathbb{R}$ is continuously differentiable and convex.

Let the step size be diminishing given by:

$${\gamma }_t:=\frac{{\gamma }_0}{\sqrt{t+1}}\text{for all }t\ge 0.$$

Consider the averaged iterate defined by:

$${\overline{x}}_T\triangleq \frac{\sum _{t=0}^{T-1}{x}_t}{T}\text{for all }T\ge 1.$$

Then, there exists some scalars $M>0$ and $C>0$ and an optimal solution vector $x^{\ast }\in X$ such that:

$$\mathbb{E}[f({\overline{x}}_T)]–f(x^{\ast })\le (\frac{M}{2{\gamma }_0}+{\gamma }_0(C^2+{\sigma }^2))\frac{1}{\sqrt{T}}\text{for all }T\ge 2.$$

Proof:

From the update rule of the algorithm (Step 3) and by using the non-expansivity of the projection operator, we obtain:

$$\begin{array}{rl}{\Vert {\mathbf{x}}_{t+1}–{\mathbf{x}}^{\ast }\Vert }^2& \le {\Vert {\mathcal{P}}_X({\mathbf{x}}_t–{\gamma }_t\mathrm{\nabla }F({\mathbf{x}}_t,{\mathit{\xi }}_t))–{\mathcal{P}}_X({\mathbf{x}}^{\ast })\Vert }^2\\ & \le {\Vert {\mathbf{x}}_t–{\mathbf{x}}^{\ast }\Vert }^2+{\gamma }_t^2{\Vert \mathrm{\nabla }F({\mathbf{x}}_t,{\mathit{\xi }}_t)\Vert }^2–2{\gamma }_t\mathrm{\nabla }F({\mathbf{x}}_t,{\mathit{\xi }}_t{)}^T({\mathbf{x}}_t-{\mathbf{x}}^{\ast }).\end{array}$$

We now define stochastic error as ${\mathbf{w}}_t\triangleq \mathrm{\nabla }F({\mathbf{x}}_t,{\mathit{\xi }}_t)-\mathrm{\nabla }f({\mathbf{x}}_t)$ for all $t\ge 0$. Invoking this by replacing $\mathrm{\nabla }F({\mathbf{x}}_t,{\mathit{\xi }}_t)$ for $\mathrm{\nabla }f({\mathbf{x}}_t)+{\mathbf{w}}_t$ we obtain:

$$\begin{array}{rl}{\Vert {\mathbf{x}}_{t+1}-{\mathbf{x}}^{\ast }\Vert }^2& \le {\Vert {\mathbf{x}}_k–x^{\ast }\Vert }^2+{\gamma }_t^2{\Vert \mathrm{\nabla }f({\mathbf{x}}_t)+{\mathbf{w}}_t\Vert }^2–2{\gamma }_t{(\mathrm{\nabla }f({\mathbf{x}}_t)+{\mathbf{w}}_t)}^T({\mathbf{x}}_t-{\mathbf{x}}^{\ast })\\ & ={\Vert {\mathbf{x}}_t–{\mathbf{x}}^{\ast }\Vert }^2+{\gamma }_t^2({\Vert \mathrm{\nabla }f({\mathbf{x}}_t)\Vert }^2+{\Vert {\mathbf{w}}_t\Vert }^2+2\mathrm{\nabla }f({\mathbf{x}}_t{)}^T{\mathbf{w}}_t)–2{\gamma }_t{(\mathrm{\nabla }f({\mathbf{x}}_t)+{\mathbf{w}}_t)}^T({\mathbf{x}}_t-{\mathbf{x}}^{\ast }).\end{array}$$

Taking conditional expectation w.r.t Filtration operator on both sides, we obtain:

$$\begin{array}{rl}\mathbb{E}[{\Vert {\mathbf{x}}_{t+1}-{\mathbf{x}}^{\ast }\Vert }^2\mid {\mathcal{F}}_t]& \le {\Vert {\mathbf{x}}_k–{\mathbf{x}}^{\ast }\Vert }^2+{\gamma }_t^2({\Vert \mathrm{\nabla }f({\mathbf{x}}_t)\Vert }^2+\mathbb{E}[{\Vert {\mathbf{w}}_t\Vert }^2\mid {\mathcal{F}}_t]+2\mathrm{\nabla }f({\mathbf{x}}_t{)}^T\mathbb{E}[{\mathbf{w}}_t\mid {\mathcal{F}}_t])\\ & -2{\gamma }_t{(\mathrm{\nabla }f({\mathbf{x}}_t)+\mathbb{E}[{\mathbf{w}}_t\mid {\mathcal{F}}_t])}^T({\mathbf{x}}_t-{\mathbf{x}}^{\ast }).\end{array}$$

We now briefly detour and define the following Lemma, which captures the statistical properties of stochastic errors.

Lemma 1: Consider the definition of stochastic error ${\mathbf{w}}_t$ and under the defined assumption 1 we have for all $t\ge 0$:

$$\begin{array}{rl}& \mathbb{E}[{\mathbf{w}}_t\mid {\mathcal{F}}_t]={\mathbf{0}}_n\\ & \mathbb{E}[{\Vert {\mathbf{w}}_t\Vert }^2\mid {\mathcal{F}}_t]\le {\sigma }^2.\end{array}$$

(Now continuing the proof): Since $X$ is compact and that $f$ is continuous (due to convexity), there exists $C>0$ such that $\underset{\mathbf{x}\in X}{sup}\Vert \mathrm{\nabla }f(\mathbf{x})\Vert \le C$. Invoking Assumption 1 and Lemma 1, we obtain:

$$\begin{array}{rl}\mathbb{E}[{\Vert {\mathbf{x}}_{t+1}–{\mathbf{x}}^{\ast }\Vert }^2\mid {\mathcal{F}}_t]& \le {\Vert {\mathbf{x}}_t–{\mathbf{x}}^{\ast }\Vert }^2+{\gamma }_t^2(C^2+{\sigma }^2)–2{\gamma }_t\mathrm{\nabla }f({\mathbf{x}}_t{)}^T({\mathbf{x}}_t-{\mathbf{x}}^{\ast }).\end{array}$$

From convexity of $f$, we have that $\mathrm{\nabla }f({\mathbf{x}}_t{)}^T({\mathbf{x}}_t-{\mathbf{x}}^{\ast })\ge f({\mathbf{x}}_t)-f({\mathbf{x}}^{\ast })$. We obtain:

$$\begin{array}{rl}\mathbb{E}[{\Vert {\mathbf{x}}_{t+1}–{\mathbf{x}}^{\ast }\Vert }^2\mid {\mathcal{F}}_t]& \le {\Vert {\mathbf{x}}_k–{\mathbf{x}}^{\ast }\Vert }^2+{\gamma }_t^2(C^2+{\sigma }^2)-2{\gamma }_t(f({\mathbf{x}}_t)–f({\mathbf{x}}^{\ast })).\end{array}$$

Dividing both sides by $2{\gamma }_t$ and rearranging the terms, we obtain:

$$\begin{array}{r}f({\mathbf{x}}_t)–f({\mathbf{x}}^{\ast })\le \frac{1}{2{\gamma }_t}({\Vert {\mathbf{x}}_t-{\mathbf{x}}^{\ast }\Vert }^2–\mathbb{E}[{\Vert {\mathbf{x}}_{t+1}-{\mathbf{x}}^{\ast }\Vert }^2\mid {\mathcal{F}}_t])+\frac{{\gamma }_t}{2}(C^2+{\sigma }^2).\end{array}$$

Taking expectation w.r.t ${\mathcal{F}}_t$ on both sides and using the tower property of expectation, we obtain:

$$\begin{array}{rl}\mathbb{E}[f({\mathbf{x}}_t)]–f({\mathbf{x}}^{\ast })& \le \frac{1}{2{\gamma }_t}(\mathbb{E}[{\Vert {\mathbf{x}}_t–{\mathbf{x}}^{\ast }\Vert }^2]–\mathbb{E}[\mathbb{E}[{\Vert {\mathbf{x}}_{t+1}-{\mathbf{x}}^{\ast }\Vert }^2\mid {\mathcal{F}}_t]])+\frac{{\gamma }_t}{2}(C^2+{\sigma }^2)\\ & =\frac{1}{2{\gamma }_t}(\mathbb{E}[{\Vert {\mathbf{x}}_k–{\mathbf{x}}^{\ast }\Vert }^2]–\mathbb{E}[{\Vert {\mathbf{x}}_{t+1}–{\mathbf{x}}^{\ast }\Vert }^2])+\frac{{\gamma }_t}{2}(C^2+{\sigma }^2).\end{array}$$

Adding and subtracting $\frac{1}{2{\gamma }_{t-1}}\mathbb{E}[{\Vert {\mathbf{x}}_t–{\mathbf{x}}^{\ast }\Vert }^2]$, we obtain:

From boundedness of the set $X$, there exists a scalar $M>0$ such that $\underset{x\in X}{max}|x-x^{\ast }{|}^2\le M$. Thus, we have:

$$\begin{array}{rl}\mathbb{E}[f({\mathbf{x}}_t)]–f({\mathbf{x}}^{\ast })& \le (\frac{1}{2{\gamma }_{t-1}}\mathbb{E}[{\Vert {\mathbf{x}}_t–{\mathbf{x}}^{\ast }\Vert }^2]–\frac{1}{2{\gamma }_t}\mathbb{E}[{\Vert {\mathbf{x}}_{t+1}-{\mathbf{x}}^{\ast }\Vert }^2])\\ & +\left(\frac{1}{2{\gamma }_t}–\frac{1}{2{\gamma }_{t-1}}\right)M+\frac{{\gamma }_t}{2}(C^2+{\sigma }^2).\end{array}$$

Summing both sides over $t$ and dividing by $T$. On L.H.S., we can note that the function is convex, so we can apply Jensen’s inequality

$$\begin{array}{rl}\sum _{t=0}^T\mathbb{E}[f({\mathbf{x}}_t)]–Tf({\mathbf{x}}^{\ast })& \le \sum _{t=0}^T(\frac{1}{2{\gamma }_{t-1}}\mathbb{E}[{\Vert {\mathbf{x}}_t–{\mathbf{x}}^{\ast }\Vert }^2]–\frac{1}{2{\gamma }_t}\mathbb{E}[{\Vert {\mathbf{x}}_{t+1}-{\mathbf{x}}^{\ast }\Vert }^2])\\ & +\sum _{t=0}^T\left(\frac{1}{2{\gamma }_t}–\frac{1}{2{\gamma }_{t-1}}\right)M+\sum _{t=0}^T\frac{{\gamma }_t}{2}(C^2+{\sigma }^2).\\ \mathbb{E}f\left[\frac{1}{T}\sum _{t=0}^T{\mathbf{x}}_t\right]–f({\mathbf{x}}^{\ast })& \le \frac{1}{T}\sum _{t=0}^T[\frac{1}{2{\gamma }_{t-1}}\mathbb{E}[{\Vert {\mathbf{x}}_t–{\mathbf{x}}^{\ast }\Vert }^2]–\frac{1}{2{\gamma }_t}\mathbb{E}[{\Vert {\mathbf{x}}_{t+1}–{\mathbf{x}}^{\ast }\Vert }^2]]\\ & +\frac{1}{T}\sum _{t=0}^T\left(\frac{1}{2{\gamma }_t}–\frac{1}{2{\gamma }_{t-1}}\right)M+\frac{C^2+{\sigma }^2}{2T}\sum _{t=0}^T{\gamma }_t\end{array}$$

Now we apply the fact that \bar x_T = \frac{1}{T} \sum_{t=0}^T \bx_t$and$\gamma_t = \frac{\gamma_0}{\sqrt{t+1}}$

$$\begin{array}{rl}\mathbb{E}f({\overline{x}}_T)–f({\mathbf{x}}^{\ast })& \le \frac{1}{T}\sum _{t=0}^T\left[\frac{\sqrt{T}\mathbb{E}[{\Vert {\mathbf{x}}_t–{\mathbf{x}}^{\ast }\Vert }^2]–\sqrt{T+1}\mathbb{E}[{\Vert {\mathbf{x}}_{t+1}–{\mathbf{x}}^{\ast }\Vert }^2]}{2{\gamma }_0}\right]\\ & +\frac{1}{T}\sum _{t=0}^T\left(\frac{\sqrt{T+1}–\sqrt{T}}{2{\gamma }_0}\right)M+\frac{C^2+{\sigma }^2}{2T}\sum _{t=0}^T{\gamma }_t\end{array}$$

By noticing the telescoping series and applying the inequality $\sum _{t=0}^T{\gamma }_t\le 2{\gamma }_0\sqrt{T}$ and further noticing the fact that:

$$\frac{M}{2{\gamma }_0}+{\gamma }_0(C^2+{\sigma }^2)\ge \frac{M}{2{\gamma }_0}+{\gamma }_0(C^2+{\sigma }^2)–\frac{\mathbb{E}[{\Vert {\mathbf{x}}_t–{\mathbf{x}}^{\ast }\Vert }^2]}{2{\gamma }_0}$$

We obtain:

$$\begin{array}{rl}\mathbb{E}[f({\overline{x}}_T)]–f({\mathbf{x}}^{\ast })& \le \frac{1}{\sqrt{T}}[\frac{M}{2{\gamma }_0}+{\gamma }_0(C^2+{\sigma }^2)]\mathrm{\forall }T\ge 2\end{array}$$

This concludes the proof.