What is Two-Stream Self-Attention in XLNet | by Xu LIANG | Towards Data Science

In my previous post What is XLNet and why it outperforms BERT, I mainly talked about the difference between XLNet (AR language model) and BERT (AE language model) and the Permutation Language Modeling.
I believe that having an intuitive understanding of XLNet is far important than the implementation detail, so I only explained the Permutation Language Modeling and don’t mention another important part, the Two-Stream Self-Attention architecture. But as Jiaming Chen mentioned in the comment, the Two-Stream Self-Attention is another highlight in the XLNet paper, so I wrote this post to explain the Two-Stream Self-Attention as clearly as possible.
The content is structured as follows.
A quick review of Permutation Language Modeling
What problems the permutation brought?
Does BERT have such problem?
How XLNet solve the problem?
Attention mask: How XLNet implement permutation?
Special term:
AR language model: autoregressive language model
AE language model: autoencoder language model
In order to make this post more independent, here I give a brief summary of What is XLNet and why it outperforms BERT.
XLNet proposed that using Permutation Language Modeling to make the AR language model learn from bi-directional context. By this way, it can avoid the disadvantages brought by the MASK method in AE language model.
The permutation method is to get permutations of a sequence, and using the previous t-1 tokens as the context to predict the t-th position token. For example, we have a sentence[x1, x2, x3, x4]and x3 is the t-th position token that we want to predict. First, we get permutations of the sentence.
[('x1', 'x2', 'x3', 'x4'), ('x1', 'x2', 'x4', 'x3'), ('x1', 'x3', 'x2', 'x4'), ('x1', 'x3', 'x4', 'x2'), ('x1', 'x4', 'x2', 'x3'), ('x1', 'x4', 'x3', 'x2'), ('x2', 'x1', 'x3', 'x4'), ('x2', 'x1', 'x4', 'x3'), ('x2', 'x3', 'x1', 'x4'), ('x2', 'x3', 'x4', 'x1'), ('x2', 'x4', 'x1', 'x3'), ('x2', 'x4', 'x3', 'x1'), ('x3', 'x1', 'x2', 'x4'), ('x3', 'x1', 'x4', 'x2'), ('x3', 'x2', 'x1', 'x4'), ('x3', 'x2', 'x4', 'x1'), ('x3', 'x4', 'x1', 'x2'), ('x3', 'x4', 'x2', 'x1'), ('x4', 'x1', 'x2', 'x3'), ('x4', 'x1', 'x3', 'x2'), ('x4', 'x2', 'x1', 'x3'), ('x4', 'x2', 'x3', 'x1'), ('x4', 'x3', 'x1', 'x2'), ('x4', 'x3', 'x2', 'x1')]
Then we select some samples as the training data. (In the paper, it randomly selects samples from permutations)
('x1', 'x2', 'x4', 'x3'),('x1', 'x4', 'x3', 'x2'),('x2', 'x3', 'x4', 'x1'),('x4', 'x2', 'x3', 'x1'),('x3', 'x2', 'x4', 'x1'),
We can see every token have the chance that appears before x3. So the AR model can learn the bi-directional information from these context tokens. If you are still not clear about permutations, you could read the previous post.
The permutation can make AR model see the context from two directions, but it also brought problems that the original transformer cannot solve.
Z: a factorization order
p_θ: likelihood function
x_zt: the t-th token in the factorization order
x_z<t: the tokens before t-th token
This is the objective function for permutation language modeling, which means takes t-1 tokens as the context and to predict the t-th token.
There are two requirements that a standard Transformer cannot do:
to predict the token x_t, the model should only see the position of x_t, not the content of x_t (I will explain what is content in the next section)to predict the token x_t, the model should encode all tokens before x_t as the content
to predict the token x_t, the model should only see the position of x_t, not the content of x_t (I will explain what is content in the next section)
to predict the token x_t, the model should encode all tokens before x_t as the content
Especially the first requirements, transformer incorporate position encoding into the token embedding. So it cannot separate the position information from the token embedding.
BERT is an AE language model, it does not need separate position information like the AR language model. Unlike the XLNet need position information to predict t-th token, BERT uses [MASK] to represent which token to predict (we can think [MASK] is just a placeholder). For example, if BERT use x2, x1 and x4 to predict x3, the embedding of x2, x1, x4 contains the position information and other information related to [MASK]. So the model has a high chance to predict that [MASK] is x3.
Here I will give a more detail explanation about information. BERT embedding(information learned by BERT) contains two kinds of information, the position information, and the content information (I just split it into two parts for simplicity).
The position information is easy to understand that it tells the model the position of the current token. The content information (semantics and syntactic) contains the “meaning” of the current token. An intuitive example you have seen is kind — man + woman = queen .
XLNet proposes Two-Stream Self-Attention to solve the problem.
As the name indicates, it contains two kinds of self-attention. One is the content stream attention, which is the standard self-attention in Transformer. Another is the query stream attention. XLNet introduces it to replace the [MASK] token in BERT.
For example, if BERT wants to predict x3 with knowledge of the context words x1 and x2, it can use [MASK] to represent the x3 token. The [MASK] is just a placeholder. And the embedding of x1 and x2 contains the position information to help the model to “know” [MASK] is x3.
Things are different come to XLNet. One token x3 will server two kinds of roles. When it is used as content to predict other tokens, we can use the content representation (learned by content stream attention) to represent x3. But if we want to predict x3, we should only know its position and not its content. That’s why XLNet uses query representation (learned by query stream attention) to preserve context information before x3 and only the position information of x3.
In order to intuitively understand the Two-Stream Self-Attention, we can just think XLNet replace the [MASK] in BERT with query representation. They just choose different approaches to do the same thing.
When I first read the paper, I can not stop wondering about the implementation details about permutation during training. So I will talk this a little in case you are interested.
In the first part, “A quick review of Permutation Language Modeling”, I give an example that permutations of a sentence[x1, x2, x3, x4] will look like below.
[('x1', 'x2', 'x3', 'x4'), ('x1', 'x2', 'x4', 'x3'), ('x1', 'x3', 'x2', 'x4'), ('x1', 'x3', 'x4', 'x2'), ('x1', 'x4', 'x2', 'x3'), ('x1', 'x4', 'x3', 'x2'), ('x2', 'x1', 'x3', 'x4'), ('x2', 'x1', 'x4', 'x3'), ('x2', 'x3', 'x1', 'x4'), ('x2', 'x3', 'x4', 'x1'), ('x2', 'x4', 'x1', 'x3'), ('x2', 'x4', 'x3', 'x1'), ('x3', 'x1', 'x2', 'x4'), ('x3', 'x1', 'x4', 'x2'), ('x3', 'x2', 'x1', 'x4'), ('x3', 'x2', 'x4', 'x1'), ('x3', 'x4', 'x1', 'x2'), ('x3', 'x4', 'x2', 'x1'), ('x4', 'x1', 'x2', 'x3'), ('x4', 'x1', 'x3', 'x2'), ('x4', 'x2', 'x1', 'x3'), ('x4', 'x2', 'x3', 'x1'), ('x4', 'x3', 'x1', 'x2'), ('x4', 'x3', 'x2', 'x1')]
It is very easy to misunderstand that we need to get the random order of a sentence and input it into the model. But this is not true. The order of input sentence is [x1, x2, x3, x4] , and XLNet use attention mask to permute the factorization order.
The original order of the sentence is [x1, x2, x3, x4]. And we randomly get a factorization order as [x3, x2, x4, x1].
The upper left corner is the calculation of content representation. If we want to predict the content representation of x1, we should have all 4 tokens content information. KV = [h1, h2, h3, h4] and Q = h1.
The lower-left corner is the calculation of query representation. If we want to predict the query representation of x1, we cannot see the content representation of x1 itself. KV = [h2, h3, h4] and Q = g1.
The right corner is the whole calculation process. I explain it from bottom to top. First the h and g are initialized as e(xi) and w. And after the content mask and query mask, the two-stream attention will output the first layer output h^(1) and g^(1) and then calculate the second layer.
Notice the right content mask and query mask. Both of them are matrices. In the content mask, the first row has 4 red points. It means that the first token (x1) can see (attend to) all other tokens including itself( x3->x2->x4->x1). The second row has 2 red points. It means that the second token (x2) can see (attend to) two tokens (x3->x2). And so on other rows.
The only difference between the content mask and query mask is those diagonal elements in the query mask are 0, which means the tokens cannot see themselves.
Let’s sum it up. The input sentence has only one order. But we can use different attention mask to implement different factorization order.
In this post, I mainly explained what is the problem that XLNet face and how to use Two-Stream Self-Attention to solve it. I also mention some implement details about attention mask for the permutation. As for the result comparison, you can find the latest post from XLNet team, which do a more fair comparison than the paper does.
Check out my other posts on Medium with a categorized view!GitHub: BrambleXuLinkedIn: Xu LiangBlog: BrambleXu
paper: https://arxiv.org/abs/1906.08237
implementation: pytorch_transformers/modeling_xlnet.py
What is XLNet and why it outperforms BERT
Paper Dissected: “XLNet: Generalized Autoregressive Pretraining for Language Understanding” Explained, this mainly talked about transformer-xl.
A Fair Comparison Study of XLNet and BERT with Large Models