Chapter 5: Optimization Algorithms

An artificial neural network (ANN) consists of several parameters, which can be broadly categorized into two types, namely, trainable (or learnable) parameters and hyperparameters. Training a neural network is the process of determining appropriate values of the trainable parameters (weights and biases) such that the ANN function (defined in Definition 1.2) approximates the required input-output mapping. Often, such mappings are implicitly represented in training data and the trainable parameters of the network are iteratively updated to minimize the error between the predicted outputs and the true outputs. The error is quantified by an appropriately chosen loss function, and when averaged over the entire dataset, it is referred to as a cost function (see Definition 4.1).

The minimization of a cost function depends on the knowledge of how the changes in the network parameters affect the cost. This requires the computation of the gradient of the cost function with respect to the parameters. The backpropagation algorithm discussed in the previous chapter can be used to efficiently compute these gradients by propagating errors backward through the network, starting from the output layer to the first hidden layer. Once gradients are available, the parameters are updated using gradient-based optimization algorithms. The most basic method is the gradient descent method (GD), where parameters are iteratively updated in the direction opposite to the gradient in order to reduce the cost. We also have some important variants of the GD method, such as the stochastic gradient descent (SGD) method and the SGD method with momentum to improve training efficiency and stability.

In this chapter, we discuss the GD method and some of its variants. We first introduce the fundamental idea of the GD method in Section 5.1 and give three ways to implement this method, namely, the batch GD method, the SGD method, and the mini-batch GD method. In Section 5.2, we discuss some important modifications of the mini-batch GD method mainly to accelerate the convergence.

5.1 Gradient Descent Methods

Consider a cost function \(\mathcal{C}(\boldsymbol{\Theta})\) as given in Definition 4.1. The aim is to obtain a \(\boldsymbol{\Theta}^*\) such that

\[\boldsymbol{\Theta}^* = \text{argmin}_{\boldsymbol{\Theta}} \mathcal{C}(\boldsymbol{\Theta}).\]

In simple cases (like linear regression), we can obtain \(\boldsymbol{\Theta}^*\) analytically. But in deep learning, \(\mathcal{C}(\boldsymbol{\Theta})\) is nonlinear and has a large number of parameters. Hence, we need an efficient method to obtain (at least approximately) \(\boldsymbol{\Theta}^*\).

To iteratively minimize the cost function, we need to know in which direction to move \(\boldsymbol{\Theta}\) to decrease the cost most efficiently. This direction is determined by the gradient \(\nabla_{\boldsymbol{\Theta}}\mathcal{C}\).

5.1.1 Batch Gradient Descent Method

Let us consider a simple case of \(\boldsymbol{\Theta}=(w_1, w_2, \ldots, w_n)\in \mathbb{R}^n\) to motivate why we update the weights in the gradient direction \(\nabla_{\boldsymbol{\Theta}}\mathcal{C}\).

Let us start with an arbitrarily chosen vector \(\boldsymbol{w}_0\in \mathbb{R}^n\). We want to move to a new vector

\[\boldsymbol{w}_1 = \boldsymbol{w}_0 + \eta \boldsymbol{v},\]

for a small step size \(\eta>0\), where the direction \(\boldsymbol{v}\) is chosen such that \(\mathcal{C}(\boldsymbol{w}_1) < \mathcal{C}(\boldsymbol{w}_0)\).

Using the first-order Taylor expansion, we have

\begin{eqnarray} \mathcal{C}(\boldsymbol{w}_1) = \mathcal{C}(\boldsymbol{w}_0) + \eta \, \langle \nabla_{\boldsymbol{w}} \mathcal{C}(\boldsymbol{w}_0), \boldsymbol{v} \rangle + O(\eta^2). \end{eqnarray}(5.1)

Neglecting higher-order terms of \(O(\eta^2)\), the change in the cost function is approximately

\[ \mathcal{C}(\boldsymbol{w}_1) - \mathcal{C}(\boldsymbol{w}_0) \approx \eta \, \langle \nabla_{\boldsymbol{w}} \mathcal{C}(\boldsymbol{w}_0), \boldsymbol{v} \rangle. \]

To achieve a rapid decrease, we need to choose the direction vector \(\boldsymbol{v}\) such that \(\mathcal{C}(\boldsymbol{w}_1) - \mathcal{C}(\boldsymbol{w}_0)\) gives the largest decrease in the linear approximation (neglecting \(O(\eta^2)\) terms).

