Implementing Two Optimizers in Keras: Lookahead and LazyOptimizer

Keras实现两个优化器：Lookahead和LazyOptimizer 苏剑林 · 2019-07-30 · Machine Learning Keras optimizer

I recently implemented two optimizers in Keras, and since there were a few implementation tricks involved, I figured I'd put them together in a post to give a brief introduction (if there were only one, I probably wouldn't bother writing this up). Both optimizers have pretty interesting names — one is "look ahead" (looking forward?), the other is "lazy" (slacking off?). Are these two completely different optimization ideas? Not at all — it's just that their inventors were quite creative with naming. more

Lookahead

First up is the Lookahead optimizer, which comes from the paper Lookahead Optimizer: k steps forward, 1 step back, a fairly recent proposal. Interestingly, both Hinton and Jimmy Ba (one of the authors of Adam) appear on the author list, and with these two heavyweights backing it, the optimizer has drawn quite a bit of attention.

The idea behind Lookahead is very simple — in fact, strictly speaking it isn't really an optimizer itself, but rather a scheme for wrapping an existing optimizer. In short, it's a loop over the following three steps:

1. Back up the model's current weights $\theta$;
2. Starting from $\theta$, use the specified optimizer to take $k$ update steps, obtaining new weights $\tilde{\theta}$;
3. Update the model's weights to $\theta \leftarrow \theta + \alpha\left(\tilde{\theta} - \theta\right)$.

Below is my Keras implementation. The style I used was already mentioned in my earlier post "Making Keras Cooler!": Custom Optimizers for the Few — it's a kind of "invasive" approach:

https://github.com/bojone/keras_lookahead

Usage is very simple:

model.compile(optimizer=Adam(1e-3), loss='mse') # 用你想用的优化器
lookahead = Lookahead(k=5, alpha=0.5) # 初始化Lookahead
lookahead.inject(model) # 插入到模型中

As for its effectiveness, the original paper ran quite a few experiments — some showed slight improvements (the cifar10 and cifar100 ones), while others showed a more noticeable boost (the LSTM language model one). I did a quick test of my own and found essentially no change. I've always felt that optimizers are somewhat mysterious creatures — sometimes only SGD reaches the best results, sometimes only Adam manages to converge at all. In any case, you can't expect simply swapping in a different optimizer to dramatically improve your model. Lookahead's arrival just gives us one more option to try; readers with plenty of training time to spare are welcome to experiment with it.

Also see: Synced's introduction to Lookahead

LazyOptimizer

The LazyOptimizer is basically designed for NLP — or more precisely, for the Embedding layer.

LazyOptimizer points out that all momentum-based optimizers (which of course includes Adam as well as SGD with momentum) share a common problem: words that aren't sampled in the current batch still get updated using historical momentum, which can cause the Embedding layer to overfit (see this Zhihu discussion). Specifically, once a word has been sampled, the gradient of its embedding is nonzero, and this gradient gets recorded in the momentum — the actual update is driven by the momentum. In later batches, even if that word isn't sampled again, its gradient will be zero, but its momentum won't be, so the word still gets updated. As a result, even words that aren't sampled repeatedly end up having their embeddings updated repeatedly, leading to overfitting.

So, one improved approach is to only update a word's embedding when it has actually been sampled — this is the basic idea behind LazyOptimizer.

In terms of implementation, how do we determine whether a word has been sampled? The ultimate approach, of course, would be to pass in the indices of the sampled words directly, but that's not very user-friendly. Here I used an approximate method instead: check whether the gradient corresponding to that word's embedding is zero. If it is zero, that "most likely" means the word wasn't sampled in the current batch. The reasoning behind this: if the word wasn't sampled, its gradient is guaranteed to be zero; if it was sampled, the probability of the gradient being exactly zero is extremely small — after all, with so many components, the chance of all of them being zero simultaneously is negligible. So this approximation is good enough in practice.

My Keras implementation can be found at:

https://github.com/bojone/keras_lazyoptimizer

Usage here is also simple — you just wrap a momentum-based optimizer, pass in all the Embedding layers, and it becomes a new "lazy" version of that optimizer:

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

The GitHub repo also includes an IMDB example. In this example, if you directly use Adam(1e-3) as the optimizer, the best validation accuracy tops out at around 83.7%, whereas if you use LazyOptimizer(Adam(1e-3), embedding_layers), the best validation accuracy generally reaches above 84.9% — the improvement is quite noticeable. Overall, I think this is worth trying for any model with a large Embedding layer (especially word-level models). Basically, since the Embedding layer has so many parameters, reducing its update frequency lets the model focus its optimization effort on the rest of the network.

Note: this LazyOptimizer is a bit different from the standard LazyOptimizer. In the standard version, for words that haven't been sampled, all related cached quantities (such as momentum) are also left unchanged. In my implementation, however, even if a word hasn't been sampled, all of its corresponding cached quantities still get updated. Some evaluations suggest that this actually works a bit better.

Summary

Nothing too fancy here — just implementations of two optimizers in Keras, so that fellow Keras users can give them a try right away, or just have a bit more fun using Keras.