Viewing Optimization Algorithms from a Dynamical Systems Perspective (II): Adaptive Learning Rate Algorithms

从动力学角度看优化算法（二）：自适应学习率算法 苏剑林 · 2018-12-20 · Mathematics differential equation dynamics numerical computation optimizer

In Viewing Optimization Algorithms from a Dynamical Systems Perspective (I): From SGD to Momentum Acceleration], we proposed that SGD-style optimization algorithms actually correspond to numerical methods for solving ordinary differential equations (ODEs), and from this correspondence one can quite naturally analyze the convergence properties of SGD, the principle behind momentum acceleration, and so on.

In this article, we continue along the same line of thought, aiming to understand adaptive learning rate algorithms.

RMSprop

Let's start with a very classic adaptive learning rate optimization algorithm: RMSprop. Although RMSprop was not the earliest adaptive learning rate algorithm proposed, it is a remarkably practical one, and it serves as the cornerstone of more comprehensive algorithms such as Adam. Through it, we can observe how adaptive learning rate algorithms actually work.

Overview of the Algorithm

Standard gradient descent looks like this:

$$\begin{equation}\boldsymbol{\theta}_{n+1}=\boldsymbol{\theta}_{n} - \gamma \nabla_{\boldsymbol{\theta}} L(\boldsymbol{\theta}_{n})\end{equation}$$

Clearly, $\gamma$ here is a hyperparameter—the learning rate—which may need to be adjusted differently at different stages.

RMSprop, on the other hand, is

$$\begin{equation}\begin{aligned}\boldsymbol{g}_{n+1} =& \nabla_{\boldsymbol{\theta}} L(\boldsymbol{\theta}_{n})\\ \boldsymbol{G}_{n+1}=&\lambda \boldsymbol{G}_{n} + (1 - \lambda) \boldsymbol{g}_{n+1}\otimes \boldsymbol{g}_{n+1}\\ \boldsymbol{\theta}_{n+1}=&\boldsymbol{\theta}_{n} - \frac{\tilde{\gamma}}{\sqrt{\boldsymbol{G}_{n+1} + \epsilon}}\otimes \boldsymbol{g}_{n+1} \end{aligned}\end{equation}$$more

Analysis of the Algorithm

Comparing this with vanilla SGD, we can see that in the update for $\boldsymbol{\theta}$, RMSprop replaces the originally scalar learning rate $\gamma$ with a vector

$$\begin{equation}\boldsymbol{\gamma}=\frac{\tilde{\gamma}}{\sqrt{\boldsymbol{G}_{n+1} + \epsilon}}\end{equation}$$

If we regard this vector as a learning rate as well, then RMSprop has found a scheme that assigns a different learning rate to each component of the parameters.

This adjustment of the learning rate is achieved through the factor $\frac{1}{\sqrt{\boldsymbol{G}_{n+1} + \epsilon}}$, where $\boldsymbol{G}_{n+1}$ is the moving average of the squared gradient. Fundamentally, the "moving average" operation just makes the training process smoother—it is not the actual mechanism behind the adjustment. The main component doing the work is the "gradient" itself, i.e., we can use the magnitude of the gradient to adjust the learning rate.

Adaptive Learning Rate

Why can the magnitude of the gradient be used to adjust the learning rate? The idea behind this is actually quite simple.

Minima and ODEs

Without further ado, for simplicity let's start with a one-dimensional example: suppose we want to find a minimum of $L(\theta)$. We introduce a virtual time parameter $t$ and convert the problem into an ODE:

$$\begin{equation}\frac{d\theta}{dt}=\dot{\theta} = - L'(\theta)\end{equation}$$

It is not hard to see that a minimum of $L(\theta)$ is precisely a stable fixed point of this equation. Starting from an arbitrary $\theta_0$, we numerically solve this ODE, and we would expect it to eventually converge to this fixed point, thereby giving us a minimum.

The simplest Euler method approximates $\frac{\theta_{t+\gamma}-\theta_t}{\gamma}$ with $\dot{\theta}$, giving us

$$\begin{equation}\frac{\theta_{t+\gamma}-\theta_t}{\gamma} = - L'(\theta_t)\end{equation}$$

