Six Optimizer Variants: A Brief Introduction and Implementation

6个派生优化器的简单介绍及其实现 苏剑林 · 2019-11-25 · Machine Learning Python Keras optimizer

The optimizer may be one of the most "mysterious" components of deep learning: sometimes just swapping in a different optimizer brings a noticeable improvement; sometimes an optimizer that someone else claims gives huge gains turns out to be utterly useless on your own task. An optimizer with nice theoretical properties doesn't necessarily work well in practice, and an optimizer that seems to have been dreamed up out of thin air isn't necessarily bad either. Regardless, optimizers do give those of us who enjoy "deep-learning alchemy" one more option to play with.

In recent years, work on optimizers seems to have been gradually picking up, with many papers proposing improvements—large and small—to commonly used optimizers (especially Adam). This post gathers together a number of optimizer tricks and techniques, and provides a unified code implementation, for readers to call as needed.

Basic Form

By "derived" here, I mean that the techniques discussed are all built on top of an existing optimizer—any existing optimizer can adopt these techniques and thereby become a new optimizer.

The basic form of an existing optimizer is:

\begin{equation}\begin{aligned}\boldsymbol{g}_t =&\, \nabla_{\boldsymbol{\theta}} L\\ \boldsymbol{h}_t =&\, f(\boldsymbol{g}_{\leq t})\\ \boldsymbol{\theta}_{t+1} =&\, \boldsymbol{\theta}_t - \gamma \boldsymbol{h}_t \end{aligned}\end{equation}

Here $\boldsymbol{g}_t$ is the gradient, and $\boldsymbol{g}_{\leq t}$ refers to all the gradient information up to the current step. After some operation $f$ (such as accumulating momentum, accumulating second-order moments to rescale the learning rate, etc.) is applied to them, we get $\boldsymbol{h}_t$, and then the parameters are updated via $\boldsymbol{h}_t$, where $\gamma$ is the learning rate. more

A Grab-Bag of Variants

Below I introduce six variants of the optimizer, which can also be thought of as tricks for using an optimizer. These tricks are sometimes very effective, and sometimes ineffective or even counterproductive—there's no one-size-fits-all rule, so just think of them as extra options that give you extra possibilities.

Weight Decay

Weight decay refers to directly appending a decay term after each optimizer update step:

\begin{equation}\begin{aligned}\boldsymbol{g}_t =&\, \nabla_{\boldsymbol{\theta}} L\\ \boldsymbol{h}_t =&\, f(\boldsymbol{g}_{\leq t})\\ \boldsymbol{\theta}_{t+1} =&\, \boldsymbol{\theta}_t - \gamma \boldsymbol{h}_t - \gamma \lambda \boldsymbol{\theta}_t \end{aligned}\end{equation}

Here $\lambda$ is called the "decay rate." In SGD, weight decay is equivalent to adding an $l_2$ regularization term $\frac{1}{2}\lambda \Vert \boldsymbol{\theta}\Vert_2^2$ to the loss, but in optimizers with adaptive learning rates such as Adagrad and Adam, $f$ becomes nonlinear, so the two are no longer equivalent. The paper Decoupled Weight Decay Regularization specifically points out that weight decay's ability to prevent overfitting is superior to the corresponding $l_2$ regularization, and recommends using weight decay rather than $l_2$ regularization.

Layer-wise Adaptivity

In an optimizer, the final update is determined by $\boldsymbol{h}_t$ and the learning rate $\gamma$. Sometimes the norm of $\boldsymbol{h}_t$ is larger than the norm of the parameter $\boldsymbol{\theta}_t$, which can cause the update to be unstable. So a straightforward idea is: the update magnitude for each layer's parameters should be regulated by the norm of $\boldsymbol{\theta}_t$. This straightforward idea leads to the following optimizer variant:

