Optimization Algorithms from a Dynamics Perspective (I): From SGD to Momentum Acceleration

Analysis of Training Objectives

Assume the set of all training samples is \boldsymbol{S}, and the loss measure is L(\boldsymbol{x};\boldsymbol{\theta}), where \boldsymbol{x} represents a single sample and \boldsymbol{\theta} represents the optimization parameters. We can then construct the loss function: L(\boldsymbol{\theta}) = \frac{1}{|\boldsymbol{S}|}\sum_{\boldsymbol{x}\in\boldsymbol{S}} L(\boldsymbol{x};\boldsymbol{\theta}) \tag{1} The ultimate goal of training is to find a global optimum of L(\boldsymbol{\theta}) (where "optimum" here means "minimum").

GD and ODE

To achieve this goal, we can use the Gradient Descent (GD) method: \boldsymbol{\theta}_{n+1}=\boldsymbol{\theta}_{n} - \gamma \nabla_{\boldsymbol{\theta}} L(\boldsymbol{\theta}_{n}) \tag{2} where \gamma > 0 is called the learning rate, which here also happens to be the step size of the iteration (later we will see that the step size is not necessarily equal to the learning rate). There are many ways to understand gradient descent. Since we generally have \gamma \ll 1, we can rewrite it as: \frac{\boldsymbol{\theta}_{n+1}-\boldsymbol{\theta}_{n}}{\gamma} = - \nabla_{\boldsymbol{\theta}} L(\boldsymbol{\theta}_{n}) \tag{3} The left side approximates the derivative of \boldsymbol{\theta} (assuming it is a function of time t). Thus, we obtain an ODE (Ordinary Differential Equation) dynamical system: \dot{\boldsymbol{\theta}} = - \nabla_{\boldsymbol{\theta}} L(\boldsymbol{\theta}) \tag{4} Equation (2) is essentially an Euler method solution for (4). In this way, gradient descent is actually using the Euler method to solve the dynamical system (4). Since (4) is a conservative dynamical system, it can eventually converge to a fixed point (where \dot{\boldsymbol{\theta}}=0), and it can be proven that a stable fixed point is a local minimum (though not necessarily the global minimum).

From GD to SGD

Equation (2) is generally called "Batch Gradient Descent" because it requires all samples to calculate the gradient. This is a significant disadvantage: when there are thousands or millions of samples, the cost of each iteration is too high, potentially making it unfeasible. Therefore, we randomly select a subset \boldsymbol{R}\subseteq \boldsymbol{S} from \boldsymbol{S} and use only \boldsymbol{R} to calculate the gradient for a single iteration. We denote: L_{\boldsymbol{R}}(\boldsymbol{\theta}) = \frac{1}{|\boldsymbol{R}|}\sum_{\boldsymbol{x}\in\boldsymbol{R}} L(\boldsymbol{x};\boldsymbol{\theta}) \tag{5} Then equation (2) becomes: \boldsymbol{\theta}_{n+1}=\boldsymbol{\theta}_{n} - \gamma \nabla_{\boldsymbol{\theta}} L_{\boldsymbol{R}}(\boldsymbol{\theta}_{n}) \tag{6} Note that the minimum of L is our goal, while L_{\boldsymbol{R}} is a random variable, and \nabla_{\boldsymbol{\theta}} L_{\boldsymbol{R}}(\boldsymbol{\theta}_{n}) is merely an estimate of the original \nabla_{\boldsymbol{\theta}} L(\boldsymbol{\theta}_{n}). Can this approach yield reasonable results?

From SGD to SDE

Here we assume: \nabla_{\boldsymbol{\theta}} L(\boldsymbol{\theta}_{n}) - \nabla_{\boldsymbol{\theta}} L_{\boldsymbol{R}}(\boldsymbol{\theta}_{n})=\boldsymbol{\xi}_n \tag{7} follows a normal distribution with variance \sigma^2. Note that this is only an approximate description, mainly for semi-qualitative and semi-quantitative analysis. Under this assumption, stochastic gradient descent is equivalent to introducing Gaussian noise into the dynamical system (4): \dot{\boldsymbol{\theta}} = - \nabla_{\boldsymbol{\theta}} L(\boldsymbol{\theta}) + \sigma \boldsymbol{\xi} \tag{8} where \boldsymbol{\xi} follows a standard normal distribution. The original dynamical system was an ODE; now it has become an SDE (Stochastic Differential Equation), which we call the "Langevin Equation." Of course, the source of noise is not only the estimation error from random subsets; the learning rate of each iteration also introduces noise.