Using the Cauchy-Schwarz inequality, we get

\[ \langle \nabla_{\boldsymbol{w}} \mathcal{C}(\boldsymbol{w}_0), \boldsymbol{v} \rangle \ge -\| \nabla_{\boldsymbol{w}} \mathcal{C}(\boldsymbol{w}_0)\| \, \|\boldsymbol{v}\|, \]

and equality occurs when

\[\boldsymbol{v} = -\nabla_{\boldsymbol{w}} \mathcal{C}(\boldsymbol{w}_0).\]

Thus, the direction of steepest descent is the negative gradient, and the update rule becomes

\[ \boldsymbol{w}_1 = \boldsymbol{w}_0 - \eta \, \nabla_{\boldsymbol{w}} \mathcal{C}(\boldsymbol{w}_0). \]

In general, we may not reach the optimal vector \(\boldsymbol{w}^*\) in a single step. So, we need to apply the update rule repeatedly in order to move step by step in the direction of the negative gradient, until the cost function stops decreasing significantly. This leads to the general iterative gradient descent method (GDM), given by

\begin{eqnarray} \boldsymbol{\Theta}_{k+1} = \boldsymbol{\Theta}_{k} - \eta \, \nabla_{\boldsymbol{\Theta}} \mathcal{C}(\boldsymbol{\Theta}_{k}), \quad k = 0,1,2,\dots, \end{eqnarray}(5.2)

where \(\eta > 0\) is the learning rate and \(\boldsymbol{\Theta}_{0}\) is an initial guess.

Remark 5.1 (Gradient System):

Again, for the sake of simplicity, let us consider \(\boldsymbol{\Theta} = \boldsymbol{w} \in \mathbb{R}^n\), an \(n\)-dimensional vector. The gradient descent iteration (5.2) can be written as

\[ \frac{\boldsymbol{w}_{k+1} - \boldsymbol{w}_{k}}{\eta} = - \nabla_{\boldsymbol{w}} \mathcal{C}(\boldsymbol{w}_{k}), \quad k = 0,1,2,\dots \]

with step size \(\eta>0\).

Interpreting \(t_k=k\eta\) as a discrete time, the above scheme is the explicit (forward) Euler discretization of the system of ordinary differential equations

\[ \dot{\boldsymbol{w}} = - \nabla_{\boldsymbol{w}} \mathcal{C}(\boldsymbol{w}). \]

Such a system is called a gradient system. Thus, the gradient descent update rule can be viewed as an approximation of the gradient system or the gradient flow.

Further, assume that \(\mathcal{C}\in C^2(\mathbb{R}^n)\) and the Hessian \(\nabla_{\boldsymbol{w}}^2 \mathcal{C}\) is bounded (the gradient \(\nabla_{\boldsymbol{w}}\mathcal{C}\) is Lipschitz) in a neighborhood of a nondegenerate local minimum \(\boldsymbol{w}^*\) (strict minimum). Then, by choosing \(\boldsymbol{w}_0\) in this neighborhood, the piecewise linearly interpolated trajectory computed using the above Euler scheme converges (uniformly on any compact time interval) to the solution of the above gradient system as \(\eta \rightarrow 0\).

As a consequence, a minimizer of a \(C^2\) cost function can be viewed as an equilibrium point of the above gradient system in the parameter space. It is known that strict (nondegenerate) minima of \(\mathcal{C}\) are asymptotically stable equilibrium points of the above gradient system.

Problem 5.1:

Consider the gradient flow

\[ \dot{\boldsymbol{w}}(t) \;=\; -\,\nabla_{\boldsymbol{w}} \mathcal{C}(\boldsymbol{w}(t)), \qquad \boldsymbol{w}(0)=\boldsymbol{w}_0, \]

associated with a cost function \(\mathcal{C}:\mathbb{R}^n \to \mathbb{R}\). Show that the cost function is non-increasing along the trajectories of this gradient system, i.e.

\[ \frac{d}{dt}\,\mathcal{C}(\boldsymbol{w}(t)) \;\leq\; 0, \qquad t\ge0. \]

Algorithm 5.1 (Batch Gradient Descent Method):

Input:

Initial guess \(\boldsymbol{\Theta}_0 \in \mathbb{R}^n\);
Learning rate (step size) \(\eta>0\);
Maximum iterations \(K\).

