A Brief Introduction and Implementation of 6 Derived Optimizers

Weight Decay

Weight decay refers to adding a decay term directly after each update step of the optimizer: \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} where \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. However, in optimizers with adaptive learning rates like Adagrad and Adam, f becomes non-linear, so the two are no longer equivalent. The paper "Decoupled Weight Decay Regularization" specifically points out that the anti-overfitting ability of weight decay is superior to the corresponding l_2 regularization, and recommends using weight decay instead of l_2 regularization.

Layer Adaptation

In an optimizer, the final update amount is determined by \boldsymbol{h}_t and the learning rate \gamma. Sometimes the magnitude of \boldsymbol{h}_t can be much larger than the magnitude of the parameters \boldsymbol{\theta}_t, which may lead to unstable updates. Therefore, a direct idea is that the update magnitude of each layer’s parameters should be regulated by the norm of \boldsymbol{\theta}_t. This idea leads to the following optimizer variant: \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} If the base optimizer is Adam, the above optimizer is LAMB. The paper "Large Batch Optimization for Deep Learning: Training BERT in 76 minutes" points out that LAMB performs better than Adam when the batch size is large (thousands or tens of thousands).

Piecewise Linear Learning Rate

The learning rate is also a mysterious entity in optimizers. Generally, fine-tuning the learning rate strategy can yield certain improvements, while an inappropriate learning rate may even prevent the model from converging. Common learning rate strategies include warmup, exponential decay, step decay (e.g., dropping to 1/10 after a certain epoch), etc. More exotic ones include cosine learning rate strategies and polynomial learning rate strategies.

Considering that common functions can be approximated by piecewise linear functions, I have introduced a piecewise linear learning rate strategy for everyone to experiment with. The form is as follows: \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} where \rho_t is a piecewise linear function with the step count t as the independent variable.

Gradient Accumulation

Gradient accumulation was previously introduced in the article "Trading Time for Results: Keras Gradient Accumulation Optimizer". It is not strictly a variant of an optimizer, but it can be implemented within one. It achieves the effect of a large batch size using a small batch size, trading time for memory space. A larger batch size can sometimes improve results, especially when the baseline batch size is too small (below 8?).

Restating the description of gradient descent from the article "Trading Time for Results: Keras Gradient Accumulation Optimizer":

The concept of gradient accumulation is quite simple. The gradient we use for gradient descent is actually the average of gradients calculated from multiple samples. Taking batch_size=128 as an example, you can calculate the gradients of 128 samples at once and average them, or I can calculate the average gradient of 16 samples each time, cache and accumulate them, and after doing this 8 times, divide the total gradient by 8 before executing the parameter update. Of course, the update must only occur after accumulating 8 times using the 8-step average gradient; you cannot update every 16 samples, otherwise, it would just be batch_size=16.

Lookahead

The Lookahead optimizer comes from the paper "Lookahead Optimizer: k steps forward, 1 step back" and was also introduced in the previous article "Implementing Two Optimizers in Keras: Lookahead and LazyOptimizer". The idea of Lookahead is to use a standard optimizer to explore a few steps forward and then update based on the exploration results. The process is as follows: \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{i.e., overwrite the original }\boldsymbol{\theta}_{t+1}) \end{aligned}

In fact, this optimizer could just as well be called "Lookback," as it looks back every few steps and performs an interpolation with the weights from several steps ago.

Lazy Optimizer

The Lazy optimizer was also introduced in the article "Implementing Two Optimizers in Keras: Lookahead and LazyOptimizer". Its original intention is that updates to the Embedding layer should be sparsified, which helps prevent overfitting (refer to the Zhihu discussion).

Ultimate Grafting

Although some variants were implemented before, they are re-implemented here using a new method. This implementation method stems from an "ultimate grafting" (monkey patching) trick I accidentally discovered.

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 inherit from class A to create class B, where class B replaces all np.ones in the __init__ method with np.zeros, while keeping everything else the same. Since __init__ might be a very complex process, copying it entirely and modifying it would be too redundant.

Is there a way to replace everything with just a few lines of code? Yes, there is!

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

With this demo, we can "mod" existing optimizers. In Keras, parameter updates are all implemented through K.update (refer to Keras’s optimizers.py). We only need to redefine K.update using the method above.

What about tf.keras? Unfortunately, this method does not work because the iteration processes of common optimizers in tf.keras are written in C++ (refer to tf.keras’s adam.py). We cannot see the code, so we cannot modify it this way. One solution is to re-implement an optimizer like Adam ourselves, exposing the iteration process so that we can use the grafting method to modify it.

Usage Example

The six optimizer variants implemented uniformly according to the above logic are placed in my bert4keras project: bert4keras.optimizers.

All functions will perform the correct imports based on whether Keras or tf.keras is used, allowing both to be used in the same way. It includes a built-in Adam implementation specifically written for tf.keras. For tf.keras, if you want to implement the above variants, you must use the optimizers provided by bert4keras (currently only Adam); you cannot use the built-in optimizers of tf.keras.

Reference code:

from bert4keras.optimizers import *

# Turn into Adam with Weight Decay
AdamW = extend_with_weight_decay(Adam, 'AdamW')
optimizer = AdamW(learning_rate=0.001, weight_decay_rate=0.01)

# Turn into Adam with Piecewise Linear Learning Rate
AdamLR = extend_with_piecewise_linear_lr(Adam, 'AdamLR')
# Implement warmup: learning rate increases from 0 to 0.001 over the first 1000 steps
optimizer = AdamLR(learning_rate=0.001, lr_schedule={1000: 1.})

# Turn into Adam with Gradient Accumulation
AdamGA = extend_with_gradient_accumulation(Adam, 'AdamGA')
optimizer = AdamGA(learning_rate=0.001, grad_accum_steps=10)

# Combined usage
AdamWLR = extend_with_piecewise_linear_lr(AdamW, 'AdamWLR')
# Optimizer with both Weight Decay and Warmup
optimizer = AdamWLR(learning_rate=0.001,
                    weight_decay_rate=0.01,
                    lr_schedule={1000: 1.})

(Note: Implementing so many optimizers while maintaining compatibility with both Keras and tf.keras may lead to errors or omissions. If any are found, please do not hesitate to point them out.)