Under the Gaussian noise assumption, what is the solution to this equation? The solution to the original ODE is a deterministic trajectory. Now, due to the introduction of random noise, the solution is also random. The probability distribution at equilibrium can be solved as: P(\boldsymbol{\theta}) \sim \exp \left(-\frac{L(\boldsymbol{\theta})}{\sigma^2}\right) \tag{9} The derivation process can be found in general stochastic dynamics textbooks; we will simply use this result.

Heuristic Insights

From equation (8), we can derive some meaningful results. First, we see that in principle, \boldsymbol{\theta} is no longer a deterministic value but a probability distribution. Furthermore, the local minima of L(\boldsymbol{\theta}) now correspond exactly to the local maxima of P(\boldsymbol{\theta}). This indicates that if we execute gradient descent for an infinitely long time, theoretically \boldsymbol{\theta} can traverse all possible values and will have a higher probability of being in the various "pits" (minima) of L(\boldsymbol{\theta}).

\sigma^2 is the variance of the gradient. Obviously, this variance depends on the batch size. According to definition (7), the larger the batch size, the smaller the variance. In equation (9), a larger \sigma^2 means the graph of P(\boldsymbol{\theta}) is flatter, i.e., closer to a uniform distribution, allowing \boldsymbol{\theta} to wander everywhere. When \sigma^2 is smaller, the regions of the original local minima of L(\boldsymbol{\theta}) become more prominent, and \boldsymbol{\theta} is more likely to fall into a "pit" and not come out.

Based on this analysis, theoretically, we should start with a smaller batch size so that the noise variance \sigma^2 is larger. This makes the distribution closer to uniform, allowing the algorithm to traverse more regions. As the number of iterations increases, it gradually approaches the optimal region, at which point the variance should decrease to make the extrema more prominent. That is, if possible, the batch size should slowly increase with the number of iterations. This partially explains the results proposed by Google last year in "Don’t Decay the Learning Rate, Increase the Batch Size". However, increasing the batch size significantly increases computational costs, so we generally stop adjusting it after it reaches a certain amount.

Of course, as seen from the figures, when entering a stable descent region, the step size \gamma (learning rate) of each iteration should not exceed the width of the "pit." As \sigma^2 decreases, the pit becomes narrower, which also suggests that the learning rate should decrease as iterations increase. Additionally, a larger \gamma also partially contributes to noise; therefore, reducing \sigma^2 effectively implies that we should also reduce the learning rate.

The analysis result is:

When conditions permit, when using SGD, start with a small batch size and a large learning rate, then slowly increase the batch size and slowly decrease the learning rate.

As for the specific strategies for increasing or decreasing, that depends on your "alchemy" skills.

From First-Order to Second-Order

From the previous text, we know there is no difference in the iterative format between SGD and GD. Therefore, to seek acceleration for SGD, we only need to seek acceleration for GD and then replace the full gradient with the stochastic gradient. From the transition of equation (2) to (4), we know that GD transforms the problem of finding the minimum of L(\boldsymbol{\theta}) into the iterative solution of the ODE \dot{\boldsymbol{\theta}} = - \nabla_{\boldsymbol{\theta}} L(\boldsymbol{\theta}).

So, why must it be first-order? Would a second-order \ddot{\boldsymbol{\theta}} = - \nabla_{\boldsymbol{\theta}} L(\boldsymbol{\theta}) work?

To this end, we consider a general: \ddot{\boldsymbol{\theta}} + \lambda \dot{\boldsymbol{\theta}} = - \nabla_{\boldsymbol{\theta}} L(\boldsymbol{\theta}) \tag{10} This truly corresponds to a (Newtonian) mechanical system, where \lambda > 0 introduces an effect similar to friction. Let’s analyze the difference between such a system and the original first-order ODE (4) of GD.

First, from the perspective of fixed points, the stable fixed point to which (10) eventually converges (where \ddot{\boldsymbol{\theta}}=\dot{\boldsymbol{\theta}}=0) is indeed a local minimum of L(\boldsymbol{\theta}). Imagine a small ball rolling down from a mountaintop; it will automatically fall into a valley and rise again, but due to the existence of frictional resistance, it eventually stops in the valley. Note that unless the algorithm never stops, once it does stop, it must be in a valley (it could also be a peak or a saddle point, but the probability is relatively small). It can never stop on a hillside because on a hillside, potential energy still exists, which, according to the law of conservation of energy, can be converted into kinetic energy, so it will not stop.

