BERT-of-Theseus: A Model Compression Method Based on Module Replacement

BERT-of-Theseus：基于模块替换的模型压缩方法 苏剑林 · 2020-07-17 · Machine Learning model attention model compression

I recently came across a BERT compression method called "BERT-of-Theseus", from the paper BERT-of-Theseus: Compressing BERT by Progressive Module Replacing. It's a compression scheme built around the notion of "replaceability", and compared with the usual pruning or distillation approaches, the whole process feels much more elegant and simple. This post gives a brief introduction to the method, provides an implementation based on bert4keras, and verifies its effectiveness.

BERT-of-Theseus, illustration from the original paper

Model Compression

Let's first briefly go over model compression in general. Since I'm not a specialist in model compression and haven't done a particularly systematic survey of the field, this introduction may be a bit unprofessional — please bear with me.

Basic Concepts

Simply put, model compression means "simplifying a large model to obtain a smaller model with faster inference". Of course, model compression generally comes at some cost — most obviously, the final evaluation metrics will drop somewhat, since a free lunch that's "better and faster" is rare. So the premise of choosing model compression is that you can tolerate some loss of accuracy. Second, the speed-up from model compression usually only shows up at inference time; in other words, it typically requires spending more time on training. So if your bottleneck is training time, model compression isn't for you either.

The reason model compression takes longer is that it requires "training a large model first, then compressing it into a small one". Readers might wonder: why not just train a small model directly? The answer is that many experiments have already shown that training a large model first and then compressing it usually yields higher final accuracy than training a small model directly. That is, for the same inference speed, the compressed model tends to be better. Related discussions can be found in the paper Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers, and there's also a discussion on Zhihu, Why compress a model instead of just training a small CNN directly?.

Common Approaches

Common model compression techniques fall into two broad categories: 1) directly simplifying a large model into a small one; 2) using the large model to help retrain a small model. What these two approaches have in common is that both require first training a reasonably good large model, and only then proceeding with the subsequent steps.

The representative methods of the first category are pruning and quantization. Pruning, as the name suggests, tries to remove some components of the original large model to turn it into a small model, while keeping the model's performance within an acceptable range. Quantization, on the other hand, keeps the original model structure unchanged but switches to a different numerical format, again without severely degrading performance. Typically we build and train models using float32, and switching to float16 can speed things up and save GPU memory; if we can further convert to 8-bit integers or even 2-bit integers (binarization), the speed-up and memory savings become even more pronounced.

The representative method of the second category is distillation. The basic idea of distillation is to use the large model's outputs as labels when training the small model. Take classification as an example: the actual labels are one-hot, whereas the large model's outputs (e.g., logits) carry richer signal, so the small model can learn better features from them. Beyond just learning the large model's outputs, in many cases, to push performance further, the small model is also made to learn the large model's intermediate-layer outputs, attention matrices, correlation matrices, and so on. So a good distillation process usually involves multiple loss terms, and how to design these losses sensibly and balance their weights is one of the research topics in the field of distillation.

Theseus

The compression method this post introduces, called "BERT-of-Theseus", belongs to the second category of compression methods described above — that is, it also uses a large model to help train a small model, except that it is designed based on the replaceability of modules.

The name BERT-of-Theseus comes from the thought experiment "the Ship of Theseus": if the planks of Theseus's ship are gradually replaced until none of the original wood remains, is it still the same ship?

Core Idea

As mentioned above, when using distillation for model compression, we often want not only the small model's output to match the large model's output, but also the intermediate-layer outputs to match. What does "matching" mean here? It means replaceability! So the idea behind BERT-of-Theseus is: why bother laboriously achieving replaceability by adding all sorts of loss terms? Why not just replace modules of the large model with modules of the small model directly, and train the whole thing?

Here's a concrete analogy:

Suppose there are two teams, A and B, each with five players. Team A is a star team with outstanding ability; Team B is a rookie team that needs training. To train Team B, we pick 1 player from Team B and use them to replace 1 player in Team A, then let this "4+1" version of Team A keep practicing and playing matches. After a while, the newly added member's actual ability improves, and this "4+1" team has strength close to the original Team A. Repeat this process until all of Team B's players have been sufficiently trained, and eventually Team B's players can form a team with outstanding strength on their own.
By contrast, if from the very start there had only been Team B, with its members training and playing matches on their own, then even if their ability gradually improves, without the help of the star Team A, their final strength isn't guaranteed to be outstanding.

Process Details

Back to BERT compression: suppose we have a 6-layer BERT that we fine-tune directly on a downstream task to get a reasonably good model, which we call the Predecessor. Our goal is to obtain a 3-layer BERT whose performance on the downstream task is close to that of the Predecessor — or at least better than simply fine-tuning the first 3 layers of BERT directly (otherwise the whole effort would be pointless). We call this small model the Successor. So how does BERT-of-Theseus achieve this? See the figure below (right):

Illustration of the Predecessor and Successor modelsIllustration of the BERT-of-Theseus training process

