Trading Time for Performance: A Keras Gradient Accumulation Optimizer

用时间换取效果：Keras梯度累积优化器 苏剑林 · 2019-07-08 · Machine Learning Keras gradient optimizer

Trading Time for Performance: A Gradient-Accumulation Optimizer for Keras

Now in Keras you too can achieve the effect of large batch sizes using small batch sizes — as long as you're willing to spend $n$ times as much time, you can get the effect of $n$ times the batch size, without needing any extra GPU memory. GitHub repo: https://github.com/bojone/accum_optimizer_for_keras

Idle Chatter #

A year or two ago, NLP tasks didn't really require worrying about OOM problems, because compared to models in the CV field, most NLP models were actually quite shallow, and running out of GPU memory was rare. Fortunately, or unfortunately, Bert was born, and then it took off. Bert and its successors (GPT-2, XLNet, etc.) are all built on sufficiently large Transformer models, pretrained on sufficiently large corpora, and then fine-tuned to accomplish specific NLP tasks.

more

Even if you'd really rather not use Bert, the reality nowadays is: the elaborate, complex model you've carefully designed might still perform worse than simply fine-tuning Bert. So one way or another, to keep up with the times, you eventually need to learn how to fine-tune Bert. The trouble is, "you don't know how bad it is until you try it" — as soon as the task gets even slightly more complex, or the sentence length gets even slightly longer, you run out of GPU memory, and the batch size plummets — 32? 16? 8? It's entirely possible to keep dropping and dropping. This isn't hard to understand: Transformers are based on attention, and attention theoretically has both space and time complexity of $\mathcal{O}(n^2)$. Even though, when compute is powerful enough, attention can still perform reasonably fast thanks to its parallelism, the memory footprint can't be avoided — $\mathcal{O}(n^2)$ means that when your sentence length doubles, the memory usage basically needs to become 4 times as much. With that growth rate, OOM is bound to happen sooner or later. And the even more unfortunate news is that, with everyone fine-tuning pretrained Bert, your batch_size=8 might score several parts-per-thousand or even several percentage points lower than someone else's batch_size=80. Obviously, this is quite painful for readers who are trying to climb the leaderboard. Is there really no way out except buying more GPUs?

Getting Down to Business #

Yes there is! By caching and accumulating gradients, we can trade time for space, and the final training effect ends up equivalent to a larger batch size. So, as long as you can run with batch_size=1, and as long as you're willing to spend $n$ times as much time, you can achieve the effect of $n$ times the batch size.

The idea of gradient accumulation was already introduced in an earlier post, "Making Keras Cooler!": Some Lesser-Known Custom Optimizers, where it was called "soft batch." In this post, I'll go with the more mainstream term and call it "gradient accumulation" instead.

The idea behind gradient accumulation is actually very simple. The gradient we use in gradient descent is, in fact, the average of the gradients computed from multiple samples. Take batch_size=128 as an example: you could compute the gradients of all 128 samples at once and average them, or I could compute the average gradient of 16 samples at a time, cache and accumulate it, and after doing this 8 times, divide the total accumulated gradient by 8 before actually performing the parameter update. Of course, you must accumulate for 8 rounds before using the averaged gradient from those 8 rounds to update the parameters — you can't update the parameters every time you compute 16 samples' worth of gradient, otherwise you'd effectively just have batch_size=16.

As I mentioned earlier, the implementation in that previous post was actually flawed, because it used K.switch(cond, K.update(p, new_p), p) to control the update. But in fact this approach cannot control the update, because K.switch only guarantees the selectivity of the result, not the selectivity of the execution — in fact it's equivalent to cond * K.update(p, new_p) + (1 - cond) * p, which means that regardless of cond, both branches actually get executed. In fact, Keras or TensorFlow "almost" never has a conditional construct that executes only one branch (I say "almost" because under some fairly restrictive conditions it can be done), so this path is a dead end.

Since we can't write it that way, the only option is to work on the "update amount" instead: as mentioned before, we compute the gradient for 16 samples each time, and update the parameters every single time, except that in 7 out of every 8 rounds the update amount is zero, and only once is there an actual gradient-descent update applied. Fortunately, this approach can be seamlessly plugged into existing Keras optimizers, meaning we don't need to rewrite the optimizer from scratch! For the detailed implementation, see: https://github.com/bojone/accum_optimizer_for_keras

The specific implementation is really nothing more than a few programming tricks for redirecting things behind the scenes — there isn't much technical depth to it. I won't go into further detail on the implementation itself; feel free to discuss it in the comments if you have questions.

(Note: this optimizer modification, which lets a small batch size achieve the effect of a large batch size, assumes the model does not contain Batch Normalization, because Batch Normalization must use the mean and variance of the entire batch during gradient descent. So if your network uses Batch Normalization and you want to precisely achieve the effect of a larger batch size, the only current solution is more GPU memory / more GPUs.)

Experiments #

The usage is very simple:

`opt = AccumOptimizer(Adam(), 10) # 10 is the number of accumulation steps
model.compile(loss='mse', optimizer=opt)
model.fit(x_train, y_train, epochs=10, batch_size=10)`

This is then equivalent to an Adam optimizer with batch_size=100, at the cost that each epoch runs slower (since the batch size is smaller); the benefit is that you only need the amount of GPU memory required for batch_size=10.

A question readers might want to ask is: how do you prove that your implementation actually works? That is, how do you prove that your result really corresponds to batch_size=100 rather than batch_size=10? To answer this, I ran a fairly extreme experiment; the code is here:

https://github.com/bojone/accum_optimizer_for_keras/blob/master/mnist_mlp_example.py

The code is very simple: it's just an MLP doing MNIST classification with the Adam optimizer, where fit uses batch_size=1. There are two choices for the optimizer: the first is plain Adam(), and the second is AccumOptimizer(Adam(), 100).

With plain Adam(), the loss kept hovering around 0.4 and then got even larger later on (this was true even on the training set), and the validation accuracy never exceeded 97%. With AccumOptimizer(Adam(), 100), the training-set loss kept decreasing, eventually dropping to about 0.02, and the best validation accuracy reached 98%+. Finally, I compared this against plain Adam() but with batch_size=100 directly, and found the results were roughly the same as AccumOptimizer(Adam(), 100) with batch_size=1. This result is sufficient to show that the implementation works as intended and achieves the expected goal.

If that's still not convincing enough, here's one more training result for reference: in a certain Bert fine-tuning experiment, using plain Adam() with batch_size=12, I got 70.33% accuracy; using AccumOptimizer(Adam(), 10) with batch_size=12 (expected to be equivalent to an effective batch size of 120), I got 71% accuracy — a 0.7-point improvement. If you're chasing leaderboard rankings, that 0.7% could well be decisive.

Conclusion #

At last, gradient accumulation (soft batch) has been properly implemented. From now on, when using Bert, you can consider using a large batch_size too~