Keras Implementation of Two Optimizers: Lookahead and LazyOptimizer

Lookahead

First up is the Lookahead optimizer, which originates from the paper "Lookahead Optimizer: k steps forward, 1 step back". It is a recently proposed optimizer, and interestingly, the legendary Geoffrey Hinton and one of the authors of Adam, Jimmy Ba, are among the authors. With the backing of these two giants, this optimizer has attracted significant attention.

The idea behind Lookahead is very simple. To be precise, it is not an optimizer itself, but rather a scheme for using existing optimizers. Simply put, it iteratively executes the following three steps:

1. Back up the model’s existing weights \theta;

2. Starting from \theta, update the weights for k steps using a specified optimizer to obtain new weights \tilde{\theta};

3. Update the model weights as \theta \leftarrow \theta + \alpha(\tilde{\theta} - \theta).

Below is my Keras implementation. The coding style was mentioned in a previous article, "Making Keras a Bit Cooler: Niche Custom Optimizers", which follows an "invasive" implementation approach:

https://github.com/bojone/keras_lookahead

The usage is very straightforward:

model.compile(optimizer=Adam(1e-3), loss='mse') # Use any optimizer you want
lookahead = Lookahead(k=5, alpha=0.5) # Initialize Lookahead
lookahead.inject(model) # Inject it into the model

Regarding its effectiveness, the original paper conducted several experiments. Some showed slight improvements (like the CIFAR-10 and CIFAR-100 experiments), while others showed more significant gains (like the LSTM language model). I did some simple experiments myself, and the results showed no significant change. I have always felt that optimizers are somewhat mysterious; sometimes you must use SGD to achieve the best results, and other times you must use Adam to get the model to converge. In short, one cannot expect to significantly improve model performance simply by switching an optimizer. The emergence of Lookahead just gives us one more choice. Readers with sufficient training time should feel free to experiment with it.

Note: For a more detailed introduction to Lookahead, see the article by Jiqizhixin.

LazyOptimizer

The LazyOptimizer is essentially designed for NLP, or more accurately, for Embedding layers.

LazyOptimizer points out a problem that exists in all optimizers with momentum (which naturally includes Adam and SGD with momentum): words that are not sampled in the current batch are still updated using historical momentum, which may lead to overfitting of the Embedding layer (refer to the Zhihu discussion). Specifically, after a word is sampled, the gradient of its Embedding is non-zero, and this gradient is recorded in the momentum. The actual update uses this momentum. In subsequent batches, if the word is not sampled, its Embedding gradient is zero, but its momentum is not zero, so the word is still updated. Consequently, even words that are not repeatedly sampled have their corresponding Embeddings updated repeatedly, leading to overfitting.

Therefore, an improved scheme is to update only when the word has been sampled. This is the basic principle of LazyOptimizer.

In terms of implementation, how do we determine if a word has been sampled? The ultimate method would be to pass the indices of the sampled words, but that is not user-friendly. I used an approximate method here: check if the gradient corresponding to the word’s Embedding is zero. If it is zero, it means it is "very likely" that it was not sampled in the current batch. The logic is that if it wasn’t sampled, the gradient must be zero; if it was sampled, the probability of the gradient being exactly zero is extremely small, given the many components involved. Thus, this implementation is sufficient.

My Keras implementation is located at:

https://github.com/bojone/keras_lazyoptimizer

The usage is also very simple: wrap an optimizer that has momentum and pass in all the Embedding layers to create a new "Lazy" version of the optimizer:

model.compile(
    loss='mse',
    optimizer=LazyOptimizer(Adam(1e-3), embedding_layers)
)

The GitHub repository also includes an IMDB example. In this example, if you use Adam(1e-3) directly as the optimizer, the highest validation accuracy is around 83.7%. However, if you use LazyOptimizer(Adam(1e-3), embedding_layers), the optimal validation accuracy can basically reach over 84.9%. The effect is quite evident. Overall, I think models with very large Embedding layers (especially word-based models) are worth a try. The general idea is that because the Embedding layer has too many parameters, reducing the update frequency allows the model to focus on optimizing the remaining parts.

Note: This LazyOptimizer is slightly different from the standard one. In the standard LazyOptimizer, for words that are not sampled, all related cached values (such as momentum) are also not updated. In my implementation, even if the word is not sampled, all cached values for that word are still updated. Some evaluations suggest that this approach actually yields better results.

Summary

There isn’t much else to say—I just implemented two optimizers in Keras so that friends using Keras can try them out or make their Keras experience a bit more interesting.

Original Address: https://kexue.fm/archives/6869

For more details on reprinting, please refer to: Scientific Space FAQ