Processing:
For \(k = 0,1,2,\dots,K-1\)

Compute the gradient \(\boldsymbol{g}_k = \nabla_{\boldsymbol{\Theta}} \mathcal{C}(\boldsymbol{\Theta}_k)\)
(using the backpropagation algorithm for the neural network).
Update the parameter:
\[ \boldsymbol{\Theta}_{k+1} = \boldsymbol{\Theta}_k - \eta \, \boldsymbol{g}_k. \]

Output: Approximate minimizer \(\boldsymbol{\Theta}_K\).

Note:

The gradient descent method is also called as the steepest descent method. Since the cost function includes the loss function evaluated at all the examples of a dataset, the method in the above algorithm is also called the batch gradient descent method or batch steepest descent method.

5.1.2 Stochastic Gradient Descent (SGD)

Recall, from Definition 4.1, that the cost function is defined as

\[ \mathcal{C}(\boldsymbol{\Theta}) = \frac{1}{N}\sum_{i=1}^N \ell_i(\boldsymbol{\Theta}),~~ \ell_i(\boldsymbol{\Theta}) = \ell\big(\boldsymbol{y}_i,\tilde{\boldsymbol{y}}_i;\boldsymbol{\Theta}\big), \]

where \(\ell\) is the loss corresponding to a single example \((\boldsymbol{x}_i, \boldsymbol{y}_i)\) of the dataset \(\mathcal{D}\), and \(\tilde{\boldsymbol{y}}_i\)(\(=\tilde{\boldsymbol{y}}_i(\boldsymbol{\Theta})\)) is the network prediction for the input \(\boldsymbol{x}_i\). The gradient is therefore given by

\[ \nabla_{\boldsymbol{\Theta}} \mathcal{C}(\boldsymbol{\Theta}) = \frac{1}{N}\sum_{i=1}^N \nabla_{\boldsymbol{\Theta}} \ell\big(\boldsymbol{y}_i,\tilde{\boldsymbol{y}}_i;\boldsymbol{\Theta}\big). \]

Since, for each iteration of GDM, the gradient has to be computed using the entire dataset of size \(N\), the computational cost of evaluating (5.2) per iteration grows linearly with \(N\). Since deep learning often involves large-scale datasets, the batch gradient descent method becomes prohibitively expensive.

An alternate approach is the stochastic gradient descent method (SGD) which replaces the gradient of the cost function by the gradient of the per-sample cost \(\ell_{i_k}(\boldsymbol{\Theta})\) for a randomly chosen sample \((\boldsymbol{x}_{i_k},\boldsymbol{y}_{i_k})\). For an initial choice \(\boldsymbol{\Theta}_{0}\), the update rule for each iteration is

\begin{eqnarray} \boldsymbol{\Theta}_{k+1} = \boldsymbol{\Theta}_k - \eta \, \nabla_{\boldsymbol{\Theta}} \ell_{i_k}(\boldsymbol{\Theta}_k), \end{eqnarray}(5.3)

where \(i_k \sim \operatorname{Uniform}\{1,\ldots,N\}\) (i.i.d.), for each \(k=0,1,2, \ldots.\)

Problem 5.2:

Write the algorithm for the stochastic gradient descent method.

Since \(i_k\) is uniformly sampled, we have

\[ \mathbb{E}\!\left[\nabla_{\boldsymbol{\Theta}} \ell_{i_k}(\boldsymbol{\Theta}_k)\, | \, \boldsymbol{\Theta}_k \right] = \nabla_{\boldsymbol{\Theta}} \mathcal{C}(\boldsymbol{\Theta}_k). \]

and therefore, \(\nabla_{\boldsymbol{\Theta}} \ell_{i_k}(\boldsymbol{\Theta}_k)\) is an unbiased estimator of the cost gradient.

Thus, SGD can be viewed as a noisy discretization of the gradient system. The stochasticity introduced by the procedure often helps in escaping saddle points and shallow local minima. However, it causes fluctuations in the trajectory and leads to slow convergence.

Figure 5.1: The two individual loss surfaces \(\ell_1\) and \(\ell_2\) corresponding to two training examples, and their average, the cost function \(\mathcal{C}\) given in Example 5.1 is depicted. Each individual loss is a parabolic cylinder with an infinitely extended valley, whereas their average is an elliptic paraboloid with a unique minimum.

