As is well known, MLM (Masked Language Model) is the pre-training method for BERT and RoBERTa. As the name suggests, it involves masking some tokens from the original sequence and then having the model predict these masked tokens. As research has deepened, it has been discovered that MLM is not only useful as a pre-training method but also has rich application value. For example, I previously found that directly loading BERT’s MLM weights allows it to function as a UniLM for Seq2Seq tasks (refer to here), and the paper "Spelling Error Correction with Soft-Masked BERT" published at ACL 2020 uses the MLM model for text correction.
However, readers who have carefully read the BERT paper or tried it themselves should 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 considered this issue and proposed a new MLM model design that offers higher training efficiency and better results.
The MLM Model
Assume the original sequence is \boldsymbol{x}=[x_1,x_2,\dots,x_T], and \boldsymbol{x}\backslash \{x_i\} represents the sequence after replacing the i-th token with \text{[MASK]}. Then the MLM model models: p\big(x_i, x_j, x_k, \cdots\big|\,\boldsymbol{x}\backslash \{x_i,x_j,x_k,\cdots\}\big) We say its efficiency is low because only a small portion of tokens can be selected for masking each time, such as 15%. This means that only 15% of the tokens in each sample are trained, so the same sample needs to be trained repeatedly multiple times. In BERT, each sample is masked multiple times and stored as tfrecords, which is inefficient and increases disk space usage.
If all tokens of each sample could serve as prediction targets during training, the training efficiency would naturally improve. Unidirectional language models like GPT can achieve this, but MLM is a bidirectional model and cannot do this directly. To achieve this goal, we need to simplify the above formula. Assume we mask only one token at a time, meaning we want to construct the distribution: p\big(x_i\big|\,\boldsymbol{x}\backslash \{x_i\}\big),\,i=1,2,\dots,T We then hope to obtain 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\}) simultaneously through a single model pass. How can this be achieved? This brings us to the results of the paper introduced in this article, which proposes a design called T-TA (Transformer-based Text Autoencoder) that allows us to predict the distribution of all tokens at once.
Introduction to T-TA
First, we know 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 Self-Attention. In MLM, since we want to model p(x_i|\,\boldsymbol{x}\backslash \{x_i\}), the i-th output must not contain information from the i-th token. To this end, the first modification is: remove the token input from \boldsymbol{Q}. That is, the \boldsymbol{Q} of the first layer of Attention cannot contain token information; it can only contain position vectors. This is because we aggregate information from \boldsymbol{K} and \boldsymbol{V} through \boldsymbol{Q}. If \boldsymbol{Q} itself contains token information, it will cause information leakage. Then, we must prevent information leakage from \boldsymbol{K} and \boldsymbol{V}, which requires modifying the Attention Mask to mask out the diagonal part of the Attention (i.e., self-attention), as shown in the figure.
If this is still unclear, we can understand it from the general form of Attention. The general definition of Attention is: 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} It is obvious that Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V})_i must be related to \boldsymbol{q}_i, so \boldsymbol{q}_i absolutely cannot contain information about the i-th token. However, it does not necessarily have to be related to \boldsymbol{k}_i, \boldsymbol{v}_i, because if \text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_i)=0, then \boldsymbol{k}_i, \boldsymbol{v}_i are effectively non-existent. Therefore, the diagonal part of the Attention needs to be masked.
However, this leak-proof Attention Mask can only be maintained for one layer! That is to say, even after doing this, Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V})_j has already incorporated information from the i-th token. Thus, starting from the second layer, if you still use the output of the first layer as \boldsymbol{K} and \boldsymbol{V}, information leakage will occur even with the aforementioned Attention Mask.
The original paper’s solution is blunt, but it seems to be the only way: every layer of Attention shares the original input as \boldsymbol{K} and \boldsymbol{V}! Therefore, let \boldsymbol{E} be the token embedding sequence and \boldsymbol{P} be the corresponding position vectors. The calculation processes of T-TA and BERT can be summarized as:
\begin{array}{c} \boxed{\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 Calculation Schematic} \end{array} \qquad \begin{array}{c} \boxed{\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 Calculation Schematic} \end{array}
Of course, details like residuals and FFN are omitted, leaving only the core operations. During the pre-training phase, T-TA’s Attention uses a diagonal Attention Mask. For downstream task fine-tuning, this mask can be removed.
Experimental Results
Based on the above design, T-TA can predict all tokens at once, resulting in high training efficiency. Furthermore, it does not require additional \text{[MASK]} symbols, thus achieving consistency between pre-training and fine-tuning. However, it is not difficult to understand that T-TA is actually a simplification of the standard Transformer, so theoretically, its fitting ability is weakened. With these trade-offs, is there still an improvement in performance? Naturally, the paper’s experimental results show there is. The original paper conducted multiple experiments, showing that the T-TA design can basically match or even exceed models trained with standard MLM under the same parameter constraints. The authors also generously open-sourced the code for reproducibility (link).
When it comes to modifying Transformer structures, one might imagine a massive array of GPUs and TPUs running in parallel. However, although the authors did not specifically list their experimental equipment, it can be seen from the paper that the setup was not "luxurious." To this end, the authors only trained a 3-layer T-TA and reproduced 3-layer versions of MLM and GPT (unidirectional language model) in the same manner to compare effects. Indeed, all T-TA results in the paper are from 3-layer models, some of which even exceeded the BERT-base version. Thus, the authors vividly taught us a lesson: Without "tycoon" equipment, one can still work on modifying the Transformer, and one can still publish at ACL; the key is having a truly effective idea.
Personal Analysis
Finally, let’s briefly discuss why T-TA is effective. Readers might question how effectiveness can be guaranteed for more layers since the authors only performed 3-layer experiments. Well, let’s look at this model from another perspective.
From a design standpoint, for T-TA, once the input is given, \boldsymbol{K} and \boldsymbol{V} remain constant across all Attention layers; only \boldsymbol{Q} changes. It is not surprising that readers might doubt its effectiveness. However, don’t forget that Google recently proposed Synthesizer (refer to "Google’s New Work Synthesizer: We Don’t Know Self-Attention Well Enough"), which explored several Attention variants. One of them, referred to as "R," essentially fixes \boldsymbol{Q} and \boldsymbol{K} as constants, and it actually works quite well! Note that \boldsymbol{Q} and \boldsymbol{K} in "R" are absolute constants, independent of the input.
Therefore, since having \boldsymbol{Q} and \boldsymbol{K} as constants works reasonably well, why can’t \boldsymbol{K} and \boldsymbol{V} be constants? Moreover, T-TA’s \boldsymbol{K} and \boldsymbol{V} depend dynamically on the input; they are only "constant" once the input is fixed. Thus, theoretically, T-TA’s fitting ability is stronger than the "R" model in Synthesizer. Since "R" can perform well, it is not surprising that T-TA can too.
Of course, I look forward to seeing experimental results with deeper models in the future.
Original article address: https://kexue.fm/archives/7661
For more details on reprinting, please refer to: Scientific Space FAQ