Throughout the BERT-of-Theseus training process, the Predecessor's weights are kept frozen. The 6-layer Predecessor is divided into 3 modules, each corresponding to one layer of the 3-layer Successor. During training, a Successor layer randomly replaces its corresponding Predecessor module, and then the model is directly fine-tuned using the downstream task's training objective (only training the Successor's layers). After sufficient training, the Successor is then separated out on its own and fine-tuned further on the downstream task until the validation metric stops improving.

The equivalent model for the setup above

In implementation, this is essentially similar to a Dropout-style process: both the Predecessor and Successor models are run simultaneously, and the output of one of the two corresponding modules is zeroed out, then the results are summed and passed to the next layer, i.e.

\begin{equation}\begin{aligned} &\varepsilon^{(l)}\sim U(\{0, 1\})\\ &x^{(l)} = x_p^{(l)} \times \varepsilon^{(l)} + x_s^{(l)} \times \left(1 - \varepsilon^{(l)}\right)\\ &x_p^{(l+1)} = F_p^{(l+1)}\left(x^{(l)}\right)\\ &x_s^{(l+1)} = F_s^{(l+1)}\left(x^{(l)}\right) \end{aligned}\end{equation}

Since $\varepsilon$ is either 0 or 1 (no need to fine-tune this — randomly choosing between the two with equal probability 0.5 already works well), each branch effectively has only one module selected at any time. So the diagram on the right above is equivalent to the model structure shown on the right. Since the zeroing-out is random each time, after enough training steps every layer of the Successor gets adequately trained.

Analysis of the Method

What advantage does BERT-of-Theseus have over distillation? Well, since it got published, its performance is presumably comparable at least, so let's not compare performance and instead compare the methods themselves. It's clear that the main characteristic of BERT-of-Theseus is: simplicity.

As mentioned earlier, distillation often also needs to match intermediate-layer outputs, which brings in many training objectives: downstream-task loss, intermediate-layer output loss, correlation-matrix loss, attention-matrix loss, and so on — just thinking about balancing all these losses is a headache. By contrast, BERT-of-Theseus, through the replacement operation, directly forces the Successor to produce outputs similar to the Predecessor's, and the final training objective is just the downstream-task loss — about as simple as it gets. Moreover, BERT-of-Theseus has another special advantage: many distillation methods need to be applied during both pretraining and fine-tuning to achieve strong results, whereas BERT-of-Theseus can achieve comparable results by acting directly during downstream fine-tuning alone. This advantage isn't apparent from the algorithm itself — it's an empirical finding.

Formally, the random-replacement idea of BERT-of-Theseus is somewhat reminiscent of the image data-augmentation schemes SamplePairing and mixup (see From SamplePairing to mixup: The Magic of the Regularization Term), both of which augment the original model by randomly sampling two objects and taking a weighted sum. It also somewhat resembles the progressive training scheme of PGGAN, where transitioning between two models is achieved by mixing them to some degree. Readers familiar with these methods might raise some extensions or questions about BERT-of-Theseus: does $\varepsilon$ have to be strictly 0 or 1? Would an arbitrary random number in $0\sim 1$ work? Or, instead of randomness, could we just let $\varepsilon$ decrease gradually from 1 to 0? None of these ideas have been thoroughly experimented with; readers who are interested can modify the code below and try them out.

Experimental Results

The original authors have open-sourced their PyTorch implementation, JetRunner/BERT-of-Theseus, and Qiu Zhenyu has also shared his own write-up as well as a TensorFlow implementation based on the original BERT, qiufengyuyi/bert-of-theseus-tf. Naturally, since I decided to write this introduction, there had to be a bert4keras-based Keras implementation as well:

https://github.com/bojone/bert-of-theseus

This is probably the simplest and most readable implementation of BERT-of-Theseus currently available, bar none.

You can look at the original paper yourself for its reported results. I ran experiments on a few text classification tasks, and the results were broadly similar, consistent with Qiu's experimental conclusions as well. On the CLUE iflytek dataset, the results were as follows:

$$\begin{array}{c|c|c} \hline & \text{direct fine-tuning} & \text{BERT-of-Theseus}\\ \hline \begin{array}{c}\text{number of layers} \\ \text{effect}\end{array} & \begin{array}{ccc}\text{full 12 layers} & \text{first 6 layers} & \text{first 3 layers} \\ 60.11\% & 58.99\% & 57.96\%\end{array} & \begin{array}{cc}\text{6 layers} & \text{3 layers} \\ 59.61\% & 59.36\% \end{array}\\ \hline \end{array}$$

As we can see, compared to directly fine-tuning the first few layers, BERT-of-Theseus does bring a certain performance improvement. Regarding the random zeroing scheme, besides choosing between 0/1 with equal probability, the original paper also tried other strategies, which gave a slight improvement but introduced extra hyperparameters, so I didn't experiment with them — interested readers are welcome to modify and try it themselves.

Also, for distillation, if the Successor has the same structure as the Predecessor (same-model distillation), the Successor's final performance is usually even better than the Predecessor's. Does BERT-of-Theseus have this property too? I tested this idea as well, and found that the answer is no — that is, under the same-model setting, the Successor trained via BERT-of-Theseus does not outperform the Predecessor. So it seems that, good as BERT-of-Theseus is, it cannot completely replace distillation.

Summary

This post introduced and experimented with a BERT compression method called "BERT-of-Theseus". Its defining feature is being clean and straightforward: purely through a replacement operation, it lets the small model learn the large model's behavior, achieving state-of-the-art model compression results while relying on just a single loss term.