Example 5.1:

Consider a linear model

\[ \widehat{y}=w_1x+w_2 \]

and a training dataset with two examples

\[\mathcal{D} = \big\{ (x_1,y_1)=(-1,1),~ (x_2,y_2)=(1,2) \big\}. \]

Using the squared-error loss

\[ \ell_i(w_1,w_2) = \frac{1}{2}\left(y_i-\widehat{y}_i\right)^2, ~i=1,2,\]

the losses corresponding to the two examples are

\[\begin{array}{rcl} \ell_1(w_1,w_2) &=& \frac12(1+w_1-w_2)^2, \\ \ell_2(w_1,w_2) &=& \frac12(2-w_1-w_2)^2. \end{array}\]

The cost function is given by

\[\begin{array}{rcl} \mathcal{C}(w_1,w_2) &=& \dfrac{\ell_1(w_1,w_2)+\ell_2(w_1,w_2)}{2} \\ &=& \dfrac14\left[(1+w_1-w_2)^2 +(2-w_1-w_2)^2\right]. \end{array}\]

Geometrically, each individual loss is a parabolic cylinder, as depicted in Figure 5.1, where there is a whole line along which the loss remains zero. Thus, each individual loss surface has a valley that is infinitely extended in one direction. In contrast, the cost function surface is an elliptic paraboloid with a unique minimum, as observed in Figure 5.1.

We observed that the mean-square cost function in the above example has a unique minimum. However, this is not always the case, even for linear models.

Problem 5.3:

Consider a linear model

\[ \widehat{y}_i=\mathbf{x}_i^{T}\mathbf{w}, \qquad i=1,\ldots,N, \]

and the squared-error cost function

\[ \mathcal{C}(\mathbf{w}) = \frac{1}{2N}\sum_{i=1}^{N} \left(y_i-\mathbf{x}_i^{T}\mathbf{w}\right)^2. \]

Suppose that the Hessian matrix

\[ H=\nabla^2\mathcal{C}(\mathbf{w}) \]

is positive definite. Prove that the cost function \(\mathcal{C}\) has a unique minimum.

The uniqueness of the minimizer \(\boldsymbol{w}^*\) established above is a consequence of the cost function \(\mathcal{C}\) being strict convexity.

Problem 5.4:

With the assumptions of the above problem, show that the cost function \(\mathcal{C}\) is strictly convex.

Remark 5.2 (Global versus Local Convexity):

For the linear model considered above, the Hessian

\[ \nabla^2\mathcal{C}(\boldsymbol{w})=H=\frac{1}{N}X^TX \]

is independent of \(\boldsymbol{w}\), and it is positive definite at every point of \(\mathbb{R}^n\) whenever \(X\) has full column rank. Consequently, \(\mathcal{C}\) is strictly convex on the whole of \(\mathbb{R}^n\). Hence \(\boldsymbol{w}^*\) is the global minimizer.

For a genuinely nonlinear model, the Hessian of \(\mathcal{C}\) typically varies with \(\boldsymbol{w}\), and hence need not be positive definite everywhere. Consequently, \(\mathcal{C}\) need not be globally convex. Therefore, as in Remark 5.1, the positive definiteness of the Hessian near a nondegenerate critical point \(\boldsymbol{w}^*\) gives strict convextiy of \(\mathcal{C}\) only locally, in a small neighborhood of \(\boldsymbol{w}^*\). So, we can only expect \(\boldsymbol{w}^*\) to be a local minimizer.

5.1.3 Mini-batch Gradient Descent

The batch GD method uses the gradient of the cost surface, which is the average of the gradient of the loss function evaluated at all the training examples. This results in a path that smoothly descends towards the minimum, but it is computationally expensive. On the other hand, the stochastic GD method uses the gradient corresponding to a randomly selected training example at each iteration. This improves computational efficiency, but its path is more irregular. Figure 5.2 depicts a visual illustration of the typical paths followed by the batch GD and stochastic GD methods.

Figure 5.2: Typical paths followed by batch GD and stochastic GD from the same starting point. The ellipses represent level curves of the cost function, with the innermost curve surrounding its minimum.

A compromise between batch GD and SGD is to average the gradients over a small randomly chosen subset (mini-batch) \(\mathcal{B}_k\subset \{1,\dots,N\}\) with \(|\mathcal{B}_k|=M\ll N\). The update reads