Therefore, the convergence situation should be no different from first-order GD, so we only need to compare their convergence speeds.

GD + Momentum

We rewrite (10) equivalently as: \dot{\boldsymbol{\theta}}=\boldsymbol{\eta},\quad \dot{\boldsymbol{\eta}}=-\lambda \boldsymbol{\eta} - \nabla_{\boldsymbol{\theta}} L(\boldsymbol{\theta}) \tag{11}

We discretize \dot{\boldsymbol{\theta}} as: \dot{\boldsymbol{\theta}}\approx \frac{\boldsymbol{\theta}_{n+1}-\boldsymbol{\theta}_{n}}{\gamma} \tag{12} How should we handle \boldsymbol{\eta}? \boldsymbol{\eta}_n? No, (\boldsymbol{\theta}_{n+1}-\boldsymbol{\theta}_{n})/\gamma as the derivative at time n is only a first-order approximation (\mathcal{O}(\gamma)), while as the derivative at time n+1/2, it is a second-order approximation (\mathcal{O}(\gamma^2)). Therefore, more accurately: \frac{\boldsymbol{\theta}_{n+1}-\boldsymbol{\theta}_{n}}{\gamma}=\boldsymbol{\eta}_{n+1/2} \tag{13} Similarly, from the second part of equation (11), we derive the following result, which also has second-order accuracy: \frac{\boldsymbol{\eta}_{n+1/2}-\boldsymbol{\eta}_{n-1/2}}{\gamma}=-\lambda\left(\frac{\boldsymbol{\eta}_{n+1/2}+\boldsymbol{\eta}_{n-1/2}}{2}\right)- \nabla_{\boldsymbol{\theta}} L(\boldsymbol{\theta}_n) \tag{14} In summary, to pursue high precision, (\boldsymbol{\theta}_{n+1}-\boldsymbol{\theta}_{n})/\gamma is the derivative of \boldsymbol{\theta} at time n+1/2, (\boldsymbol{\eta}_{n+1/2}-\boldsymbol{\eta}_{n-1/2})/\gamma is the derivative of \boldsymbol{\eta} at time n, and (\boldsymbol{\eta}_{n+1/2}+\boldsymbol{\eta}_{n-1/2})/2=\boldsymbol{\eta}_n. They all have \mathcal{O}(\gamma^2) precision. Let: \boldsymbol{v}_{n+1}=\gamma\boldsymbol{\eta}_{n+1/2},\quad \beta = \frac{1-\lambda\gamma/2}{1+\lambda\gamma/2},\quad \alpha = \frac{\gamma^2}{1+\lambda\gamma/2} \tag{15} Then combining (13) and (14), we get: \begin{aligned} &\boldsymbol{\theta}_{n+1} = \boldsymbol{\theta}_{n} + \boldsymbol{v}_{n+1} \\ & \boldsymbol{v}_{n+1} = \beta\boldsymbol{v}_{n} - \alpha \nabla_{\boldsymbol{\theta}} L(\boldsymbol{\theta}_n) \end{aligned} \tag{16} This is the GD algorithm with Momentum. In mathematics, it also has a special name: the "Leapfrog Integration Method."

How Does It Accelerate?

Combining equations (15) and (16), we can observe how Momentum accelerates GD.

In GD (2), the learning rate is \gamma, the step size is also \gamma, and the precision is \mathcal{O}(\gamma). In Momentum, the learning rate is \alpha\approx \gamma^2, the step size is \gamma, and the precision is \mathcal{O}(\gamma^2)=\mathcal{O}(\alpha). Thus, after choosing a learning rate \alpha, under the same precision, Momentum actually advances with a step size of \sqrt{\alpha}, while pure GD advances with a step size of \alpha.

Since the learning rate is generally less than 1, \sqrt{\alpha} > \alpha. Therefore:

One of the principles of Momentum acceleration is that it allows the entire algorithm to advance with a larger step size under the same learning rate without losing precision.

