Modifying the Transformer Architecture to Design a Faster, Better MLM Model

修改Transformer结构，设计一个更快更好的MLM模型 苏剑林 · 2020-08-07 · Machine Learning language model attention

As everyone knows, MLM (Masked Language Model) is the pretraining scheme used by BERT and RoBERTa: as the name suggests, it masks out some tokens from the original sequence and then has the model predict those masked tokens. As research has progressed, people have found that MLM isn't just useful as a pretraining objective — it also has rich application value. For example, I previously found that simply loading BERT's MLM weights lets it be used as UniLM for Seq2Seq tasks (see here), and similarly, the ACL 2020 paper Spelling Error Correction with Soft-Masked BERT applies an MLM model to text error correction.

However, anyone who has read the BERT paper carefully, or tried it hands-on, will know that the training efficiency of the original MLM is relatively low, because only a small portion of tokens can be masked for training at a time. The ACL 2020 paper Fast and Accurate Deep Bidirectional Language Representations for Unsupervised Learning also considers this problem and proposes a new MLM model design that achieves higher training efficiency and better performance. more

The MLM Model

Suppose the original sequence is $\boldsymbol{x}=[x_1,x_2,\dots,x_T]$, and $\boldsymbol{x}\backslash \{x_i\}$ denotes the sequence obtained by replacing the i-th token with $\text{[MASK]}$. Then the MLM model amounts to modeling

\begin{equation}p\big(x_i, x_j, x_k, \cdots\big|\,\boldsymbol{x}\backslash \{x_i,x_j,x_k,\cdots\}\big)\end{equation}

We say this is inefficient because only a small fraction of tokens can be selected for masking each time — say 15% — meaning that only 15% of the tokens in each sample get trained on in a given pass, so the same sample has to be trained on repeatedly many times. In BERT, each sample is masked multiple times and the results are stored as tfrecords, which not only lowers training efficiency but also increases disk space usage.

Illustration of the MLM task

If, during training, every token in a sample could serve as a prediction target, training efficiency would naturally improve. Unidirectional language models like GPT can do this, but MLM is a bidirectional model and cannot directly achieve this. To reach this goal, we need to simplify the formula above by assuming that only one token is masked out at a time, i.e., the distribution to be constructed is

\begin{equation}p\big(x_i\big|\,\boldsymbol{x}\backslash \{x_i\}\big),\,i=1,2,\dots,T\end{equation}

We then want a single model, in a single forward pass, to simultaneously predict $p(x_1|\,\boldsymbol{x}\backslash \{x_1\}),p(x_2|\,\boldsymbol{x}\backslash \{x_2\}),\dots,p(x_T|\,\boldsymbol{x}\backslash \{x_T\})$. How can this be achieved? This brings us to the paper this post is about, which proposes a design called T-TA (Transformer-based Text Autoencoder) that lets us predict the distributions of all tokens at once.

Introducing T-TA

T-TA's Attention Mask pattern

First, recall that the core operation of the Transformer is $Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V})$; in BERT, $\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}$ are all the same, which is exactly Self-Attention. In MLM, since we want to model $p(x_i|\,\boldsymbol{x}\backslash \{x_i\})$, the $i$-th output must clearly not contain any information about the $i$-th token. To achieve this, the first modification we make is: remove the token input from $\boldsymbol{Q}$, meaning that the $\boldsymbol{Q}$ of the first Attention layer must not contain token information, only positional vectors. This is because it is through $\boldsymbol{Q}$ that we aggregate the information of $\boldsymbol{K},\boldsymbol{V}$; if $\boldsymbol{Q}$ itself already carried token information, that would leak information. Next, we need to prevent the leakage of $\boldsymbol{K},\boldsymbol{V}$'s own information, which requires modifying the Attention Mask to mask out the attention along the diagonal (i.e., attention to oneself), as shown in the figure.

If this still isn't clear, we can understand it from the general form of Attention. Attention is generally defined as

\begin{equation}Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V})_i = \frac{\sum\limits_{j=1}^n \text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)\boldsymbol{v}_j}{\sum\limits_{j=1}^n \text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)}\label{eq:gen-att}\end{equation}

So it's clear that $Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V})_i$ is necessarily related to $\boldsymbol{q}_i$, which means $\boldsymbol{q}_i$ absolutely must not contain information about the $i$-th token; but it need not be related to $\boldsymbol{k}_i,\boldsymbol{v}_i$, since as long as $\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_i)=0$, $\boldsymbol{k}_i,\boldsymbol{v}_i$ is effectively nonexistent — hence we need to mask out the diagonal attention entries.

However, this leak-preventing Attention Mask can only be maintained for one layer! That is, even after doing this, $Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V})_j$ has already absorbed information about the $i$-th token, so starting from the second layer, if you still take the output of the first layer as $\boldsymbol{K},\boldsymbol{V}$, information leakage will occur even with the above Attention Mask in place.