\begin{eqnarray} \boldsymbol{\Theta}_{k+1} = \boldsymbol{\Theta}_k - \eta \nabla_{\boldsymbol{\Theta}} \mathcal{C}_{\mathcal{B}_k}(\boldsymbol{\Theta}_k) , \end{eqnarray}(5.4)

where

\[ \mathcal{C}_{\mathcal{B}_k}(\boldsymbol{\Theta}_k) = \frac{1}{M}\sum_{i\in \mathcal{B}_k} \ell_i(\boldsymbol{\Theta}_k). \]

This method is called the mini-batch gradient descent method. Mini-batches reduce the variance of the stochastic gradient and provide a better computational efficiency than GD method. In practice, batch sizes such as \(32\), \(64\), or \(128\) are widely used.

Another approach in mini-batch gradient descent method, which is common in practice, is the epoch-wise training. Here, in each epoch, the entire training dataset is first randomly shuffled and then divided into mini-batches of size \(M\ll N\). This results in at least \(N/M\) mini-batches, which are used to update \(\boldsymbol{\Theta}\) sequentially within that epoch. If \(N\) is not exactly divisible by \(M\), then the last mini-batch of the epoch contains the remaining number of examples. Thus, each mini-batch produces one parameter update, and one epoch corresponds to a complete pass over the training dataset, where all mini-batches are processed once.

Problem 5.5:

Write the algorithm for the epoch-wise training of a mini-batch gradient descent method.

Problem 5.6:

Consider the mini-batch gradient estimator

\[ \boldsymbol{g}_{\mathcal{B}_k}(\boldsymbol{\Theta}) = \frac{1}{M}\sum_{i\in \mathcal{B}_k} \boldsymbol{g}_i(\boldsymbol{\Theta}), \qquad \boldsymbol{g}_i(\boldsymbol{\Theta}) := \nabla_{\boldsymbol{\Theta}} \ell_i(\boldsymbol{\Theta}), \]

where the indices in \(\mathcal{B}_k\subset\{1,\ldots,N\}\), with \(|\mathcal{B}_k|=M\), are drawn independently and uniformly from \(\{1,\ldots,N\}\), as in the SGD sampling scheme. Let

\[ \sigma^2(\boldsymbol{\Theta}) := \mathbb{E}\left[\|\boldsymbol{g}_{i_k}(\boldsymbol{\Theta}) - \nabla_{\boldsymbol{\Theta}}\mathcal{C}(\boldsymbol{\Theta})\|^2 \,|\, \boldsymbol{\Theta}\right] \]

denote the variance of the single-sample SGD gradient estimator. Show that

\[ \mathbb{E}\left[\|\boldsymbol{g}_{\mathcal{B}_k}(\boldsymbol{\Theta}) - \nabla_{\boldsymbol{\Theta}}\mathcal{C}(\boldsymbol{\Theta})\|^2 \,|\, \boldsymbol{\Theta}\right] = \frac{\sigma^2(\boldsymbol{\Theta})}{M}. \]

5.2 Momentum and Adaptive Gradient Methods

In this section, we discuss some variants of the GD method discussed in the previous section. These variants are proposed to accelerate the convergence of the GD method.

5.2.1 Momentum Method

Recall, in physics, the momentum of a moving particle is defined as the product of mass and velocity of the particle. Taking the physical motivation, we add an extra term \(\mu \boldsymbol{v}\) to the mini-batch SGD method to get

\begin{eqnarray} \boldsymbol{\Theta}_{k+1} = \boldsymbol{\Theta}_{k} - \eta \nabla_{\boldsymbol{\Theta}} \mathcal{C}_{\mathcal{B}_k}(\boldsymbol{\Theta}_k) + \mu \boldsymbol{v}_k,~k=1,2,\ldots, \end{eqnarray}(5.5)

is the mini-batch cost function, \(\mu\in [0, 1)\) is the momentum parameter, and

\[ \boldsymbol{v}_k = \boldsymbol{\Theta}_{k} - \boldsymbol{\Theta}_{k-1}, \]

which carries the memory of the previous update direction into the future updates. As will be seen from the physical derivation below, \(\mu\) plays the role of a friction (damping) coefficient, while \(\boldsymbol{v}_k\) resembles the velocity of a moving particle.

