SGD with Momentum

Recall from previous lecture, the running time of GD on a strongly convex (PL-condition) function depends only on condition number $\kappa$. When the condition number is high, convergence can become slow.

So how can we speed up gradient descent when the condition is high? There are three common solutions:

• Momentum (this lecture)
• Adaptive learning rates
• Preconditioning

We introduce momentum here. A direct analysis would be messy, so we use a very simple example to give some intuition.

Simple Quadratic Function¶

The simplest possible setting with a non-1 condition number is a 2D quadratic. To get a high condition number, we want in the Hessian matrix, the biggest (bigger) value to be very big and the smallest (smaller) to be very small. Consider the following example with $a \gt b$. Here $a$ is just the curvature of the first dimension and $b$ is the curvature of the second dimension.

The second derivative matrix is the constant

By definition of a condition number in linear algebra, we know $\kappa = \frac a b$. By its definition in ML optimization problem, we know $\kappa = \frac L \mu$, so we can just set $a = L, b = \mu$, so

With this, we can write our update step as

If we run the update $T$ steps, we will have (from now on, $T$ will most likely denote number of iterations instead of transpose.)

Feed it into $f$, we have

Therefore, the final value of $f(w_t)$ will be dominated by the exponential term: one of $(1 - \alpha L)^{2T}$ or $(1 - \alpha \mu)^{2T}$

To minimize $f(w_t)$, we have to minimize $| 1 - \alpha L|$ and $|1 - \alpha \mu|$ at the same time. That is to minimize $\max(| 1 - \alpha L|, |1 - \alpha \mu|)$ Note if we just look at the first dimension, so we just minimize $| 1 - \alpha L|$, we set $\alpha = \frac 1 L$. With $L \gt \mu$, this will always be a small number. Therefore, for the first dimension with a respective larger curvature $L$, we want a smaller learning rate. On the other hand, if we only look at the second dimension, where the respective curvature is the smaller $\mu$, we should want a higher learning rate.

Now that we have to minimize both at the same time, we are forced to choose something in the middle. We will always reach minimum when we have $\alpha L - 1 = 1 - \alpha \mu$, so $\alpha = \frac 2 {L + \mu}$. Substitute this $\alpha$ in,

As we said, the bigger of these two terms will dominate the final value of $f(w_T)$, so

We know $1-x \approx e^{-x}$ around $x = 1$, so (here $1 - \frac{2}{\kappa + 1} \approx 1$ with $\kappa$ being large)

Therefore, even in the simplest setting, we can’t get rid of this $\frac {} {\kappa + 1}$ term.

Polyak Momentum¶

We want to sort of detect high vs low curvature during GD. The idea is:

• If this dim has high curvature, we tend to overshoot, so in next step, gradient will point in the opposite direction
• If this dim has low curvature, we are more likely to stay in the same direction, which means the step size is too small

Therefore, we want to make steps smaller when gradients reverse sign and larger when gradients are consistently in the same direction. Polyak momentum does this

The intuition is that

• If the current gradient is in the same direction as the previous step, move a little further in the same direction
• If it’s in the opposite direction, move a little less far

This is equivalent to

Stage Transition Matrix¶

Go back to our simple example, denote $A = \begin{bmatrix} L & 0 \\ 0 & \mu \end{bmatrix}$, so

We can write this update process as a matrix operation too. The matrix here will be a block matrix, where each entry is actually a $2 \times 2$ matrix.

This block matrix in whole is actually a $4 \times 4$ matrix that transforms a vector in $\R^4$ to $\R^4$, so if we write in the basis form

And if we find the eigenvalues of this block matrix, denote them as $\lambda_1, \lambda_2, \lambda_3, \lambda_4$, we can write

As before, what dominates the value of $f(w_t)$ will be the biggest among the exponential $\lambda^T$ term. Therefore, we want to minimize all these eigenvalues at the same time.

In addition, recall the square term in $f(w)$

Therefore, we actually only care about the magnitude of $w_T$. That is, we want to minimize the magnitude of all these eigenvalues at the same time.

Analyzing Eigenvalues¶

We start to analyze the eigenvalues of this block matrix.

We write this block matrix out:

Since we are analyzing the eigenvalues and the basis are not of any importance, we swap the 2nd column with 3rd column and swap 2nd row with 3rd row. This is equivalent to swap the 2nd and 3rd basis in domain vector space and also swap 2nd and 3rd basis in the codomain space. Again, this is fine because we only care about the eigenvalues.

We write this new matrix also in block matrix form,

This new block matrix is also in diagonal form, so to solve for its eigenvalues, we just have to solve for $B$’s eigen values.

Recall that we want to minimize all of eigen values’ norms all at the same time. We achieve this when they all have the same norm. For $B$ specifically, it has two eigenvalues $\lambda_1, \lambda_2$ and we want $|\lambda_1| = |\lambda_2|$.

Write out the characteristic polynomial of matrix $B$

Solve this quadratic equation, we have

Two solutions to this equation have the same norm when we have them to be the same or they are complex numbers. That is when

To find the exact value of $|\lambda|$, we don’t have to go through the process of finding norm of a complex number. In fact, we just recall that product of all eigenvalues is equal to the determinant of this matrix, so

Therefore, to minimize $|\lambda|$, we actually need to minimize $\sqrt \beta$. So we have this new linear optimization problem to solve:

This is a special case of Karush–Kuhn–Tucker conditions (KTT). When we solve a linear programming problem of $k$ variables and $n$ inequalities, $k$ of these $n$ inequalities will actually hold as equality. In this case, this minimization problem is solved when 2 of these 2 inequality constraints achieve equality. . So we solve

which gives us

Since we have $L \gt \mu$, to make the two absolute values to be equal, we must have

Solve for both $\alpha$ and $\beta$, we have

Recall that the norm of $w_t$ will be dominated with the one with max eigenvalue ($C = \max \|c\|$)

Therefore, we have ($C' = \frac {\max(L, \mu)} 2 C^2$)

Therefore, our function converges to a given interval $\epsilon$ (remember this is a quadratic function), when we have $T$ satisfies the following condition

So

Recall a normal GD convergence rate is

Therefore, we have shown that using momentum with GD on this simple example of quadratic function has a faster convergence rate than a vanilla GD. The result does generalize to more general cases with other kinds of functions.

When we use momentum with SGD, it is not guaranteed that it gives a better result, but people still use it.

Nesterov Momentum¶

One disadvantage of Polyak momentum is that the momentum term is not guaranteed to point to the right direction. Also, it is only guaranteed to have this nice acceleration for quadratics. Therefore, we introduce Nesterov Momentum, which works for general strongly convex objectives.

Polyak:

Nesterov:

Difference: instead of calculating the momentum term at the current position, we pretend to have already taken one step and calculate the momentum term there.