Furthermore, as shown in the figure, if starting from point A, gradient descent will slowly descend to point B and finally stay at point C. Of course, if the noise or learning rate is large, it might cross point C and reach point D. But with momentum acceleration, (10) is a dynamical equation, and all concepts of Newtonian mechanics can be applied here. Starting from point A, the initial velocity is 0, then it slowly descends, potential energy is converted into kinetic energy, and after passing point B, it slowly rises, kinetic energy is converted back into potential energy. If the friction is small, it will still have kinetic energy upon reaching point C, allowing it to reach point D directly. This is guaranteed by the conservation of energy, even without noise. In SGD, one must rely on a large learning rate and small batch size (enhanced noise) to achieve a similar effect.

Thus, we can also say:

Momentum acceleration provides the dynamical possibility to "overshoot" suboptimal local minima.

So how should \lambda be chosen? Can’t we just let \lambda=0 or \beta=1?

As mentioned earlier, the \lambda > 0 term acts like friction to dissipate energy. Without this term, no matter how small the learning rate is (as long as it’s not zero), the Momentum algorithm will not stay at the local minimum but will keep moving. It’s like a pendulum that would swing forever without friction; this is determined by energy conservation. Energy conservation tells us that at the point of minimum energy (the minimum we desire), kinetic energy is at its maximum—meaning the velocity is fastest. Simply put, the algorithm wouldn’t stop. But with friction dissipating energy, energy is no longer conserved, and the pendulum eventually stops at the point of minimum energy. Therefore, introducing \lambda is necessary for the algorithm’s convergence, and from (15) we have \beta < 1. However, \lambda should not be too large; excessive friction would cause the motion to stop before reaching the minimum. To ensure the existence of the acceleration effect, we also have \beta > 0.

Finally, from the definition of \beta in (15), we see that when \lambda is fixed (i.e., the friction coefficient is fixed), if the learning rate \alpha decreases (meaning \gamma also decreases), then \beta should increase accordingly. The ratio of increase can be estimated simply. From (16), we get approximately 1-\lambda\sqrt{\alpha}=\beta, from which we can solve for \lambda, and then substitute the new \alpha to calculate the new \beta.

Thus, we obtain a reference tuning scheme for \beta in the SGD+Momentum optimizer:

When using SGD+Momentum, if you decrease the learning rate, you should slightly increase \beta. When the learning rate decreases from \alpha to r\alpha, \beta can be considered for an increase to 1 - (1-\beta)\sqrt{r}.

Nesterov Momentum

The Momentum algorithm essentially solves (10) numerically, but solving (10) is not limited to explicit iterative formats like (13) and (14); there are also implicit iterations. For example, we can approximate (10) as: \begin{aligned} &\frac{\boldsymbol{\theta}_{n+1}-\boldsymbol{\theta}_{n}}{\gamma} = \frac{\boldsymbol{\eta}_{n+1}+\boldsymbol{\eta}_{n}}{2}\\ &\frac{\boldsymbol{\eta}_{n+1}-\boldsymbol{\eta}_{n-1}}{2\gamma} = -\lambda \frac{\boldsymbol{\eta}_{n}+\boldsymbol{\eta}_{n-1}}{2} - \nabla_{\boldsymbol{\theta}} L\left(\frac{\boldsymbol{\theta}_{n+1}+\boldsymbol{\theta}_n}{2}\right) \end{aligned} \tag{17} Let: \boldsymbol{v}_{n+1}=\frac{\gamma}{2}(\boldsymbol{\eta}_{n+1}+\boldsymbol{\eta}_{n}),\quad \beta = 1-\lambda\gamma,\quad \alpha=\gamma^2 \tag{18} We get: \begin{aligned} &\boldsymbol{\theta}_{n+1} = \boldsymbol{\theta}_{n} + \boldsymbol{v}_{n+1} \\ & \boldsymbol{v}_{n+1} = \beta\boldsymbol{v}_{n} - \alpha \nabla_{\boldsymbol{\theta}} L\left(\frac{\boldsymbol{\theta}_{n+1}+\boldsymbol{\theta}_n}{2}\right) \end{aligned} \tag{19} This is an implicit iteration formula. Theoretically, to find \boldsymbol{\theta}_{n+1}, we would need to solve a system of nonlinear equations. But approximately, we can replace \boldsymbol{\theta}_{n+1} with \boldsymbol{\theta}_{n}+\beta \boldsymbol{v}_{n} to get: \begin{aligned} &\boldsymbol{\theta}_{n+1} = \boldsymbol{\theta}_{n} + \boldsymbol{v}_{n+1} \\ & \boldsymbol{v}_{n+1} = \beta\boldsymbol{v}_{n} - \alpha \nabla_{\boldsymbol{\theta}} L\left(\boldsymbol{\theta}_n + \frac{\beta}{2}\boldsymbol{v}_n\right) \end{aligned} \tag{20} If we replace the \beta/2 inside the parentheses with \beta, it becomes the standard GD algorithm with Nesterov momentum. However, I feel the above formula seems more reasonable, as the Nesterov algorithm intends to use the gradient at \boldsymbol{\theta}_{n+1} instead of \boldsymbol{\theta}_n to give the algorithm "look-ahead capability," but it is actually biased. Intuitively, using the gradient at (\boldsymbol{\theta}_n+\boldsymbol{\theta}_{n+1})/2 is more "balanced."