By taking \(\boldsymbol{v}_0=\boldsymbol{0}\), the iterative formula (5.5) can be rewritten as

\begin{eqnarray} \left.\begin{array}{lcl} \boldsymbol{v}_{k+1} &=& \mu \boldsymbol{v}_k - \eta \nabla_{\boldsymbol{\Theta}} \mathcal{C}_{\mathcal{B}_k}(\boldsymbol{\Theta}_k), \\ \boldsymbol{\Theta}_{k+1} &=& \boldsymbol{\Theta}_k + \boldsymbol{v}_{k+1}, ~k=1,2,\ldots, \end{array}\right\} \end{eqnarray}(5.6)

and is referred to as the momentum method.

Problem 5.7:

Recall, Newton's second law states

mass \(\times\) acceleration = force

For a particle moving on the cost surface, we have

\[\begin{array}{rcl} \text{force} &=& \text{frictional force}\\ &&~~ + \text{gravitational force driving the particle towards regions of lower cost} \end{array}\]

Mathematically, Newton's second law can be written as

\[ m \ddot{\boldsymbol{\Theta}} = - m\gamma\dot{\mathbf{\Theta}} -\nabla \mathcal{C}_{\mathcal{B}_k}(\mathbf{\Theta}), \]

where \(\gamma>0\) is the damping coefficient and \(m>0\) is the mass of the particle.

Considering the approximations

\[ \ddot{\boldsymbol{\Theta}} \approx \frac{ \mathbf{\Theta}_{k+1} -2\mathbf{\Theta}_k +\mathbf{\Theta}_{k-1} }{h^2} ~\text{and}~ \dot{\mathbf{\Theta}} \approx \frac{\mathbf{\Theta}_k-\mathbf{\Theta}_{k-1}}{h}, \]

for some sufficiently small \(h>0\), find the learning rate \(\eta\) and the momentum parameter \(\mu\), such that the Newton's second law reduces to the iterative formula (5.6) of the momentum method.

The momentum term can also be interpreted from a control-theoretic perspective. Treating \(\boldsymbol{v}_k\) as a control input, the term \(\mu \boldsymbol{v}_k\) acts as a feedback term with \(\mu\) playing the role of a feedback gain. This feedback regulates how strongly past velocity influences the current update, smoothing oscillations and accelerating convergence along consistent descent directions.

Problem 5.8:

Let \(\{\boldsymbol{\Theta}_k\}\) be computed using the momentum method, and suppose the gradient \(\nabla_{\boldsymbol{\Theta}} \mathcal{C}_{\mathcal{B}_k}(\boldsymbol{\Theta}_k) \approx \boldsymbol{g}\), for every \(k\), where \(\boldsymbol{g}\) is a constant vector. Show that

\[ \lim_{k\to \infty} \boldsymbol{v}_k \approx -\frac{\eta}{1-\mu}\,\boldsymbol{g}. \]

In other words, show that the momentum method reduces to the mini-batch gradient descent method with an effective learning rate \(\eta/(1-\mu)\) in regions of low curvature.

5.2.2 Nesterov Accelerated GD Method

The momentum method discussed in the above subsection evaluates the gradient of the cost function at the current iteration value \(\boldsymbol{\Theta}_k\) while updating the next iteration. Since the momentum is in the direction \(\mu \boldsymbol{v}_k\), the convergence can be further improved by evaluating the cost gradient at the look-ahead point

\[\widetilde{\boldsymbol{\Theta}}_k = \boldsymbol{\Theta}_k + \mu \boldsymbol{v}_k.\]

This method is called the Nesterov Accelerated Gradient Descent (NAGD) method, and is given by

\begin{eqnarray} \left.\begin{array}{lcl} \boldsymbol{v}_{k+1} &=& \mu \boldsymbol{v}_k - \eta \nabla_{\boldsymbol{\Theta}} \mathcal{C}_{\mathcal{B}_k}(\widetilde{\boldsymbol{\Theta}}_k), \\ \boldsymbol{\Theta}_{k+1} &=& \boldsymbol{\Theta}_k + \boldsymbol{v}_{k+1}, ~k=1,2,\ldots, \end{array}\right\} \end{eqnarray}(5.7)

Problem 5.9:

Let \(\Theta \in \mathbb{R}\) and consider the quadratic cost \(\mathcal{C}(\Theta) = \frac12 a \Theta^2\), \(a>0\). Let \((\Theta_k, v_k)\) be obtained using the Nesterov update rule (5.7).

