From a Dynamical Systems Perspective on Optimizers (V): Why Shouldn't the Learning Rate Be Too Small?

从动力学角度看优化算法（五）：为什么学习率不宜过小？ 苏剑林 · 2020-10-10 · Machine Learning differential equation dynamics gradient optimizer

The theme of this post is "why we need a finite learning rate." By "finite" here, we mean neither too large nor too small — moderate is fine. It's not hard to understand why too large a learning rate is problematic, since it can cause the algorithm to diverge, but why is too small a learning rate also undesirable? An easy-to-understand answer is that too small a learning rate requires too many iteration steps, which is an unnecessary waste; so from the perspective of "saving energy" and "speeding things up," we don't want to use an overly small learning rate. But if we set aside considerations of compute and time, is an arbitrarily small learning rate then acceptable? A paper recently posted on Arxiv by Google, Implicit Gradient Regularization, attempts to answer exactly this question. It points out that a finite learning rate implicitly introduces a gradient penalty term into the optimization process, and that this gradient penalty term is helpful for improving generalization performance. Therefore, even setting aside compute and time, one should not use an overly small learning rate.

This blog has discussed gradient penalties on several occasions before. In the posts A Brief Discussion on Adversarial Training: Meaning, Methods, and Reflections (with Keras Implementation) and Random Thoughts on Generalization: From Random Noise and Gradient Penalties to Virtual Adversarial Training, we analyzed how adversarial training is, to some extent, equivalent to a gradient penalty on the input; while the Flooding trick introduced in Do We Really Need to Push the Training Loss Down to Zero? is equivalent to a gradient penalty on the parameters. In general, whether applied to the input or to the parameters, gradient penalties help improve generalization ability to some degree. more

The Direction of Steepest Descent

Like this series of posts, this paper treats the optimization process as solving a differential equation. Recalling our earlier post From a Dynamical Systems Perspective on Optimizers (III): A More Holistic View, let the loss function be $L(\boldsymbol{\theta})$, and let us treat $\boldsymbol{\theta}$ as a trajectory $\boldsymbol{\theta}(t)$ varying along some time parameter $t$. Now let's consider its rate of change:

\begin{equation}\frac{d}{dt}L(\boldsymbol{\theta}(t))=\left\langle\nabla_{\boldsymbol{\theta}}L(\boldsymbol{\theta}(t)),\, \dot{\boldsymbol{\theta}}(t)\right\rangle\end{equation}

We want $L(\boldsymbol{\theta}(t))$ to decrease over time (smaller loss is better), so we want the above expression to be less than 0. When the norm $\Vert\dot{\boldsymbol{\theta}}(t)\Vert$ is fixed, the right-hand side of the above equation attains its minimum in the direction opposite to the gradient $-\nabla_{\boldsymbol{\theta}}L(\boldsymbol{\theta}(t))$, which is why we say the negative gradient direction is the direction of steepest descent. For simplicity, we can directly set

\begin{equation}\dot{\boldsymbol{\theta}}(t) = -\nabla_{\boldsymbol{\theta}}L(\boldsymbol{\theta}(t))\triangleq - \boldsymbol{g}(\boldsymbol{\theta}(t))\label{eq:odes}\end{equation}

Then solving for the parameters $\boldsymbol{\theta}$ reduces to solving the above system of ordinary differential equations — this is the basic starting point of the "From a Dynamical Systems Perspective on Optimizers" series.

The Regularization Hidden in the Learning Rate

However, the practical issue is that we cannot actually solve the system of differential equations $\eqref{eq:odes}$ exactly; we can only use numerical iteration. Using the simplest method, Euler's method, for instance, we get

\begin{equation}\boldsymbol{\theta}_{t+\gamma} = \boldsymbol{\theta}_{t} - \gamma \boldsymbol{g}(\boldsymbol{\theta}_t)\label{eq:gd}\end{equation}

This is, in fact, the most basic form of gradient descent, where $\gamma$ is what we usually call the learning rate. The equation above is essentially a difference equation.

