Trading Time for Performance: Keras Gradient Accumulation Optimizer

Background

A year or two ago, there was little need to worry about OOM (Out of Memory) issues in NLP tasks. Compared to models in the CV field, most NLP models were quite shallow and rarely exhausted video memory (VRAM). Fortunately, or unfortunately, BERT arrived and became a sensation. BERT and its successors (GPT-2, XLNet, etc.) are all based on sufficiently large Transformer models, pre-trained on massive corpora, and then fine-tuned to complete specific NLP tasks.

Even if you are reluctant to use BERT, the current reality is: the complex models you carefully design might not perform as well as simply fine-tuning BERT. Therefore, to keep up with the times, one must learn BERT fine-tuning. The problem is that once you start, you realize that if the task is slightly complex or the sentence length is a bit long, VRAM becomes insufficient, and the batch size drops sharply—32? 16? 8? It is possible for it to keep falling.

This is understandable. Transformers are based on Attention, and Attention theoretically has a space and time complexity of \mathcal{O}(n^2). Although Attention can perform fast enough due to its parallelism when computing power is sufficient, VRAM consumption cannot be avoided. \mathcal{O}(n^2) means that when the sentence length doubles, the VRAM usage essentially quadruples. This growth rate makes it very easy to trigger an OOM error.

Even more unfortunate is that while everyone is fine-tuning pre-trained BERT, your batch_size=8 might be several tenths of a percent or even several percentage points lower than someone else’s batch_size=80. This is clearly frustrating for readers aiming for the top of the leaderboards. Is there no other way besides adding more GPUs?

Implementation

There is! Through gradient caching and accumulation, we can trade time for space. The final training effect is equivalent to using a larger batch size. Therefore, as long as you can run batch_size=1 and are willing to spend n times the time, you can achieve the results of n times the batch size.

The idea of gradient accumulation was introduced in a previous article, "Making Keras Cooler: Niche Custom Optimizers", where it was referred to as a "soft batch." In this article, I will follow the mainstream terminology and call it "accumulate gradients." 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 then average them. Alternatively, I can calculate the average gradient of 16 samples at a time, cache and accumulate them, and after doing this 8 times, divide the total gradient by 8 before performing a parameter update. Of course, the update must only occur after accumulating 8 times using the average gradient of those 8 iterations; you cannot update every 16 samples, or it would simply be batch_size=16.

As mentioned earlier, the implementation in the previous article was incorrect because it used:

K.switch(cond, K.update(p, new_p), p)

to control the update. However, this approach does not actually control the execution of the update because K.switch only guarantees the selectivity of the result, not the selectivity of the execution. In fact, it is equivalent to:

cond * K.update(p, new_p) + (1 - cond) * p

This means that regardless of the value of cond, both branches are executed. In fact, in Keras or TensorFlow, there is "almost" no conditional syntax that executes only one branch (I say "almost" because it can be achieved under very specific, stringent conditions), so this path is a dead end.

Since we cannot write it that way, we must work on the "update amount." As mentioned before, we calculate the gradient for 16 samples and update the parameters every time, but in 7 out of 8 times, the update amount is 0, and only the 8th time is a real gradient descent update. Fortunately, this approach can be seamlessly integrated into existing Keras optimizers, so we don’t need to rewrite the optimizers! For the detailed implementation, please see:

https://github.com/bojone/accum_optimizer_for_keras

The specific implementation involves some programming tricks to "graft" the logic onto the optimizer. There isn’t much high-level technical complexity in the code itself. I won’t go into detail about the implementation here; if you have questions, feel free to discuss them in the comments.

(Note: This optimizer modification allows a small batch size to achieve the effect of a large batch size, provided the model does not contain Batch Normalization. This is because Batch Normalization must use the mean and variance of the entire batch during gradient descent. Therefore, if your network uses Batch Normalization and you want to accurately achieve the effect of a large batch size, the only current method is to add more VRAM or 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)

In this way, it is equivalent to an Adam optimizer with batch_size=100. The cost is that the speed of each epoch will be slower (because the batch size is smaller), but the benefit is that you only need the VRAM required for batch_size=10.

Readers might ask: how do you prove that your implementation works? That is, how do you prove the result is indeed equivalent to batch_size=100 and not batch_size=10? To this end, I conducted an extreme experiment. The code is available here:
https://github.com/bojone/accum_optimizer_for_keras/blob/master/mnist_mlp_example.py

The code is simple: use a multi-layer MLP for MNIST classification with the Adam optimizer and fit with batch_size=1. There are two choices for the optimizer: the first is direct Adam(), and the second is AccumOptimizer(Adam(), 100):

• If using direct Adam(), the loss fluctuates around 0.4 and eventually increases (even on the training set), and the validation accuracy does not exceed 97%.

• If using AccumOptimizer(Adam(), 100), the training loss decreases steadily, eventually reaching around 0.02, and the peak validation accuracy is 98%+.

• Finally, I compared the results of direct Adam() with batch_size=100 and found that the performance was similar to AccumOptimizer(Adam(), 100) with batch_size=1.

These results are sufficient to show that the implementation is effective and achieves the intended purpose. If this is not convincing enough, I will provide another training result as a reference: In a certain BERT fine-tuning experiment, using direct Adam() with batch_size=12, I achieved 70.33% accuracy; using AccumOptimizer(Adam(), 10) with batch_size=12 (expected equivalent batch size of 120), I achieved 71% accuracy, an increase of 0.7%. If you are competing on a leaderboard, this 0.7% could be decisive.

Conclusion

I have finally officially implemented gradient accumulation (soft batch). In the future, when using BERT, we can consider using large batch sizes!

When reposting, please include the original article address: https://kexue.fm/archives/6794

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