Show that \((\Theta_k, v_k)\) satisfies the linear recursion
\[ \begin{pmatrix}\Theta_{k+1}\\ v_{k+1}\end{pmatrix} = M_{\mathrm{NAG}} \begin{pmatrix}\Theta_{k}\\ v_{k}\end{pmatrix}, \qquad M_{\mathrm{NAG}} = \begin{pmatrix}1-\eta a & \mu(1-\eta a)\\ -\eta a & \mu(1-\eta a)\end{pmatrix}. \]
Compare \(M_{\mathrm{NAG}}\) with the matrix \(M\) obtained for the classical momentum method on the same quadratic cost. Identify the extra factor contributed by the look-ahead point.
Fix \(a=1\), \(\eta=0.1\), and \(\mu=0.9\). Compute the eigenvalues and the spectral radius of both \(M\) and \(M_{\mathrm{NAG}}\) for these values, and verify that \(\rho(M_{\mathrm{NAG}}) < \rho(M)\), where \(\rho(M)\) is the spectral radius of the matrix M defined as
\[ \rho(M) := \max\{ |\lambda| : \lambda \text{ is an eigenvalue of } M \}. \]

Note:

For any iterative method of the form

\[ \boldsymbol{x}_{k+1} = M \boldsymbol{x}_k, \qquad k=0,1,2,\ldots, \]

where \(M\) is the iteration matrix. Let \(\lambda\) be an eigenvalue of \(M\) with \(|\lambda|=\rho(M)\). If \(\boldsymbol{x}_0\) has a nonzero component along the eigenspace of \(\lambda\), then

\[ \|\boldsymbol{x}_k\| \approx C\, \rho(M)^k, \]

for large values of \(k\), for some constant \(C>0\) depending on \(\boldsymbol{x}_0\). If \(\boldsymbol{x}_0\) has zero component along this eigenspace, the same asymptotic form holds with \(\rho(M)\) replaced by the largest \(|\lambda'|\) among the eigenvalues \(\lambda'\) for which \(\boldsymbol{x}_0\) has a nonzero component.

The following statement, however, holds for every initial vector \(\boldsymbol{x}_0\), with no such qualification: \(\boldsymbol{x}_k \to \boldsymbol{0}\) as \(k\to\infty\) if and only if \(\rho(M) < 1\). Further, whenever the geometric estimate above applies, a smaller spectral radius means a faster rate of convergence, and \(\rho(M)\) is the number that governs this rate.

5.2.3 Adaptive Step Size Methods

So far, we have kept the learning rate \(\eta\) constant for all the parameters and across all the iterations. This is the simplest form of a GD method, whereas the learning rate can depend on the iteration as

\[ \boldsymbol{\Theta}_{k+1} = \boldsymbol{\Theta}_k - \eta_k \nabla_{\boldsymbol{\Theta}} \mathcal{C}(\boldsymbol{\Theta}_k), k=0,1,2,\ldots. \]

In fact, it has been proved (omitted for this course) that the stochastic GD method

\[ \boldsymbol{\Theta}_{k+1} = \boldsymbol{\Theta}_k - \eta_k\, \nabla_{\boldsymbol{\Theta}} \ell_{i_k}(\boldsymbol{\Theta}_k), k=0,1,2,\ldots. \]

converges if the sequence \(\{\eta_k\}\) satisfies the Robbins–Monro conditions

\[ \sum_{k=0}^\infty \eta_k = \infty,~~\sum_{k=0}^\infty \eta_k^2 < \infty. \]

For instance, the sequence of learning rate

\[\eta_k = \frac{C_0}{C_1+k},\]

for some positive constants \(C_0\) and \(C_1\), satisfies the above two conditions (prove this!).

Further, to accelerate the convergence, it is also possible to vary the learning rate across parameters. In this subsection, we present some of the important such methods.

5.2.3.1 AdaGrad.

Each parameter \(\Theta_j\) has its own effective learning rate scaled by the accumulated squared gradients given by

\[ \Theta_{j,k+1} = \Theta_{j,k} - \frac{\eta}{\sqrt{G_{j,k} + \epsilon}} \, g_{j,k}, \]

where \(g_{j,k}\) is the \(j\)th component of the gradient \(\boldsymbol{g}_k = \nabla_\Theta \mathcal{C}(\Theta_k)\) at step \(k\),