that is,

$$\begin{equation}\theta_{t+\gamma} = \theta_t - \gamma L'(\theta_t)\end{equation}$$

This is exactly gradient descent: $\theta_{t+\gamma}$ corresponds to $\theta_{n+1}$, and $\theta_t$ corresponds to $\theta_n$, meaning each step advances by an amount $\gamma$.

The Idea of Variable Learning Rates

The question is: what value should we choose for $\gamma$? Of course, from the perspective of "approximating $\frac{\theta_{t+\gamma}-\theta_t}{\gamma}$ with $\dot{\theta}$," the smaller $\gamma$ is, the more accurate the approximation. But the smaller $\gamma$ is, the more iterations are needed, meaning greater computational cost. So while "smaller is better" is ideal, it isn't practical.

Thus, the most sensible approach is: use just enough at each step. But how do we know when it's "enough"?

Since we are approximating $\dot{\theta}$ with $\frac{\theta_{t+\gamma}-\theta_t}{\gamma}$, we must analyze the degree of approximation. By the Taylor series, we have

$$\begin{equation}\theta_{t+\gamma} = \theta_t + \gamma \dot{\theta}_t + \mathcal{O}(\gamma^2)\end{equation}$$

In our case we have $\dot{\theta} = - L'(\theta)$, so we get

$$\begin{equation}\theta_{t+\gamma} = \theta_t - \gamma L'(\theta_t) + \mathcal{O}(\gamma^2)\end{equation}$$

We can expect that when $\gamma$ is small, the error term $\mathcal{O}(\gamma^2) < \gamma \left|L'(\theta_t)\right|$—that is, under certain conditions, $\gamma \left|L'(\theta_t)\right|$ itself serves as a measure of the error term. If we keep $\gamma \left|L'(\theta_t)\right|$ bounded within a certain range, then the error is controlled as well. Namely:

$$\begin{equation}\gamma \left|L'(\theta_t)\right|\leq \tilde{\gamma}\end{equation}$$

where $\tilde{\gamma}$ is a constant, and we could even simply set $\gamma \left|L'(\theta_t)\right|=\tilde{\gamma}$ (for now ignoring the possibility of $L'(\theta_t)=0$, in order to first observe the core idea overall), giving us

$$\begin{equation}\gamma = \frac{\tilde{\gamma}}{\left|L'(\theta_t)\right|}\end{equation}$$

In this way we have used the gradient to adjust the learning rate.

Applying a Moving Average

Readers might object: substituting $\gamma = \tilde{\gamma} / \left|L'(\theta_t)\right|$ back into the original iteration gives us just

$$\begin{equation}\theta_{t+\tilde{\gamma} / \left|L'(\theta_t)\right|} = \theta_t - \tilde{\gamma}\cdot\text{sign}\big[L'(\theta_t)\big]\end{equation}$$

—which only uses the sign information of the gradient. Isn't that a waste? (It's overly trivial: no matter how large or small the gradient is, each iteration of $\theta$ moves a fixed step length.)

Note, however, that from the perspective of solving an ODE, there's actually nothing wrong with this, because the solution of the ODE is a trajectory $(t,\theta(t))$. In this scheme, although $\theta$ becomes trivial, $t$ becomes non-trivial instead—it's as if the roles of $t,\theta$ have simply been swapped, so it's still reasonable. However, if what we care about is the optimization problem itself, i.e., finding the minimum of $L(\theta)$, then the above formula is indeed somewhat too trivial, because if every iteration of $\theta$ moves only a fixed step length, it becomes something like a grid search, which is quite inefficient.

So, to improve on this trivial behavior while still retaining the feature of using the gradient to adjust the learning rate, we can take a "moving average" of the gradient, giving us

$$\begin{equation}\begin{aligned}G_{t+\tilde{\gamma}}=&\lambda G_{t} + (1 - \lambda) |L'(\theta_t)|^2\\ \gamma =& \frac{\tilde{\gamma}}{\sqrt{G_{t+\tilde{\gamma}} + \epsilon}}\\ \theta_{t+\gamma} =& \theta_t - \gamma L'(\theta_t) \end{aligned}\end{equation}$$