\begin{equation}\begin{aligned}\boldsymbol{g}_t =&\, \nabla_{\boldsymbol{\theta}} L\\ \boldsymbol{h}_t =&\, f(\boldsymbol{g}_{\leq t})\\ \boldsymbol{\theta}_{t+1} =&\, \boldsymbol{\theta}_t - \gamma \boldsymbol{h}_t\times \frac{\Vert\boldsymbol{\theta}_t\Vert_2}{\Vert\boldsymbol{h}_t\Vert_2} \end{aligned}\end{equation}

If the base optimizer is Adam, then the above optimizer is LAMB. The paper Large Batch Optimization for Deep Learning: Training BERT in 76 minutes shows that LAMB outperforms Adam when the batch size is large (on the order of thousands).

Piecewise Linear Learning Rate

The learning rate is also something of a "black art" within the optimizer. Generally speaking, carefully tuning the learning rate schedule can bring some improvement, while an inappropriate learning rate can even prevent the model from converging. Common learning rate schedules include warmup, exponential decay, step decay (e.g., dropping to 1/10 of the original value after a certain epoch), and more esoteric ones such as cosine schedules and polynomial schedules.

Given that common functions can all be approximated by piecewise linear functions, I simply introduced a piecewise linear learning rate schedule for anyone to play with freely. It takes the following form:

\begin{equation}\begin{aligned}\boldsymbol{g}_t =&\, \nabla_{\boldsymbol{\theta}} L\\ \boldsymbol{h}_t =&\, f(\boldsymbol{g}_{\leq t})\\ \boldsymbol{\theta}_{t+1} =&\, \boldsymbol{\theta}_t - \gamma \rho_t\boldsymbol{h}_t \end{aligned}\end{equation}

where $\rho_t$ is some piecewise linear function of the step count $t$.

Gradient Accumulation

Gradient accumulation was introduced previously in Trading Time for Effect: A Keras Gradient Accumulation Optimizer. Strictly speaking it isn't a variant of the optimizer itself, but it can be built into the optimizer, achieving the effect of a large batch size using a small batch size—trading time for space. A larger batch size sometimes improves results, especially when the baseline batch size is too small (below 8, say).

To restate the description of gradient descent from Trading Time for Effect: A Keras Gradient Accumulation Optimizer:

The idea behind gradient accumulation is actually very simple. The gradient used in gradient descent is, in fact, the average of the gradients computed from multiple samples. Take batch_size=128 as an example: you can compute the gradients for all 128 samples at once and average them, or I can compute the average gradient for 16 samples at a time, cache and accumulate it, and after doing this 8 times, divide the total gradient by 8 and only then perform the parameter update. Of course, you must accumulate for the full 8 iterations and use the average gradient over those 8 iterations to update the parameters—you can't update after every batch of 16, otherwise it's really just batch_size=16.

Lookahead

The Lookahead optimizer comes from the paper Lookahead Optimizer: k steps forward, 1 step back, and was also introduced previously in Implementing Two Optimizers in Keras: Lookahead and LazyOptimizer. The idea of Lookahead is to use a familiar optimizer to probe ahead a few steps, and then update based on the result of that probing. The procedure is as follows:

\begin{equation}\begin{aligned}&\boldsymbol{g}_t =\, \nabla_{\boldsymbol{\theta}} L\\ &\boldsymbol{h}_t =\, f(\boldsymbol{g}_{\leq t})\\ &\boldsymbol{\theta}_{t+1} =\, \boldsymbol{\theta}_t - \gamma\boldsymbol{h}_t\\ &\text{if}t\,\text{mod}\,k = 0\text{:}\\ &\qquad\boldsymbol{\Theta}_{t+1} = \boldsymbol{\Theta}_t + \alpha (\boldsymbol{\theta}_{t+1}- \boldsymbol{\Theta}_t)\\ &\qquad\boldsymbol{\theta}_{t+1} = \boldsymbol{\Theta}_{t+1} \,(\text{overwrite original}\boldsymbol{\theta}_{t+1}) \end{aligned}\end{equation}