\[ G_{j,k} = \sum_{i=0}^k g_{j,i}^2, \]

and \(\epsilon>0\) is a small constant added to avoid blow-up of the learning rate.

The main advantage of AdaGrad is its ability to assign a different effective learning rate to each parameter based on its gradients of the past iterations. Thus, parameters with large gradients in the past receive smaller learning rates while parameters with smaller gradients retain larger updates. This is particularly important for sparse-gradient problems. However, the disadvantage is that the accumulation \(G_{j,k}\) only increases and consequently, the learning rate only decreases.

5.2.3.2 Root Mean Square Propagation (RMSProp)

In order to address AdaGrad’s disadvantage of decay in learning rate, RMSProp uses an exponential moving average of squared gradients given by

\[ G_{j,k} = \rho G_{j,k-1} + (1-\rho) g_{j,k}^2, \]

where \(0<\rho<1\). In this update, the effective learning rate depends more strongly on the gradients of the recent iterations, whereas the older ones get less weightage.

5.2.3.3 Adaptive Moments (Adam).

We have seen that the AdaGrad updates the effective learning rates through cumulative history, whereas the RMSProp updates with a decaying memory.

The Adam method combines RMSProp with momentum by including the exponentially weighted moving averages of both gradients and squared gradients. The update rule for the Adam method is given by

\[\begin{array}{rcl} m_{j,k} &=& \beta_1 m_{j,k-1} + (1-\beta_1) g_{j,k}, \\ v_{j,k} &=& \beta_2 v_{j,k-1} + (1-\beta_2) g_{j,k}^2, \\ \Theta_{j,k+1} &=& \Theta_{j,k} - \frac{\eta}{\sqrt{\widehat{v}_{j,k}}+\epsilon}\, \widehat{m}_{j,k},~~k=0,1,2,\ldots, \end{array}\]

for \(\beta_1, \beta_2\in (0,1)\), and

\[ m_{j,-1} = 0,~~v_{j,-1} = 0, \]

where \(\widehat{m}_{j,k}\) and \(\widehat{v}_{j,k}\) are bias-corrected estimates given by

\[ \widehat{m}_{j,k} = \frac{m_{j,k}}{1-\beta_1^{k+1}},~~\widehat{v}_{j,k} = \frac{v_{j,k}}{1-\beta_2^{k+1}}. \]

Adam is the most widely used optimizer in deep learning practice.

The following problem clarifies the role of \(\widehat{m}\) and \(\widehat{v}\) through the concept of bias correction in statistical estimation.

Problem 5.10:

Let \(X\) be a random variable with mean \(\mu\) and let \(\{X_1,X_2,\ldots, X_N\}\) be independent and identically distributed samples from \(X\). Consider the exponential moving average

\[ \mu_0=0,~~\mu_{k} = \beta \mu_{k-1} + (1-\beta) X_{k},~~k=0,1,2,\ldots, N, \]

where \(0< \beta < 1\). Also define

\[ \widehat{\mu}_k = \frac{\mu_k}{1-\beta^{k+1}}. \]

Show that \(\mu_k\) is a biased estimator of \(\mu\).
Show that \(\widehat{\mu}_k\) is an unbiased estimator of \(\mu\).

5.2.3.4 AdamW

A commonly used variant of Adam is AdamW, which uses a decoupled weight-decay term. Formally, one may consider the regularized cost

\[ \mathcal{C}_\lambda(\boldsymbol{\Theta}) = \mathcal{C}(\boldsymbol{\Theta}) + \frac{\lambda}{2} \sum_j \Theta_j^2. \]

Note:

The details and motivation for the \(L^2\) penalty will be discussed in the next chapter.

In AdamW, the weight-decay term is applied separately from the adaptive gradient update. In componentwise form, the update is

\[ \Theta_{j,k+1} = \Theta_{j,k} - \eta\, \frac{\widehat m_{j,k}} {\sqrt{\widehat v_{j,k}}+\epsilon} - \eta\lambda\Theta_{j,k}, \]

where the bias-corrected estimates are obtained as in Adam, with the gradients computed without the penalty term.

AdamW is widely used in the training of Transformer-based models and large language models. In particular, it is the default optimizer in the current (at the time of writing these notes) Hugging Face Transformers training framework.