Here $\lambda$ is a constant close to but less than 1. This makes $G_t$ fairly stable within a certain range, while to some extent still preserving the characteristics of the gradient $L'(\theta_t)$ itself. So using it to adjust the learning rate is a rather clever move. To avoid notational confusion caused by $t+\tilde{\gamma},t+\gamma$, let's uniformly use $n,n+1$ to denote the subscript, giving us:

$$\begin{equation}\begin{aligned}G_{n+1}=&\lambda G_{n} + (1 - \lambda) |L'(\theta_n)|^2\\ \gamma =& \frac{\tilde{\gamma}}{\sqrt{G_{n+1} + \epsilon}}\\ \theta_{n+1} =& \theta_n - \gamma L'(\theta_n) \end{aligned}\end{equation}\label{eq:rmsprop-1}$$

This is exactly the RMSprop algorithm mentioned at the beginning.

There is also another important reason for using a moving average: what we actually want to estimate is the mean of the squared gradient over the full dataset, but we can only compute it over one batch at a time (the so-called mini-batch gradient descent). This means the result computed each time is actually biased, and the moving average helps correct this bias to some extent.

Analysis of the High-Dimensional Case

All the discussion above was for the one-dimensional case. How do we generalize to the multi-dimensional case?

Readers might think it's simple: just replace scalars with vectors, right? It's not quite that simple, because there are at least two reasonable choices for generalizing $\eqref{eq:rmsprop-1}$ to higher dimensions:

$$\begin{equation}\begin{aligned}G_{n+1}=&\lambda G_{n} + (1 - \lambda) \Vert \nabla_{\boldsymbol{\theta}}L(\boldsymbol{\theta}_n)\Vert^2\\ \gamma =& \frac{\tilde{\gamma}}{\sqrt{G_{n+1} + \epsilon}}\\ \boldsymbol{\theta}_{n+1} =& \boldsymbol{\theta}_n - \gamma \nabla_{\boldsymbol{\theta}}L(\boldsymbol{\theta}_n) \end{aligned}\end{equation}$$

or

$$\begin{equation}\begin{aligned}\boldsymbol{G}_{n+1}=&\lambda \boldsymbol{G}_{n} + (1 - \lambda)\big(\nabla_{\boldsymbol{\theta}}L(\boldsymbol{\theta}_n)\otimes \nabla_{\boldsymbol{\theta}}L(\boldsymbol{\theta}_n)\big)\\ \boldsymbol{\gamma} =& \frac{\tilde{\gamma}}{\sqrt{\boldsymbol{G}_{n+1} + \epsilon}}\\ \boldsymbol{\theta}_{n+1} =& \boldsymbol{\theta}_n - \boldsymbol{\gamma}\otimes \nabla_{\boldsymbol{\theta}}L(\boldsymbol{\theta}_n) \end{aligned}\end{equation}\label{eq:rmsprop-n}$$

The former accumulates using the overall norm of the gradient, ultimately preserving the scalar nature of the learning rate; the latter accumulates each component of the gradient separately, in which case the adjusted learning rate becomes a vector, equivalent to assigning a different learning rate to each parameter. From the standpoint of strict theoretical analysis, the first approach is actually more rigorous, but experimentally, the second one turns out to be more effective.

What we usually call the RMSprop algorithm refers to the latter, $\eqref{eq:rmsprop-n}$. However, many practitioners who prefer pure SGD-based training tend to criticize this vectorized learning rate approach, arguing that it actually changes the direction of the gradient, leading to inaccurate gradients and ultimately worse results. So readers who dislike vectorized learning rates might want to try the former.

Summary

This article once again analyzed optimization algorithms from the perspective of ODEs, this time providing an understanding of an adaptive learning rate algorithm (RMSprop) through the lens of error control. As for Adam, which we use more commonly, it is a combination of RMSprop and momentum acceleration, which we won't elaborate on here.

Treating the optimization problem as one of solving an ordinary differential equation essentially turns it into a dynamical systems problem, allowing us to understand optimization algorithms from a more physical point of view (even if this understanding is only intuitive rather than rigorous), and even letting us borrow theoretical results from ODE theory. The author will try to give more examples like this in the future.