It's easy to imagine that starting from $t=0$, the resulting points $\boldsymbol{\theta}_{\gamma},\boldsymbol{\theta}_{2\gamma},\boldsymbol{\theta}_{3\gamma},\cdots$ will deviate somewhat from the exact solution $\boldsymbol{\theta}(\gamma),\boldsymbol{\theta}(2\gamma),\boldsymbol{\theta}(3\gamma),\cdots$ of the system of equations $\eqref{eq:odes}$. How do we measure the extent of this deviation? Here's one way to think about it: $\boldsymbol{\theta}_{\gamma},\boldsymbol{\theta}_{2\gamma},\boldsymbol{\theta}_{3\gamma},\cdots$ is actually itself the exact solution of some differential equation system similar to $\eqref{eq:odes}$, except that the corresponding $\boldsymbol{g}(\boldsymbol{\theta}(t))$ has been replaced with some new $\tilde{\boldsymbol{g}}(\boldsymbol{\theta}_t)$. We just need to compare the difference between $\tilde{\boldsymbol{g}}(\boldsymbol{\theta}_t)$ and $\boldsymbol{g}(\boldsymbol{\theta}(t))$.

After some derivation, if we keep only terms up to first order in $\gamma$, we obtain

\begin{equation}\tilde{\boldsymbol{g}}(\boldsymbol{\theta}_t) = \boldsymbol{g}(\boldsymbol{\theta}_t) + \frac{\gamma}{4}\nabla_{\boldsymbol{\theta}}\Vert \boldsymbol{g}(\boldsymbol{\theta}_t)\Vert^2 = \nabla_{\boldsymbol{\theta}}\left(L(\boldsymbol{\theta}_t) + \frac{1}{4}\gamma\Vert \nabla_{\boldsymbol{\theta}} L(\boldsymbol{\theta}_t)\Vert^2\right)\end{equation}

We leave the derivation to the next section. As we can see, this is essentially equivalent to adding a regularization term in the form of a gradient penalty $\frac{1}{4}\gamma\Vert \nabla_{\boldsymbol{\theta}} L(\boldsymbol{\theta})\Vert^2$ to the loss function, and gradient penalty terms help the model reach flatter regions, which is beneficial for improving generalization performance. In other words, the discretized iteration process implicitly introduces a gradient penalty term, which actually helps the model's generalization. And if $\gamma\to 0$, this implicit penalty weakens or even vanishes.

The conclusion, therefore, is that the learning rate should not be too small — a larger learning rate not only speeds up convergence but also helps improve the model's generalization ability. Of course, some readers might wonder: if I directly add the gradient penalty to the loss, could I then just use a sufficiently small learning rate? In theory, yes — the original paper refers to this approach of adding the gradient penalty directly into the loss as "explicit gradient regularization."

From Difference Equation to Differential Equation

To convert a difference equation into a differential equation, we can use the ordinary "perturbation method," which this blog has briefly introduced before (see the tag "perturbation"). However, a more elegant solution is to work directly with the series expansion of operators, as in the earlier post The Art of Operators: Differences, Derivatives, and Bernoulli Numbers.

Let's expand $\boldsymbol{\theta}_{t+\gamma}$ with a Taylor series:

\begin{equation}\boldsymbol{\theta}_{t+\gamma}=\boldsymbol{\theta}_{t}+\gamma \dot{\boldsymbol{\theta}}_{t} + \frac{1}{2}\gamma^2\ddot{\boldsymbol{\theta}}_{t} + \frac{1}{6}\gamma^3\dddot{\boldsymbol{\theta}}_{t} + \cdots\end{equation}

If we denote the operation of differentiating with respect to $t$ as $D$, then the above expression is actually

\begin{equation}\boldsymbol{\theta}_{t+\gamma} = \left(1+\gamma D + \frac{1}{2}\gamma^2 D^2 + \frac{1}{6}\gamma^3 D^3 + \cdots\right)\boldsymbol{\theta}_{t} = e^{\gamma D}\boldsymbol{\theta}_{t}\end{equation}

So the difference equation $\eqref{eq:gd}$ can be written as

\begin{equation}\left(e^{\gamma D} - 1\right)\boldsymbol{\theta}_{t} = - \gamma \boldsymbol{g}(\boldsymbol{\theta}_t)\end{equation}

