Dropping the Constraint, Strengthening the Model: Boosting albert's Performance with One Line of Code

抛开约束，增强模型：一行代码提升albert表现 苏剑林 · 2020-01-29 · Machine Learning language model attention

The title of this post might sound a bit like clickbait, but the change described, when implemented under the bert4keras framework, really is just a one-line change. As for whether it actually brings an improvement, I can't make any promises — but on the handful of fairly representative tasks I tested, it consistently matched or even exceeded the baseline, so the title is basically telling the truth.

So what exactly is the change? It can be summed up in a single sentence:

In downstream tasks, drop albert's weight-sharing constraint — that is, use albert as if it were bert.

For the details of the idea, read on～more

What is albert?

This change was specifically designed for albert, so to understand it we first need to know what albert is. Let me spend a bit of space giving a quick primer on albert. I'll assume the reader already has some familiarity with bert, so the focus here is on comparing the similarities and differences between albert and bert.

Low-rank factorization

First, there's the embedding layer. Take the Chinese bert base as an example: the total number of tokens is roughly 20,000, and the embedding dimension is 768, so the embedding layer alone accounts for about 15 million parameters — roughly 1/6 of the total parameter count. The first thing albert cuts down is exactly this embedding layer: it shrinks the embedding dimension to 128, and then uses a $128\times 768$ transformation matrix to map it back up to 768 dimensions. This way the embedding layer's parameter count drops to only 1/6 of the original — this is the so-called low-rank factorization.

Parameter sharing

Second is the transformer part. In transformer-based architectures like bert, the core building block consists of self-attention, layer normalization, fully-connected layers (i.e., 1D convolutions with kernel size 1), and so on — let's call this a "transformer block." The bert model is simply a stack of many such transformer blocks. The figure on the left below is a minimal schematic of bert base, which stacks 12 transformer blocks:

Minimal schematic of bert baseMinimal schematic of albert base

Notice that in bert's design, the input and output shapes of each transformer block are identical, which means it's perfectly reasonable to feed the output of the current block back in as its own input. This tells us that the same block can, in principle, be reused iteratively, rather than requiring a fresh block at every layer. albert adopts the simplest and most direct scheme: all layers share the very same transformer block (as shown in the right-hand figure above)! As a result, in albert base, the parameter count of the transformer part drops directly to 1/12 of that in bert base.

A brief remark

Besides these two points, another notable difference between albert and bert is that during pretraining, albert replaces the NSP (Next Sentence Prediction) task with the SOP (Sentence-Order Prediction) task. But this is not an architectural matter, so it's not our concern here — readers can look this up on their own.

Overall, albert was designed to reduce the parameter count, in the hope that this reduction would also provide some regularization effect, thereby lowering the risk of overfitting and improving final performance. But did the outcome live up to the author's expectations? Judging by the "track record," when albert first came out, its largest version set a new record on the GLUE leaderboard, so it appears to have met the author's expectations. However, albert isn't always so ideal, and it's not the small model we might imagine it to be.

For a model, what we usually care about are two metrics: speed and performance. From the two figures of bert and albert above, we can already see that at inference time (the forward pass), albert is essentially no different from bert. So under the same specification (say, both base versions), albert and bert have identical inference speed — or, to be even more precise, albert is actually a bit slower, because albert's embedding part has an extra matrix multiplication. In other words, albert brings no improvement in inference speed!

So, at the same specification, which performs better, albert or bert? The albert paper itself already gives the answer: for versions up to large, albert performs worse than bert; only at the xlarge and xxlarge scales does albert start to outperform bert. But roberta-style pretraining compensates for bert's shortcomings, so the only albert version that can genuinely be said to reliably beat bert/roberta is xxlarge. Yet albert xxlarge is such a huge model that it's very hard for most of us to even get it running.

So basically, we can say: (assuming most people can actually afford to run it) at the same inference speed, albert performs worse; at the same performance, albert is slower.

What about training? One point not yet mentioned is that albert's parameter-sharing design has a strong regularization effect, which is why albert drops dropout. Parameter sharing and the removal of dropout do save some GPU memory and speed up training, but in my own evaluations the improvement is only around 10%–20%. In other words, even though albert's parameter count is reduced to 1/10 or even less of bert's, that doesn't mean its GPU memory usage drops to 1/10, nor does it mean training speed increases tenfold — instead, the gains are only modest.

Dropping the sharing constraint

From the discussion above, we can draw a couple of conclusions:

1. In terms of inference alone, albert is basically the same as bert;
2. albert's parameter sharing is, on balance, mostly a negative for performance.

Given this, we might try something new: when fine-tuning on a downstream task, what if we simply drop the parameter-sharing constraint? That is, when fine-tuning, use albert like bert — equivalent to a bert model whose transformer blocks all happen to start from the same initial weights.

Performance evaluation

Without further ado, let's look at the numbers. Here I picked four tasks to test. For reproducibility, each experiment below was run three times, and the table shows the average of the three runs. The "unshared" version refers to the model with parameter sharing removed. The "training speed" column refers to the time taken per epoch, measured on a single TITAN RTX GPU — for relative comparison only.

The experiments were run with bert4keras. For the unshared version, all that's needed is to load the albert weights when build_transformer_model and set model='albert_unshared' — this is the "one line of code" referred to in the title.

First, a relatively simple text sentiment classification task.

