Momentum (ML)

Nice Distill blog: Why Momentum Really Works

A common modification of Stochastic Gradient Descent is to add a momentum term. We update the parameters with a weighted combination of the gradient computed from the current batch and the direction moved in the previous step:

$$\begin{array}{rl}{\mathbf{m}}_{t+1}& ⟵\beta \cdot {\mathbf{m}}_t+(1-\beta )\sum _{i\in {\mathcal{B}}_t}\frac{\partial {\ell }_i[{\varphi }_t]}{\partial \varphi }\\ {\varphi }_{t+1}& ⟵{\varphi }_t-\alpha \cdot {\mathbf{m}}_{t+1}\end{array}$$

where:

• ${\mathbf{m}}_t$ is the momentum (which drives the update at iteration $t$)
• $\beta \in [0,1)$ is a decay parameter that controls the degree to which the gradient is smoothed over time, $\alpha$ is the learning rate.
• $\alpha$ is the learning rate
• $\frac{\partial {\ell }_i[{\varphi }_t]}{\partial \varphi }$ is the gradient of the loss function for data point $i$, with respect to the model parameters $\varphi$, evaluated at the current parameters ${\varphi }_t$.

Intuition

Essentially, we “remember” the direction that the past gradients took us in, and use a weighted average of the average past direction and the current gradient, resulting in a smoother path.

If we roll out the momentum expression, we can see that older gradients contribute less and less (exponential decay).

The recursive formulation of the momentum calculation means that the gradient step is an infinite weighted sum of all the previous gradients, where the weights get smaller as we move back in time. The effective learning rate increases if all these gradients are aligned over multiple iterations but decreases if the gradient direction repeatedly changes as the terms in the sum cancel out. The overall effect is a smoother trajectory and reduced oscillatory behavior in valleys.

MIT 6.036 Notes

We can use methods like Running Averages to describe strategies for computing $\eta$. Momentum is a simple method that does this by “averaging” recent gradient updates, so that if they have been bouncing back and forth in some direction, we take out that component of the motion. For momentum, we have:

$$\begin{array}{rl}{V}_0& =0\\ V_t& =\gamma V_{t-1}+\eta {\nabla }_{W}J(W_{t-1})\\ W_t& =W_{t-1}-V_t\end{array}$$

This can also be written as:

$$\begin{array}{rl}{M}_0& =0\\ M_t& =\gamma M_{t-1}+(1-\gamma ){\nabla }_{W}J(W_{t-1})\\ W_t& =W_{t-1}-{\eta }^{\prime }{M}_t\end{array}$$

These two forms are equivalent if we let $\eta ={\eta }^{\prime }(1-\gamma )$. Essentially, $V_t$ (or $M_t$) is a moving average of the gradient. We are doing an update with step size ${\eta }^{\prime }$ on a moving average of the gradients with parameter $\gamma$.

$V_t$ will be bigger in dimensions that consistently have the same sign for ${\nabla }_W$ and smaller for those that don’t. Of course we now have two parameters to set ($\eta$ and $\gamma$), but the hope is that the algorithm will perform better overall, so it will be worth trying to find good values for them. Often γ is set to be something like 0.9.