I believe that over the past year (especially the last six months), everyone has frequently seen reports on various Transformer-related works (such as BERT, GPT, XLNet, etc.), along with the continuous breaking of benchmarks for various basic evaluation tasks. At the same time, many blogs and columns have provided popular science and interpretations of these models.
As the saying goes, "the layman looks at the fun, while the professional looks at the trick." We not only need to understand these works at the level of "what they are," but we also need to think about "why." This "why" includes not only "why do it this way" but also "why is it possible to do it this way." For example, when discussing XLNet’s permutation language model, we might have understood its benefits from various introductions, but we should think further:
Why can Transformer implement a permutation language model? How is it implemented? Can an RNN implement it?
This article analyzes the fundamental reasons why many Transformer models can perform so "brilliantly" from the perspective of masking the Attention matrix. As the title suggests, "Transformer is like a play; it all depends on the Mask." This is one of the important "tricks" of various fancy Transformer models.
By reading this article, you may learn about:
The relationship between the masking methods of the Attention matrix and various pre-training schemes;
Directly using a pre-trained BERT model for Seq2Seq tasks.
Background
Since "Attention is All You Need", Transformer-like models based purely on Attention have gradually become popular, and the emergence of BERT pushed this trend to a new height. Subsequently, various works based on large-scale pre-trained Transformer models have continued to appear. Some apply existing models, some attempt to better explain and visualize these models, and others improve architectures or pre-training methods to achieve better results. Overall, these pre-training-based works are emerging one after another, creating a dazzling array of options. To some extent, if you haven’t fine-tuned BERT yet, you are already lagging behind mainstream NLP technology.
Fancy Pre-training
As is well known, the traditional means of model pre-training is the language model. For example, the ELMo model is based on the BiLSTM architecture and uses language models in two directions to pre-train two LMs respectively. Later, OpenAI’s GPT and GPT-2 also unswervingly adhered to using the traditional (standard, unidirectional) language model for pre-training.
However, there are even more varieties of pre-training gameplay. For example, BERT used a method called "Masked Language Model (MLM)" for pre-training, though this is just a variant of the ordinary language model. XLNet proposed a more thorough "Permutation Language Modeling," which we can call a "shuffled language model." There is also the UNILM model, which directly uses a single BERT architecture for Seq2Seq. You can use it as a pre-training method or simply use it for Seq2Seq tasks...
With such a variety of methods, we cannot help but wonder: Why did this phenomenon of "a hundred flowers blooming" in large-scale pre-trained models only appear in the era when Transformer became popular?
Transformer Exclusive
In fact, besides the unidirectional language model and its simple variant, the masked language model, UNILM’s Seq2Seq pre-training and XLNet’s permutation language model pre-training are essentially customized for the Transformer architecture. To put it bluntly, if it were an RNN architecture, it would be impossible to pre-train using the permutation language model method. As for the Seq2Seq pre-training method, two models (encoder and decoder) must be introduced simultaneously, whereas the Transformer architecture can handle it with a single model.
The secret lies within the Attention matrix. Attention is essentially calculating the similarity between inputs in pairs, which constitutes an n^2 similarity matrix (i.e., the Attention matrix, where n is the sentence length; in this article, Attention refers to Self-Attention). This means its space complexity is of the order \mathcal{O}(n^2). In contrast, RNN and CNN models are only \mathcal{O}(n). Therefore, Attention is usually more memory-intensive. However, there are advantages to this: larger space occupancy also means more possibilities. We can add various prior constraints to this \mathcal{O}(n^2) Attention matrix, making it capable of more flexible tasks. Simply put, only pure Attention models have enough "capacity" to carry so many "variations."
The way to add prior constraints is to perform different forms of Masking on the Attention matrix, which is the focus of this article.
Analysis
In the article "A Brief Reading of ’Attention is All You Need’", I provided a basic introduction to Attention. Here is a brief review. The mathematical form of Attention is: Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}) = \text{softmax}\left(\frac{\boldsymbol{Q}\boldsymbol{K}^{\top}}{\sqrt{d_k}}\right)\boldsymbol{V}
Here \boldsymbol{Q}\in \mathbb{R}^{l_q\times d_q}, \boldsymbol{K}\in\mathbb{R}^{l_k\times d_q}, \boldsymbol{V}\in\mathbb{R}^{l_k\times d_v} represent the query, key, and value vector sequences, respectively. We can consider key and value to be in a one-to-one correspondence. \boldsymbol{Q}\boldsymbol{K}^{\top} performs the dot product between every query and key vector, which is then normalized by softmax to obtain an l_q\times l_k Attention matrix. it describes the correlation strength between any two elements of query and key. The stories we will tell later are all based on this Attention matrix. Finally, multiplying by \boldsymbol{V} is equivalent to weighted summation of the vectors in \boldsymbol{V} according to this correlation strength, ultimately outputting an l_q\times d_v vector sequence.
The most commonly used Attention method is Self-Attention, where \boldsymbol{Q}, \boldsymbol{K}, \boldsymbol{V} are all derived from the same vector sequence through linear transformations. Transformer is a combination of Self-Attention and Position-Wise Feed-Forward layers (equivalent to a 1D convolution with a kernel size of 1). Thus, Transformer is an Attention-based transformation from vector sequence to vector sequence.
In this section, we will analyze the Masking methods of the Attention matrix in detail, which correspond to the implementation principles of unidirectional language models, permutation language models, and Seq2Seq.
Unidirectional Language Model
A language model can be seen as an unconditional text generation model. If readers are not familiar with text generation models, they can refer to relevant materials and the article "Playing with Keras: Seq2Seq Automatic Title Generation". A unidirectional language model essentially "remembers" the training corpus through the following conditional probability distribution: p(x_1,x_2,x_3,\dots,x_n)=p(x_1) p(x_2|x_1) p(x_3|x_1,x_2) \dots p(x_n|x_1,\dots,x_{n-1})
When we usually say "language model," we refer to a unidirectional (more narrowly, just forward) language model. The key point of a language model is to prevent seeing "future information." As shown in the formula above, when predicting x_1, there is no external input; when predicting x_2, only x_1 can be input; when predicting x_3, only x_1, x_2 can be input; and so on.
RNN models are naturally suitable for language models because they are inherently recursive. If using CNNs, one needs to Mask the convolution kernel, i.e., zero out the part corresponding to the right side of the kernel. What about Transformer? It requires an Attention matrix in the form of a lower triangular matrix:
As shown in the figure, each row of the Attention matrix represents the output, and each column represents the input. The Attention matrix represents the correlation between output and input. Assuming white squares represent 0, the 1st row indicates that "Bei" (North) can only be related to the start token <s>. The 2nd row indicates that "Jing" (Capital) can only be related to the start token <s> and "Bei". And so on. Therefore, one only needs to introduce a lower triangular Mask in the Transformer’s Attention matrix and shift the input and output by one position during training to implement a unidirectional language model. (For the implementation of Masking, refer to the Mask section of "Making Keras Cooler: Layers within Layers and Masking".)
Permutation Language Model
The permutation language model is a concept proposed by XLNet, primarily used for XLNet’s pre-training. Speaking of XLNet, I find its permutation language model pre-training very interesting, but I don’t like that it changed the basic architecture to Transformer-XL. I think anyone with resources could try the combination of "BERT + Permutation Language Model pre-training"; there might be unexpected discoveries.
Like the standard language model, the permutation language model performs conditional probability decomposition, but the decomposition order is random: \begin{aligned} p(x_1,x_2,x_3,\dots,x_n)=&p(x_1) p(x_2|x_1) p(x_3|x_1,x_2) \dots p(x_n|x_1,x_2,\dots,x_{n-1})\\ =&p(x_3) p(x_1|x_3) p(x_2|x_3,x_1) \dots p(x_n|x_3,x_1,\dots,x_{n-1})\\ =&\dots\\ =&p(x_{n-1})p(x_1|x_{n-1})p(x_n|x_{n-1}, x_1)\dots p(x_2|x_{n-1}, x_1,\dots,x_3) \end{aligned}
In short, any "appearance order" of x_1, x_2, \dots, x_n is possible. In principle, each order corresponds to a model, so there are n! language models. A Transformer-based model can incorporate all these orders into a single model!
How is this achieved? Taking the generation of "Beijing welcomes you" as an example, assume a random generation order is "<s> \to Ying \to Jing \to Ni \to Huan \to Bei \to <e>". We only need to Mask the Attention matrix as shown in the second sub-figure below to achieve this:
Similar to the unidirectional language model, the 4th row has only one blue square, indicating that "Ying" can only be related to the start token <s>. The 2nd row has two blue squares, indicating that "Jing" can only be related to the start token <s> and "Ying". And so on. Intuitively, this is like shuffling the lower triangular Mask of the unidirectional language model.
In other words, implementing a language model with a specific order is equivalent to shuffling the original lower triangular Mask in a certain way. Because Attention provides an n\times n matrix, we have enough degrees of freedom to Mask this matrix in different ways to achieve diverse effects.
At this point, readers might have an implementation question: the shuffled Mask doesn’t seem to have a clear pattern. Do we have to randomly generate such a Mask matrix every time? In fact, there is a simpler, mathematically equivalent training scheme. This scheme stems from the fact that pure Attention models are inherently unordered; the word order is actually added through Position Embeddings. That is, we input not only the token itself but also the position ID of the token. In other words, you think you input the sequence "[Bei, Jing, Huan, Ying, Ni]", but you actually input the set "{(Bei, 1), (Jing, 2), (Huan, 3), (Ying, 4), (Ni, 5)}".
Since it is just a set and independent of order, we can completely change the input order. For example, for the order "<s> \to Ying \to Jing \to Ni \to Huan \to Bei \to <e>", we can input in the order "(Ying, 4), (Jing, 2), (Ni, 5), (Huan, 3), (Bei, 1)". That is, shuffle the tokens to "Ying, Jing, Ni, Huan, Bei" and input them into the Transformer, but the position of the 1st token is not 1, but 4, and so on. After this change, the Mask matrix can be restored to a lower triangular matrix. Therefore, one only needs to shuffle at the input level, which is much simpler to operate.
Seq2Seq
Now we come to our "main event": combining BERT and other Transformer architectures with Seq2Seq. Why is it the main event? Because in principle, any NLP problem can be converted into Seq2Seq; it is a truly universal model. So if Seq2Seq can be achieved, any task can theoretically be implemented.
There are two well-known works combining BERT with Seq2Seq: MASS and UNILM. Both are from Microsoft and were released in the same month! MASS still uses a standard Seq2Seq architecture, using BERT-like Transformer models for the encoder and decoder respectively; its main contribution is providing a Seq2Seq pre-training scheme. What is truly interesting is UNILM, which provides an elegant way to directly use a single BERT model for Seq2Seq tasks without distinguishing between encoder and decoder. Achieving this requires almost no effort—just a special Mask.
(Interlude: The actual sequence of events is that I independently thought of the idea of using a single BERT model for Seq2Seq two weeks ago, only to find that this idea had already been implemented as UNILM.)
UNILM treats Seq2Seq as sentence completion. If the input is "What do you want to eat?" and the target sentence is "White cut chicken," UNILM concatenates them into one: [CLS] What do you want to eat [SEP] White cut chicken [SEP]. After this transformation, the simplest scheme is to train a language model and input "[CLS] What do you want to eat [SEP]" to predict "White cut chicken" word by word until "[SEP]" appears, as shown in the left figure below:
However, the left figure is just the most basic scheme. It also includes "What do you want to eat" in the prediction range (making the Attention for this part unidirectional, i.e., the corresponding Mask matrix is lower triangular). In fact, this is unnecessary and constitutes an extra constraint. What truly needs to be predicted is only the "White cut chicken" part, so we can remove the Mask for the "What do you want to eat" part, resulting in the Mask in the right figure above.
In this way, the Attention for the input part is bidirectional, while the Attention for the output part is unidirectional, satisfying the requirements of Seq2Seq without extra constraints. This is the idea provided in UNILM for completing Seq2Seq tasks with a single BERT model. By adding the aforementioned Mask shape without modifying the model architecture, one can directly use BERT’s Masked Language Model pre-trained weights, leading to faster convergence. This fits the original intention of a universal model: "One BERT in hand, the world is mine." I personally believe this is a very elegant solution.
Experiment
In fact, the Masking schemes mentioned above have mostly been integrated into bert4keras, which I wrote. Readers can directly use bert4keras to load BERT pre-trained weights and call the aforementioned Masking schemes for corresponding tasks. Below, we provide an example of a fast-converging Seq2Seq model using the UNILM approach.
Open Source Code
The test task for this code is still the previous title generation. The code is adapted from "Playing with Keras: Seq2Seq Automatic Title Generation". Thanks to the encapsulation of bert4keras, the model implementation is very simple and clean. This time, the original THUCNews dataset was used directly. Readers can download the dataset and source code to test and reproduce it.
Details: task_seq2seq_autotitle.py
How effective is it? Experiments show that in the title generation task, starting from the first epoch (1000 iterations), it can already generate basically readable titles. Correspondingly, when using LSTM previously, it took dozens of times more iterations to achieve the same effect.
Brief Explanation
Below is a brief explanation of the key parts of the code.
First, the input format is still token_id and
segment_id. For example:
tokens = ['[CLS]', 'What', 'do', 'you', 'want', 'to', 'eat', '[SEP]', 'White', 'cut', 'chicken', '[SEP]']
token_ids = [token_dict[t] for t in tokens]
segment_ids = [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1]segment_ids are used to distinguish between the input
sentence and the target sentence. 0 corresponds to the input sentence,
and 1 corresponds to the target sentence. The built-in
tokenizer.encode can generate these token_id
and segment_id pairs.
As for building the model, it only takes a few lines:
model = build_transformer_model(
config_path,
checkpoint_path,
application='unilm',
keep_tokens=keep_tokens
)
model.summary()
y_in = model.input[0][:, 1:] # Target tokens
y_mask = model.input[1][:, 1:]
y = model.output[:, :-1] # Predicted tokens, shifted by one from target
# Cross-entropy as loss, masking out predictions for the input part
cross_entropy = K.sparse_categorical_crossentropy(y_in, y)
cross_entropy = K.sum(cross_entropy * y_mask) / K.sum(y_mask)Note that in build_transformer_model, as long as
application=’unilm’ is set, it will automatically load the
MLM part of BERT and pass in the corresponding Mask. The rest is just
writing the loss function. Additionally, there is a
keep_tokens parameter, which is used to streamline the
Embedding layer. For Chinese BERT, there are about 20,000 tokens in
total, which means predicting the generated token is a 20,000-class
classification problem. However, nearly half of the tokens are never
actually predicted. Therefore, this 20,000-class classification wastes
some computation. An option is provided here to maintain our own token
list and pass in the corresponding IDs to keep only those tokens,
thereby reducing the computational load (the streamlined version is
usually half the original size or even less).
The remaining steps involve decoding via beam search, which is no different from general Seq2Seq and will not be repeated here. Readers can refer to "Playing with Keras: Seq2Seq Automatic Title Generation" and the code.
Summary
This article systematically summarizes the Masking techniques for the Attention matrix in Transformer and provides an implementation of Seq2Seq using the UNILM scheme. For Seq2Seq text generation tasks in the same language, adopting the UNILM approach and loading BERT’s MLM pre-trained weights can effectively and quickly implement and improve generation results. It is worth a try.
When reposting, please include the original address: https://kexue.fm/archives/6933
For more detailed reposting matters, please refer to: "Scientific Space FAQ"