Just as with ordinary algebraic operations, we have

\begin{equation}\begin{aligned} D\boldsymbol{\theta}_{t} =& - \gamma \left(\frac{D}{e^{\gamma D} - 1}\right)\boldsymbol{g}(\boldsymbol{\theta}_t)\\ =& - \left(1 - \frac{1}{2}\gamma D + \frac{1}{12}\gamma^2 D^2 - \frac{1}{720}\gamma^4 D^4 + \cdots\right)\boldsymbol{g}(\boldsymbol{\theta}_t) \end{aligned}\end{equation}

The left-hand side of this equation is just $\dot{\boldsymbol{\theta}}_{t}$, so the right-hand side gives us the expression for $-\tilde{\boldsymbol{g}}(\boldsymbol{\theta}_t)$. Keeping terms up to first order gives

\begin{equation} - \left(1 - \frac{1}{2}\gamma D\right)\boldsymbol{g}(\boldsymbol{\theta}_t) = - \boldsymbol{g}(\boldsymbol{\theta}_t) + \frac{1}{2}\gamma \frac{d}{dt}\boldsymbol{g}(\boldsymbol{\theta}_t) = - \boldsymbol{g}(\boldsymbol{\theta}_t) + \frac{1}{2}\gamma \nabla_{\boldsymbol{\theta}}\boldsymbol{g}(\boldsymbol{\theta}_t)\dot{\boldsymbol{\theta}}_t \end{equation}

which is

\begin{equation}\begin{aligned} \dot{\boldsymbol{\theta}}_{t} =& - \boldsymbol{g}(\boldsymbol{\theta}_t) + \frac{1}{2}\gamma \nabla_{\boldsymbol{\theta}}\boldsymbol{g}(\boldsymbol{\theta}_t)\dot{\boldsymbol{\theta}}_t\\ =&- \boldsymbol{g}(\boldsymbol{\theta}_t) + \frac{1}{2}\gamma \nabla_{\boldsymbol{\theta}}\boldsymbol{g}(\boldsymbol{\theta}_t)\left[- \boldsymbol{g}(\boldsymbol{\theta}_t) + \frac{1}{2}\gamma \nabla_{\boldsymbol{\theta}}\boldsymbol{g}(\boldsymbol{\theta}_t)\dot{\boldsymbol{\theta}}_t\right]\\ =&- \boldsymbol{g}(\boldsymbol{\theta}_t) - \frac{1}{2}\gamma \nabla_{\boldsymbol{\theta}}\boldsymbol{g}(\boldsymbol{\theta}_t)\boldsymbol{g}(\boldsymbol{\theta}_t)\quad\text{(omit second-order terms)}\\ =&- \boldsymbol{g}(\boldsymbol{\theta}_t) - \frac{1}{4}\gamma \nabla_{\boldsymbol{\theta}}\Vert\boldsymbol{g}(\boldsymbol{\theta}_t)\Vert^2 \end{aligned}\end{equation}

So, to first order, $\tilde{\boldsymbol{g}}(\boldsymbol{\theta}_t)=\boldsymbol{g}(\boldsymbol{\theta}_t) + \frac{1}{4}\gamma \nabla_{\boldsymbol{\theta}}\Vert\boldsymbol{g}(\boldsymbol{\theta}_t)\Vert^2$ — which completes the derivation.

A Perfunctory Little Summary

The development and popularization of deep learning owes much to the successful application of gradient-descent-based optimizers, yet why gradient descent works so well remains without a deep explanation. Through the process of "model training," many researchers have accumulated a certain amount of unexplained "tricks of the trade" — such as how large the batch size should be, or how to tune the learning rate — and presumably everyone has their own rules of thumb here.

Regarding the phenomenon that "the learning rate should not be too small," everyone probably has some intuitive sense of it, and it's often taken for granted as "common knowledge," with little thought given to the underlying reasons. This Google paper offers a possible explanation for this phenomenon: a moderate, rather than overly small, learning rate introduces an implicit gradient penalty term into the optimization process, which helps convergence toward flatter regions. I think the analysis process here is well worth studying and referring to.