From Language Models to Seq2Seq: Transformer Is Like an Actor — It's All About the Mask
Over the past year (especially the past six months), I'm sure everyone has frequently come across reports on various Transformer-related work (Bert, GPT, XLNet, and so on), along with a constant stream of new records on all sorts of benchmark tasks. Meanwhile, there have also been plenty of blog posts and columns popularizing and explaining these models.
As the saying goes, "laypeople watch the spectacle, insiders study the tricks." We shouldn't just understand these works at the level of "what they are" — we also need to think about "why." This "why" doesn't just mean "why do it this way," but also "why is it even possible to do it this way." For instance, when discussing XLNet's permutation language model, we may already understand from many introductions the benefits of a permutation language model, but let's go a step further and ask:
Why can Transformers implement a permutation language model? How is it done? Can RNNs do it too?
This post analyzes, from the angle of masking the attention matrix, the fundamental reason why so many Transformer models can pull off such dazzling tricks. As the title says, "Transformer is like an actor — it's all about the mask," which is one of the key "tricks" behind the many flavors of Transformer models out there.
By the end of this post, you should hopefully understand:
1. The relationship between the way the attention matrix is masked and various pretraining schemes;
2. How to directly use a pretrained Bert model to do Seq2Seq tasks.
Background
Since Attention is All You Need, Transformer models based purely on attention have gradually become popular, and the emergence of Bert pushed this trend to a new height. After that, one after another, works based on large-scale pretrained Transformer models kept appearing — some applying off-the-shelf models, some trying to better explain and visualize these models, and others improving the architecture or the pretraining method to get better results. In short, this stream of pretraining-based work has been endless, almost dazzling in its variety. To some extent, if you still haven't fine-tuned Bert, you could say you've fallen behind mainstream NLP practice.
Fancy Pretraining Schemes
As everyone knows, the traditional way of pretraining a model is via a language model. For example, ELMo uses a BiLSTM as its base architecture, pretraining the LSTMs in both directions with a language model in each direction respectively; OpenAI's GPT and GPT-2 later stuck faithfully to the time-honored (standard, unidirectional) language model for pretraining.
However, there are even more elaborate ways to pretrain. Bert, for example, used what's called a "Masked Language Model," which is really just a variant of the ordinary language model; XLNet went further and proposed a more thorough "Permutation Language Modeling," which we might call a "permuted language model"; and there's also the UNILM model, which directly uses a single Bert-like architecture to do Seq2Seq — you can think of it as a pretraining method, or you could just use it directly to do Seq2Seq tasks...
With all this variety, one can't help but wonder: why is it precisely in the era when Transformers became popular that we see this "let a hundred flowers bloom" phenomenon of large-scale pretrained models?
Exclusive to Transformers
In fact, apart from the unidirectional language model and its simple variant, the masked language model, both UNILM's Seq2Seq pretraining and XLNet's permutation language model pretraining are essentially custom-tailored for the Transformer architecture. To put it bluntly, if you were using an RNN architecture, you simply couldn't pretrain it using a permutation language model approach; and as for Seq2Seq pretraining, you'd need to introduce two separate models (an encoder and a decoder) simultaneously, rather than being able to handle everything with a single model as the Transformer architecture allows.
The secret behind this lies mainly in the attention matrix. Attention essentially computes pairwise similarities across the input, forming a similarity matrix of size $n^2$ (i.e., the attention matrix, where $n$ is the sentence length; throughout this post, "attention" refers to self-attention). This means its memory footprint is on the order of $\mathcal{O}(n^2)$, whereas RNN and CNN models are only $\mathcal{O}(n)$ — so in practice attention usually consumes more GPU memory. However, every coin has two sides: a larger footprint also means more possibilities. We can impose all sorts of prior constraints on this $\mathcal{O}(n^2)$-scale attention matrix, enabling it to handle much more flexible tasks. Put simply, only a purely attention-based model has enough "capacity" to accommodate so many "tricks."
And the way we impose these prior constraints is by masking the attention matrix in different forms — which is exactly what this post focuses on.
Analysis
I already gave a basic introduction to attention in A Brief Read of "Attention is All You Need" (Introduction + Code), so here I'll just do a quick recap. The mathematical form of attention is:
\begin{equation}Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}) = softmax\left(\frac{\boldsymbol{Q}\boldsymbol{K}^{\top}}{\sqrt{d_k}}\right)\boldsymbol{V}\end{equation}
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 vector sequences of the query, key, and value respectively. We can think of key and value as being paired one-to-one, while $\boldsymbol{Q}\boldsymbol{K}^{\top}$ takes the pairwise dot products between the query and key vectors, then normalizes them with $softmax$, giving us an attention matrix of shape $l_q\times l_k$. This describes the strength of association between any pair of elements from the query and the key — and everything we discuss below is really about tinkering with this attention matrix. Finally, it's multiplied with $\boldsymbol{V}$, which amounts to a weighted sum of the vectors in $\boldsymbol{V}$ according to these association strengths, ultimately producing an output vector sequence of shape $l_q\times d_v$.
The most commonly used form of attention nowadays is self-attention, where $\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}$ are all linear transformations of the same vector sequence, and the Transformer is a combination of self-attention with a position-wise fully-connected layer (equivalent to a 1D convolution with kernel size 1). So a Transformer is essentially an attention-based transformation from a vector sequence to another vector sequence.
In this section, we'll take a fairly detailed look at how the attention matrix is masked, corresponding respectively to how unidirectional language models, permutation language models, and Seq2Seq are implemented.
Unidirectional Language Models
A language model can be thought of as an unconditional text-generation model. If you're not yet familiar with text-generation models, you can look up related material, and refer to Playing with Keras' seq2seq for Automatic Title Generation to build intuition. A unidirectional language model essentially "memorizes" the training corpus through the following conditional probability decomposition:
\begin{equation}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})\end{equation}
What we usually mean by "language model" refers to a unidirectional one (in the narrower sense, specifically a forward one). The key point of a language model is that it must prevent seeing "future information." As in the formula above, when predicting $x_1$, there is no external input at all; when predicting $x_2$, only $x_1$ can be fed in; when predicting $x_3$, only $x_1,x_2$ can be fed in; and so on.
Illustration of a unidirectional language model. Each predicted token depends only on the tokens before it.
RNNs are naturally suited to being language models, because their computation is itself recursive. If you wanted to use a CNN instead, you'd need to mask the convolution kernel, i.e., zero out the part of the kernel corresponding to the right-hand side. What about a Transformer? There, you need an attention matrix in the form of a lower-triangular matrix:
Masking scheme for a unidirectional (forward) language model
As shown in the figure, each row of the attention matrix actually represents an output, and each column represents an input, and the attention matrix expresses the relationship between outputs and inputs. Suppose the white cells all represent 0; then the first row shows that "北" (North) can only be related to the start token \<s\>, and the second row shows that "京" (Capital) can only be related to the start token \<s\> and "北," and so on. So, all we need to do is introduce a lower-triangular mask into the Transformer's attention matrix, and train with the input and output shifted by one position, to implement a unidirectional language model. (For how the mask itself is implemented, see the "Mask" section of "Let Keras Be a Bit Cooler!": Layers Within Layers, and Masking.)
Permutation Language Models
The permutation language model is a concept proposed by XLNet, mainly used for XLNet's pretraining. Speaking of XLNet, I find its permutation-language-model pretraining scheme quite interesting, but I'm not so fond of the fact that it switched its base architecture to Transformer-XL. I think whoever has the compute resources should try the combination of "Bert + permutation-language-model pretraining" — it might yield some surprising results.
Like the ordinary language model, the permutation language model also performs a conditional-probability decomposition, but the order of decomposition is random:
\begin{equation}\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}\end{equation}
In short, any "order of appearance" of $x_1,x_2,\dots,x_n$ is possible. In principle, each ordering corresponds to a different model, so in principle there would be $n!$ different language models. But a Transformer-based model can fold all of these orderings into a single model!
How is that achieved? Let's again take the generation of "北京欢迎你" (Beijing welcomes you) as an example. Suppose a randomly chosen generation order is "\<s\> → 迎 → 京 → 你 → 欢 → 北 → \<e\>." Then we only need to mask the attention matrix in the way shown in the second sub-figure below, and we achieve our goal:
Mask for the forward language model
Mask for the permutation language model
Mask for the reverse language model
Similar to the unidirectional language model before, row 4 has only one blue cell, indicating that "迎" (welcome) can only be related to the start token \<s\>, while row 2 has two blue cells, indicating that "京" can only be related to the start token \<s\> and "迎," and so on. Intuitively, this looks like the lower-triangular mask of the unidirectional language model has simply been "shuffled."
In other words, implementing a language model with some specific ordering is equivalent to shuffling the original lower-triangular mask in a particular way. Precisely because attention gives us an attention matrix of shape $n\times n$, we have enough degrees of freedom to mask this matrix in different ways and thereby achieve a diverse range of effects.
At this point, readers might have an implementation-level question: the shuffled mask doesn't seem to follow any obvious pattern — do we really need to randomly generate such a seemingly patternless mask matrix every single time? In fact, there's a simpler training scheme that's mathematically equivalent. This scheme stems from the fact that a purely attention-based model is, at its core, an order-agnostic model — the word order is only introduced via the position embedding that gets added on. In other words, what we input is not just the tokens themselves, but also the position id of each token. Put differently, while you might think you're feeding in the sequence "[北, 京, 欢, 迎, 你]," what you're actually feeding in is the set "{(北, 1), (京, 2), (欢, 3), (迎, 4), (你, 5)}."
Reordering so that a forward language model can implement a permutation language model
Since it's really just a set, independent of order, we can just as well feed it in a different order. Take the earlier example "\<s\> → 迎 → 京 → 你 → 欢 → 北 → \<e\>": we can feed it in the order "(迎, 4), (京, 2), (你, 5), (欢, 3), (北, 1)" — that is, we shuffle the tokens into "迎,京,你,欢,北" and feed that into the Transformer, but the position of the first token is no longer 1, it's 4; and so on. After this reordering, the mask matrix reverts back to being lower-triangular, so all we need is to shuffle things at the input level, which is much simpler to implement.
Seq2Seq
Now we come to the "main event": combining Bert-style Transformer architectures with Seq2Seq. Why call it the main event? Because in principle, any NLP problem can be recast as a Seq2Seq problem — it's a truly universal model. So if we can pull off Seq2Seq, in theory we can handle any task at all.
There are two well-known pieces of work combining Bert with Seq2Seq: MASS and UNILM, both from Microsoft, and both published in the same month, no less! MASS uses an ordinary Seq2Seq architecture, employing Bert-like Transformer models as both the encoder and the decoder; its main contribution is providing a Seq2Seq-style pretraining scheme. What's really interesting, though, is UNILM, which offers an elegant way to let us do Seq2Seq tasks with a single Bert model directly, without distinguishing between an encoder and a decoder. And achieving this takes almost no effort at all — just one special mask.
(Aside: as it happens, I independently came up with the idea of doing Seq2Seq with a single Bert model a couple of weeks before writing this, and then when I went looking through the literature, I found this idea had already been done — and it was exactly UNILM.)
UNILM treats Seq2Seq directly as sentence completion. Suppose the input is "你想吃啥" (What do you want to eat) and the target sentence is "白切鸡" (poached chicken). UNILM concatenates the two sentences into one: [CLS] 你 想 吃 啥 [SEP] 白 切 鸡 [SEP]. Once transformed this way, the simplest approach is to train a language model, then feed in "[CLS] 你 想 吃 啥 [SEP]" and predict "白 切 鸡" character by character until "[SEP]" appears — as shown in the figure on the left below:
Doing Seq2Seq via a unidirectional language model
Designing a more suitable mask for Seq2Seq
However, the left figure is only the most naive scheme — it also puts "你想吃啥" within the range to be predicted (which means the attention over that part is unidirectional, i.e., the corresponding portion of the mask matrix is lower-triangular). In fact, this is unnecessary and constitutes an extraneous constraint. Only "白切鸡" actually needs to be predicted, so we can remove the mask constraint over "你想吃啥," yielding the mask shown in the right figure above.
This way, the attention over the input portion is bidirectional, while the attention over the output portion is unidirectional, satisfying the requirements of Seq2Seq without any unnecessary constraints. This is precisely the idea UNILM offers, letting us accomplish a Seq2Seq task with a single Bert model, simply by adding a mask of the above shape — no need to modify the model architecture at all — and moreover, it lets us directly reuse Bert's pretrained Masked-Language-Model weights, leading to faster convergence. This fits perfectly with the philosophy of "with one Bert in hand, the world is mine," a universal model. Personally, I think this is an extremely elegant solution.
Illustration of the UNILM Seq2Seq model. The input portion can attend bidirectionally internally, while the output portion only attends unidirectionally.
Experiments
In fact, the masking schemes above have all basically been integrated into bert4keras, which I wrote myself. Readers can use bert4keras directly to load Bert's pretrained weights and invoke the masking schemes above for the corresponding tasks. Below, I'll give an example of using the UNILM approach to build a Seq2Seq model that converges quickly.
Open-Source Code
The test task for this round of code is once again title generation, adapted from the code in Playing with Keras' seq2seq for Automatic Title Generation. Thanks to the encapsulation provided by bert4keras, the model-building code is very clean and simple this time. This time, I used the raw THUCNews dataset directly; readers can download the dataset and source code themselves to reproduce the results.
See here for details: task_seq2seq_autotitle.py
How good are the results? In experiments, on the title-generation task, starting from the very first epoch (1000 iterations), the model was already able to generate basically readable titles. By comparison, when I did this with LSTMs before, it typically took dozens of times more iterations to achieve the same effect.
Basically readable generation results are already obtained after the first epoch (1000 steps)
Brief Notes
Below is a brief explanation of the key parts of the code.
First, the input format is still fed in as token_id and segment_id, for example:
tokens = ['[ClS]', u'你', u'想', u'吃', u'啥', '[SEP]', u'白', u'切', u'鸡', '[SEP]']
token_ids = [token_dict[t] for t in tokens]
segment_ids = [0, 0, 0, 0, 0, 0, 1, 1, 1, 1]
segment_ids is used to distinguish the input sentence from the target sentence, with 0 corresponding to the input sentence and 1 to the target sentence. You just need the built-in tokenizer.encode to generate this token_id and segment_id.
As for building the model, it only takes a handful of lines:
model = build_transformer_model(
config_path,
checkpoint_path,
application='unilm',
keep_tokens=keep_tokens
)
model.summary()
y_in = model.input[0][:, 1:] # 目标tokens
y_mask = model.input[1][:, 1:]
y = model.output[:, :-1] # 预测tokens，预测与目标错开一位
# 交叉熵作为loss，并mask掉输入部分的预测
cross_entropy = K.sparse_categorical_crossentropy(y_in, y)
cross_entropy = K.sum(cross_entropy * y_mask) / K.sum(y_mask)
Note that within build_transformer_model, you just need to set application='unilm', and it will automatically load Bert's MLM component and pass in the corresponding mask — after that, all that's left is to write the loss function. There's also a keep_tokens option, used to streamline the embedding layer. For Chinese Bert, there are roughly 20,000 tokens in total, which means the final token-prediction step is a 20,000-way classification problem. But in practice, nearly half of those tokens will never actually be predicted (in principle, they shouldn't be), so this 20,000-way classification wastes some computation. So here I provide an option: we can maintain our own token list and pass in the corresponding ids, keeping only these tokens, which reduces the amount of computation (after streamlining, it's typically only half the original size, or even less).
What remains is decoding via beam search and so on, no different from an ordinary Seq2Seq setup, so I won't go over it again in detail — just refer to Playing with Keras' seq2seq for Automatic Title Generation and the code.
Summary
This post gives a relatively systematic overview of masking tricks for the attention matrix in Transformers, and presents an implementation of Seq2Seq using the UNILM approach. For text-generation tasks that are Seq2Seq within the same language, using the UNILM approach to load Bert's pretrained MLM weights is an effective and fast way to implement and improve generation quality — well worth trying.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.