The original paper's solution is somewhat brute-force, but it seems to be the only way to solve it: every Attention layer shares the original input as $\boldsymbol{K},\boldsymbol{V}$! So, letting $\boldsymbol{E}$ be the sequence of token embeddings and $\boldsymbol{P}$ the corresponding positional vectors, the computation process of T-TA versus BERT can be written concisely as:

\begin{equation} \begin{array}{c}\bbox[border: 1px dashed red; padding: 5px]{\begin{aligned}&\boldsymbol{Q}_0 = \boldsymbol{E}+\boldsymbol{P}\\ &\boldsymbol{Q}_1 = Attention(\boldsymbol{Q}_0,\boldsymbol{Q}_0,\boldsymbol{Q}_0) \\ &\boldsymbol{Q}_2 = Attention(\boldsymbol{Q}_1,\boldsymbol{Q}_1,\boldsymbol{Q}_1) \\ &\qquad\vdots\\ &\boldsymbol{Q}_n = Attention(\boldsymbol{Q}_{n-1},\boldsymbol{Q}_{n-1},\boldsymbol{Q}_{n-1}) \end{aligned}} \\ \text{BERT computation diagram}\quad\end{array}\qquad \begin{array}{c}\bbox[border: 1px dashed red; padding: 5px]{\begin{aligned}&\boldsymbol{Q}_0 = \boldsymbol{P}\\ &\boldsymbol{Q}_1 = Attention(\boldsymbol{Q}_0,\boldsymbol{E}+\boldsymbol{P},\boldsymbol{E}+\boldsymbol{P}) \\ &\boldsymbol{Q}_2 = Attention(\boldsymbol{Q}_1,\boldsymbol{E}+\boldsymbol{P},\boldsymbol{E}+\boldsymbol{P}) \\ &\qquad\vdots\\ &\boldsymbol{Q}_n = Attention(\boldsymbol{Q}_{n-1},\boldsymbol{E}+\boldsymbol{P},\boldsymbol{E}+\boldsymbol{P}) \end{aligned}} \\ \text{T-TA computation diagram}\quad\end{array}\end{equation}

Of course details like the residual connections and FFN are omitted here, keeping only the core operation. During pretraining, T-TA's Attention uses the diagonal-masked form described above; for downstream fine-tuning, this mask can be removed.

Experimental Results

One of the experimental tables from the original paper. T-TA can be seen to have its own distinctive advantages in semantic representation.

With this design, T-TA can predict all tokens in a single pass, giving it high training efficiency, and it doesn't need an extra $\text{[MASK]}$ symbol, thus achieving consistency between pretraining and fine-tuning. But it's not hard to see that T-TA is essentially a simplification of the standard Transformer, so in theory its fitting capacity should be weaker. Given this trade-off between gains and losses, does it actually perform better in practice? Indeed, the paper's experimental results say yes. The original paper conducted multiple experiments showing that, under the same parameter budget, this T-TA design can basically match or even surpass models trained with standard MLM. The authors also generously open-sourced their code so that others can reproduce the results (link).

Speaking of modifying the Transformer architecture, one might imagine this requires a large fleet of GPUs or TPUs running in parallel. But in fact, although the authors didn't explicitly list their experimental hardware, judging from the paper their setup doesn't look particularly "lavish." Accordingly, the authors only trained a 3-layer T-TA, and, following the same pattern, reproduced a 3-layer MLM and a 3-layer GPT (i.e., a unidirectional language model) for comparison. Indeed, all the T-TA results in the paper are for 3-layer models, and some of them even surpass the Base version of BERT. So the authors give us a vivid lesson: you don't need lavish hardware to do work on modifying the Transformer, and you can still publish at ACL — the key is having a genuinely effective idea.

Personal Analysis

Finally, let's briefly discuss why T-TA works. Readers might question: since the authors only ran experiments with 3 layers, how can we be sure it still works with more layers? Well, let's look at this model from another angle.

By design, for T-TA, once the input is given, $\boldsymbol{K},\boldsymbol{V}$ stays the same across all Attention layers, and only $\boldsymbol{Q}$ changes — so it's not surprising that readers might doubt its effectiveness. But don't forget, a while back Google proposed Synthesizer (see Google's New Work Synthesizer: We Still Don't Fully Understand Self-Attention), which explored several Attention variants, one of which, abbreviated "R," is equivalent to fixing $\boldsymbol{Q},\boldsymbol{K}$ as a constant — and it turned out to work rather well! Note that the $\boldsymbol{Q},\boldsymbol{K}$ in "R" is a genuine constant, entirely unrelated to the input.

So, given that fixing $\boldsymbol{Q},\boldsymbol{K}$ as a constant still works reasonably well, why couldn't $\boldsymbol{K},\boldsymbol{V}$ also be a constant? Moreover, T-TA's $\boldsymbol{K},\boldsymbol{V}$ dynamically depends on the input — it's only constant once the input has been fixed — so in theory T-TA's fitting capacity should be stronger than Synthesizer's "R" model. Given that "R" already works well, it shouldn't be surprising that T-TA works well too.

That said, I'm still looking forward to seeing deeper experimental results in follow-up work.