Actually, this optimizer could just as well be called "Lookaback," since every few steps it looks back and interpolates with the weights from a few steps earlier.

Lazy Optimizer

The Lazy optimizer was also introduced in the article just mentioned, Implementing Two Optimizers in Keras: Lookahead and LazyOptimizer. The idea behind it is that updates to an embedding layer should be sparse, which helps prevent overfitting. (See this Zhihu discussion.)

Reference Implementation

The introductions above were kept fairly simple—indeed, these variants themselves aren't hard to understand; the real crux is the code implementation. As you can see from the descriptions above, these six variants are not mutually exclusive, so a good implementation should let us combine one or more of them like building blocks. In addition, keras currently comes in two flavors—plain keras and tf.keras—so a good implementation should be compatible with both (or provide separate implementations for each).

The Ultimate Trick: Grafting

Although some of these variants have been implemented before, here I re-implement them in a new way. This implementation is based on a grafting technique I stumbled upon by accident.

Suppose we have a class like this:

import numpy as np

class A(object):
    def __init__(self):
        self.a = np.ones(1)
        self.b = np.ones(2)
        self.c = np.ones(3)

Now suppose we want to subclass A to get a class B, where B replaces every occurrence of __init__ in method np.ones with np.zeros, leaving everything else unchanged. Since __init__ might be a fairly complex procedure, copying it wholesale and then modifying it would clearly be too redundant.

Is there a way to replace all of it with just a few lines of code? There actually is!

class B(A):
    def __init__(self):
        _ = np.ones
        np.ones = np.zeros
        super(B, self).__init__()
        np.ones = _

With this demo in hand, we can now "hack" existing optimizers. In keras, parameter updates are all implemented via K.update (see keras's optimizers.py); we just need to redefine K.update using the approach above.

What about tf.keras? Unfortunately, this approach doesn't work there, because the iteration logic of the commonly used optimizers in tf.keras has been pushed down into C code (see tf.keras's adam.py)—we can't see the code, so we can't hack it this way. One workaround is to re-implement an optimizer such as Adam ourselves, exposing the iteration logic, so that we can then apply the above grafting approach.

Usage Example

Following the approach above, all six optimizer variants have been implemented in a unified way and placed in my bert4keras project: bert4keras.optimizers.

All the functions correctly detect and import either keras or tf.keras, so that the same code works with both. It comes with a built-in Adam implementation, which is written specifically for tf.keras. For tf.keras, if you want to use the variants described above, you can only use the optimizer bundled with bert4keras (currently only Adam)—you cannot use tf.keras's built-in optimizers.

Reference code:

from bert4keras.optimizers import *

# 变成带权重衰减的Adam
AdamW = extend_with_weight_decay(Adam, 'AdamW')
optimizer = AdamW(learning_rate=0.001, weight_decay_rate=0.01)

# 变成带分段线性学习率的Adam
AdamLR = extend_with_piecewise_linear_lr(Adam, 'AdamLR')
# 实现warmup，前1000步学习率从0增加到0.001
optimizer = AdamLR(learning_rate=0.001, lr_schedule={1000: 1.})

# 变成带梯度累积的Adam
AdamGA = extend_with_gradient_accumulation(Adam, 'AdamGA')
optimizer = AdamGA(learning_rate=0.001, grad_accum_steps=10)

# 组合使用
AdamWLR = extend_with_piecewise_linear_lr(AdamW, 'AdamWLR')
# 带权重衰减和warmup的优化器
optimizer = AdamWLR(learning_rate=0.001,
                    weight_decay_rate=0.01,
                    lr_schedule={1000: 1.})

(Note: implementing this many optimizers all at once, while also trying to maintain compatibility with both keras and tf.keras, inevitably risks some mistakes or omissions. If you spot any, please don't hesitate to point them out.)

Closing Remarks

Alchemy is no easy craft—practice it, and cherish it, while you can.