From an error analysis perspective, whether it is the standard Momentum or the Nesterov version, both are second-order algorithms with the same accuracy (only a constant-level difference). But theoretically, Nesterov is an implicit iteration with a wider stability region, so it usually yields better results.

How to implement equation (20) in frameworks like TensorFlow? Note that (20) involves first taking the derivative of L(\boldsymbol{\theta}) with respect to \boldsymbol{\theta}, and then replacing \boldsymbol{\theta} with \boldsymbol{\theta}_n + \frac{\beta}{2}\boldsymbol{v}_n (my version of Nesterov) or \boldsymbol{\theta}_n + \beta\boldsymbol{v}_n (standard Nesterov). This operation seems simple to us, but it is somewhat cumbersome in frameworks like TF.

Because although these frameworks have automatic differentiation, they are ultimately numerical calculation frameworks rather than symbolic calculation systems, so the result is only a numerical derivative. When we calculate the derivative of \boldsymbol{\theta}, the result is a tensor—some fixed numbers. Replacing \boldsymbol{\theta} with \boldsymbol{\theta}_n + \frac{\beta}{2}\boldsymbol{v}_n becomes difficult. It’s not impossible, but it’s not as straightforward as the Momentum version.

To facilitate implementation, we let (taking standard Nesterov as an example): \boldsymbol{\Theta}_n=\boldsymbol{\theta}_n + \beta\boldsymbol{v}_n Then we can obtain a new iteration formula: \begin{aligned} &\boldsymbol{\Theta}_{n+1} = \boldsymbol{\Theta}_{n} + \beta\boldsymbol{v}_{n+1} - \alpha \nabla_{\boldsymbol{\theta}} L\left(\boldsymbol{\Theta}_n\right) \\ & \boldsymbol{v}_{n+1} = \beta\boldsymbol{v}_{n} - \alpha \nabla_{\boldsymbol{\theta}} L\left(\boldsymbol{\Theta}_n\right) \end{aligned} This is how Nesterov is implemented in mainstream frameworks, such as Keras, which avoids the variable substitution process. As the iteration gets closer to the local minimum, the momentum \boldsymbol{v} will become smaller, so \boldsymbol{\Theta} and \boldsymbol{\theta} will eventually be equivalent—even a small error doesn’t matter, because our fundamental goal for using momentum is acceleration, and model selection still depends on the performance on the validation set.

Kramers Equation

The previous discussion only covered momentum acceleration for Batch Gradient Descent. Finally, let’s briefly analyze the case under Stochastic Gradient Descent. With the same assumptions as before, after introducing the stochastic gradient, we can consider that (10) becomes an equation with random force: \ddot{\boldsymbol{\theta}} + \lambda \dot{\boldsymbol{\theta}} = - \nabla_{\boldsymbol{\theta}} L(\boldsymbol{\theta})+\sigma\boldsymbol{\xi} \tag{21} This is called the "Kramers Equation." Like the Langevin equation, it is a core result in stochastic dynamics. When \lambda=0, the static solution can be written, and this static distribution can serve as a reference for the general case: P(\boldsymbol{\theta}, \boldsymbol{\eta}) \sim \exp\left(-\frac{\boldsymbol{\eta}^2/2 + L(\boldsymbol{\theta})}{\sigma^2}\right) \tag{22} where \boldsymbol{\eta}=\dot{\boldsymbol{\theta}}. The marginal distribution of \boldsymbol{\theta} is exactly equation (9). Therefore, the previous analysis of pure SGD holds true for Momentum and Nesterov algorithms as well.