$$\begin{array}{c|c|c|c|c} \hline \text{model} & \text{val set (valid)} & \text{training speed} & \text{metric value at end of first epoch} & \text{test set (test)} \\ \hline \text{small_unshared} & 94.66\% & 38s & 90.75\% & 94.35\% \\ \text{small} & 94.57\% & 33s & 91.02\% & 94.52\% \\ \hline \text{tiny_unshared} & 94.02\% & 23s & 88.09\% & 94.13\% \\ \text{tiny} & 94.14\% & 20s & 90.18\% & 93.78\% \\ \hline \end{array}$$

After removing parameter sharing, training time increases slightly, which is expected, while model performance shows mixed results. Since accuracy on this task is already fairly high, it may not reveal much of a gap between the models, so let's move on to a somewhat more complex task.

This time let's try CLUE's IFLYTEK long-text classification task, with results as follows:

$$\begin{array}{c|c|c|c} \hline \text{model} & \text{val set (dev)} & \text{training speed} & \text{metric value at end of first epoch} \\ \hline \text{small_unshared} & 57.73\% & 27s & 49.35\% \\ \text{small} & 57.14\% & 24s & 48.21\% \\ \hline \text{tiny_unshared} & 55.91\% & 16s & 47.89\% \\ \text{tiny} & 56.42\% & 14s & 43.84\% \\ \hline \end{array}$$

Here the advantage of the unshared version starts to show, mainly in overall faster convergence (see the metric values at the first epoch): the small version's best performance is clearly better, while the tiny version's best performance is slightly worse — though after fine-tuning the learning rate more carefully, tiny_unshared's best performance can actually surpass the shared tiny version (but that introduces too many variables, so the table shows results under strict controlled-variable conditions).

Next, let's try a more comprehensive task: information extraction. The results are as follows:

$$\begin{array}{c|c|c|c} \hline \text{model} & \text{val set (dev)} & \text{training speed} & \text{metric value at end of first epoch} \\ \hline \text{small_unshared} & 77.89\% & 375s & 61.11\% \\ \text{small} & 77.69\% & 335s & 46.58\% \\ \hline \text{tiny_unshared} & 76.44\% & 235s & 49.74\% \\ \text{tiny} & 75.94\% & 215s & 31.66\% \\ \hline \end{array}$$

Here we can see that on this more comprehensive, complex task, the unshared version consistently outperforms the original model of the same scale.

Finally, using seq2seq for reading-comprehension-style QA, results are as follows:

$$\begin{array}{c|c|c|c} \hline \text{model} & \text{val set (dev)} & \text{training speed} & \text{metric value at end of first epoch} \\ \hline \text{small_unshared} & 68.80\% & 607s & 57.02\% \\ \text{small} & 66.66\% & 582s & 50.93\% \\ \hline \text{tiny_unshared} & 66.15\% & 455s & 48.64\% \\ \text{tiny} & 63.41\% & 443s & 37.47\% \\ \hline \end{array}$$

This task is mainly meant to test the model's text generation capability. Here we can see the unshared version clearly outperforms the original model, with the unshared tiny version even approaching the original small model's performance.

Analysis and discussion

The experiments above were done on albert tiny/small; I've also run the base version, and the conclusions are basically the same as for tiny and small. However, the base version (and naturally large and xlarge as well) took too long to run, so I didn't complete a full set of experiments for it (nor repeat them three times), and so I haven't included those results here. Overall though, the tiny/small results already seem fairly representative.

The experiments above show that removing parameter sharing from albert makes its downstream-task performance roughly match — or even exceed — the original albert. This suggests that for many NLP tasks, parameter sharing may not actually be a particularly good constraint. Readers might be puzzled: "Why is it that at the xlarge or even xxlarge scale, the parameter-sharing albert starts to outperform the non-shared bert again?" Let me offer a possible explanation.

In theory, bert guards against overfitting via dropout and weight decay. Weight decay is also used in albert, but dropout is absent from albert, so it's worth thinking along the dropout angle. Many experiments have shown that dropout is indeed an effective strategy for reducing overfitting risk, but existing experimental models are generally nowhere near as large as bert xlarge or bert xxlarge, so dropout's effectiveness at such extreme scales remains an open question. In fact, dropout has a training/inference mismatch problem — strictly speaking, the model being trained and the model being used for inference are not exactly the same model. My personal feeling is that as models grow bigger and deeper, this inconsistency gets amplified further, so I don't think dropout is an effective way to prevent overfitting in extremely large models. What about just removing dropout from bert directly? That doesn't work well either, because without dropout there's no longer any mechanism to curb bert's overfitting, and with such a huge parameter count, overfitting would surely become severe. albert, however, removes dropout while introducing implicit regularization through parameter sharing, which lets the model grow bigger and deeper without degrading — and can even improve performance.

Conversely, for albert's parameter sharing to pay off, the model needs to be sufficiently large and deep. So if we're working with the base, small, or even tiny version, we probably shouldn't use parameter sharing at all — for small models, parameter sharing becomes an unnecessary restriction on the model's expressive power, so dropping it tends to yield better performance instead.

Summary

This post experimented with a new idea: dropping albert's parameter sharing during fine-tuning and using albert like bert. Across several tasks, this was found to match or even exceed the original albert's performance. Finally, I offered my own personal understanding of albert and this phenomenon.