A Light Reading of "Attention is All You Need" (Introduction + Code)
In mid-2017, there were two papers that were somewhat similar and that I personally admired a great deal: Facebook's Convolutional Sequence to Sequence Learning and Google's Attention is All You Need. Both can be viewed as innovations on the Seq2Seq framework, and both, in essence, abandon the RNN structure for Seq2Seq tasks.
In this post, I'll give a brief analysis of Attention is All You Need. Naturally, since both papers have received a great deal of attention, there are already many write-ups online (though most of them are essentially direct translations of the paper, with little of the authors' own understanding). So here I'll try to add as much of my own thinking as possible, and avoid repeating what other bloggers have already said.
Sequence Encoding
The typical approach for deep learning in NLP is to first tokenize the sentence and then convert each word into a corresponding word embedding, forming a sequence. This way, every sentence corresponds to a matrix $\boldsymbol{X}=(\boldsymbol{x}_1,\boldsymbol{x}_2,\dots,\boldsymbol{x}_t)$, where $\boldsymbol{x}_i$ represents the word vector (a row vector) of the $i$-th word, with dimension $d$, so $\boldsymbol{X}\in \mathbb{R}^{n\times d}$. The problem then becomes how to encode these sequences.
The first basic idea is an RNN layer. The RNN scheme is simple: proceed recursively:
\begin{equation}\boldsymbol{y}_t = f(\boldsymbol{y}_{t-1},\boldsymbol{x}_t)\end{equation}
Whether it's the now widely used LSTM, GRU, or the more recent SRU, none of them escape this recursive framework. RNNs are structurally simple and well suited to sequence modeling, but one obvious drawback of RNNs is that they cannot be parallelized, so they tend to be slow — this is an inherent defect of recursion. Also, in my personal opinion, RNNs do not do a great job of learning global structural information, since they are essentially a Markov decision process. more
The second idea is a CNN layer. The CNN scheme is also quite natural — a sliding window traversal. For example, a convolution of kernel size 3 looks like:
\begin{equation}\boldsymbol{y}_t = f(\boldsymbol{x}_{t-1},\boldsymbol{x}_t,\boldsymbol{x}_{t+1})\end{equation}
In Facebook's paper, pure convolution was used to accomplish Seq2Seq learning, and it's a refined and extreme showcase of what convolution can do — readers who are fans of CNNs really should read that paper carefully. CNNs are convenient to parallelize and are good at capturing certain kinds of global structural information. I myself am somewhat partial to CNNs, and in my current work or competition models, I've been trying to replace existing RNN models with CNNs wherever possible, and have built up my own set of practical experience along the way — but that's a topic for another time.
Google's masterpiece offers a third idea: pure attention! Attention alone is all you need! RNNs must proceed step by step recursively to obtain global information, which is why bidirectional RNNs are generally preferred; CNNs, in fact, can only capture local information, and rely on stacking layers to enlarge the receptive field. Attention's approach is the most direct of all — it grabs global information in one step! Its solution is:
\begin{equation}\boldsymbol{y}_t = f(\boldsymbol{x}_{t},\boldsymbol{A},\boldsymbol{B})\end{equation}
where $\boldsymbol{A},\boldsymbol{B}$ is another sequence (matrix). If both are taken to be $\boldsymbol{A}=\boldsymbol{B}=\boldsymbol{X}$, then this is called self-attention, meaning that $\boldsymbol{x}_t$ is directly compared with each word in the original sequence, and finally $\boldsymbol{y}_t$ is computed!
The Attention Layer
Definition of Attention
Attention
Google's generalized notion of attention is also a scheme for encoding sequences, so we can regard it, like RNNs and CNNs, as just another type of layer for sequence encoding.
What was given above was a description in a generalized, framework-level form; in fact, the scheme Google actually gave is quite concrete. First, they present the definition of attention:
\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 I use notation consistent with Google's paper, where $\boldsymbol{Q}\in\mathbb{R}^{n\times d_k}, \boldsymbol{K}\in\mathbb{R}^{m\times d_k}, \boldsymbol{V}\in\mathbb{R}^{m\times d_v}$. If we ignore the activation function $softmax$, then this is really just a multiplication of three matrices of shape $n\times d_k,d_k\times m, m\times d_v$, and the final result is a matrix of shape $n\times d_v$. So we can think of this as an attention layer that encodes a sequence $n\times d_k$ of shape $\boldsymbol{Q}$ into a new sequence of shape $n\times d_v$.
How should we understand this structure? Let's look at it vector by vector.
\begin{equation}Attention(\boldsymbol{q}_t,\boldsymbol{K},\boldsymbol{V}) = \sum_{s=1}^m \frac{1}{Z}\exp\left(\frac{\langle\boldsymbol{q}_t, \boldsymbol{k}_s\rangle}{\sqrt{d_k}}\right)\boldsymbol{v}_s\end{equation}
where $Z$ is a normalizing factor. In fact, $q,k,v$ are shorthand for $query,key,value$, and $\boldsymbol{K},\boldsymbol{V}$ are in one-to-one correspondence — they behave like key-value pairs. The meaning of the above formula is: using $\boldsymbol{q}_t$ as the query, we compute its inner product with each $\boldsymbol{k}_s$ and apply softmax to obtain the similarity between $\boldsymbol{q}_t$ and each $\boldsymbol{v}_s$, and then take a weighted sum to obtain a vector of dimension $d_v$. The factor $\sqrt{d_k}$ serves a regularizing purpose, keeping the inner products from becoming too large (if they're too large, the softmax output becomes essentially 0 or 1, losing its "softness").
In fact, this definition of attention is nothing new, but given Google's influence, we can regard this as a more formal proposal of the concept, treating it explicitly as a layer. Also, this definition is just one particular form of attention — there are other options, e.g., the operation between $query$ and $key$ need not be a dot product (it could instead be concatenation followed by an inner product with a parameter vector), and the weights don't even need to be normalized, and so on.
Multi-Head Attention
Multi-Head Attention
This is a new concept proposed by Google, refining the attention mechanism. But in terms of its actual form, it's about as simple as it gets: map $\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}$ through parameter matrices, then perform attention, repeat this process $h$ times, and concatenate the results — truly an example of "great simplicity from great sophistication." Concretely:
\begin{equation}head_i = Attention(\boldsymbol{Q}\boldsymbol{W}_i^Q,\boldsymbol{K}\boldsymbol{W}_i^K,\boldsymbol{V}\boldsymbol{W}_i^V)\end{equation}
Here $\boldsymbol{W}_i^Q\in\mathbb{R}^{d_k\times \tilde{d}_k}, \boldsymbol{W}_i^K\in\mathbb{R}^{d_k\times \tilde{d}_k}, \boldsymbol{W}_i^V\in\mathbb{R}^{d_v\times \tilde{d}_v}$, and then
\begin{equation}MultiHead(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}) = Concat(head_1,...,head_h)\end{equation}
Finally, we obtain a sequence of shape $n\times (h\tilde{d}_v)$. What's called "multi-head" is simply doing the same thing several times (without sharing parameters) and then concatenating the results.
Self-Attention
So far, the description of the attention layer has been fully general, and we can now instantiate some concrete applications. For example, in reading comprehension, $\boldsymbol{Q}$ could be the vector sequence of the passage, and we take $\boldsymbol{K}=\boldsymbol{V}$ to be the vector sequence of the question, so the output would be what's known as the Aligned Question Embedding.
In Google's paper, most of the attention used is self-attention, also called internal attention.
So-called self-attention is simply $Attention(\boldsymbol{X},\boldsymbol{X},\boldsymbol{X})$, where $\boldsymbol{X}$ is the input sequence mentioned earlier. That is, attention is performed within the sequence itself, seeking connections between elements of the sequence. One of the main contributions of Google's paper is showing that internal attention plays a substantial role in encoding sequences for machine translation (and even for Seq2Seq tasks in general), whereas previous research on Seq2Seq had largely restricted the use of attention to the decoder side. Similarly, R-Net, the model that has topped the SQuAD reading comprehension leaderboard, also incorporates self-attention, which has boosted its performance as well.
To be more precise, what Google actually uses is self multi-head attention:
\begin{equation}\boldsymbol{Y}=MultiHead(\boldsymbol{X},\boldsymbol{X},\boldsymbol{X})\end{equation}
Position Embedding
However, a moment's thought reveals that a model of this kind cannot capture the order of a sequence! In other words, if we shuffle the rows of $\boldsymbol{K},\boldsymbol{V}$ (equivalent to shuffling the word order in a sentence), the result of attention remains unchanged. This shows that, so far, the attention model is at best an extremely sophisticated "bag-of-words model."
This is a rather serious problem. As we all know, for time series — and especially for tasks in NLP — order carries very important information, representing local and even global structure. Failing to learn ordering information will significantly hurt performance (e.g., in machine translation, it's possible to translate every individual word correctly yet fail to assemble them into a coherent sentence).
So Google pulls out another trick: Position Embedding, which numbers each position and associates each number with a vector. By combining position embeddings with word embeddings, every word is given a certain amount of positional information, allowing attention to distinguish words at different positions.
Position Embedding isn't exactly new — Facebook's Convolutional Sequence to Sequence Learning also uses it. But in Google's work, their Position Embedding has a few distinguishing features:
1. In previous RNN and CNN models, Position Embedding did appear, but there it was more of an icing-on-the-cake supplementary technique — "better with it, only slightly worse without it" — because RNNs and CNNs can already capture positional information on their own. But in this pure attention model, Position Embedding is the sole source of positional information, so it is one of the model's core components, not merely a supplementary trick.
2. In previous uses of Position Embedding, the vectors were basically trained according to the task. Google instead directly gives a formula for constructing the Position Embedding:
\begin{equation}\left\{\begin{aligned}&PE_{2i}(p)=\sin\Big(p/10000^{2i/{d_{pos}}}\Big)\\ > &PE_{2i+1}(p)=\cos\Big(p/10000^{2i/{d_{pos}}}\Big) > \end{aligned}\right.\end{equation}
The meaning here is that the position with id $p$ is mapped to a $d_{pos}$-dimensional position vector, whose $i$-th component has the value $PE_i(p)$. Google states in the paper that they compared directly-trained position vectors against those computed from the above formula, and the results were close. Given that, we might as well prefer the formula-based Position Embedding, which we call the sinusoidal form of Position Embedding.
3. Position Embedding itself carries absolute positional information, but in language, relative position also matters a great deal. One important reason Google chose the aforementioned position vector formula is that, since we have $\sin(\alpha+\beta)=\sin\alpha\cos\beta+\cos\alpha\sin\beta$ as well as $\cos(\alpha+\beta)=\cos\alpha\cos\beta-\sin\alpha\sin\beta$, this shows that the vector for position $p+k$ can be expressed as a linear transformation of the vector for position $p$, which offers the possibility of expressing relative positional information.
There are a few options for combining position embeddings with word embeddings: they can be concatenated into a new vector, or the position embedding can be defined to have the same dimensionality as the word embedding, and the two added together. Both Facebook's paper and Google's paper use the latter approach. Intuitively, addition seems like it should cause information loss and thus seem undesirable, but Google's results show that addition is actually a fine approach. It seems my own understanding still isn't deep enough.
Also, although the Position Embedding given in the paper is presented in an interleaved $\sin,\cos$ form, this interleaving actually has no particular significance — you can rearrange it in any way you like (e.g., putting the first $\sin$ elements and then the last $\cos$ elements together), for the following reasons:
1. Suppose your Position Embedding is concatenated with the original word embedding — whether $\cos$ and $\sin$ are joined in interleaved order or in sequential blocks makes no difference at all, because the next step is simply a linear transformation matrix in either case;
2. If your Position Embedding is added to the original word embedding, the two approaches seem to differ slightly, but note that the word embedding itself has no local structure — that is, for a 50-dimensional word embedding, permuting the dimensions (as long as you apply the same permutation consistently) leaves it equivalent to the original word embedding. Since the object being added to (the word embedding) has no local structure to begin with, there's no need to insist on local structure (i.e., interleaved concatenation) for the object being added (the Position Embedding) either.
Some Shortcomings
At this point, the attention mechanism has largely been introduced. The advantage of an attention layer is that it can capture global connections in one step, since it directly compares the sequence pairwise (at the cost of increasing the computational complexity to $\mathcal{O}(n^2)$ — though since this is pure matrix computation, the actual cost isn't too severe). By comparison, RNNs need step-by-step recursion to capture such connections, and CNNs need to stack layers to expand the receptive field — this is a clear advantage of the attention layer.
The remaining part of Google's paper describes how attention is applied to machine translation, which is more a matter of application and hyperparameter tuning that we won't focus on particularly here. Of course, Google's results show that using pure attention for machine translation achieves state-of-the-art performance — that result is indeed impressive.
Nevertheless, I'd like to discuss some shortcomings of this paper itself and of the attention layer.
1. The paper's title is Attention is All You Need, and accordingly the paper deliberately avoids mentioning RNNs or CNNs — but I find this a bit too deliberate. In fact, the paper even coins a special name, "Position-wise Feed-Forward Networks," when in reality this is just a 1D convolution with kernel size 1. It feels a bit like renaming something purely to avoid saying "convolution," which strikes me as slightly disingenuous. (Though perhaps I'm reading too much into it.)
2. Although attention has no direct connection to CNNs, it in fact borrows heavily from CNN ideas. For instance, Multi-Head Attention is essentially doing attention multiple times and concatenating the results, which is exactly the same idea as having multiple convolution kernels in a CNN; the paper also uses residual connections, which likewise originate from CNN architectures.
3. Attention cannot model positional information particularly well, and this is a real weakness. Although Position Embedding can be introduced, I consider this only a mitigation rather than a fundamental solution. For example, training a text classification model or a machine translation model with this kind of pure attention mechanism should work reasonably well, but training a sequence labeling model (e.g., tokenization, named entity recognition) with it doesn't work nearly as well. So why does it work well for machine translation? I think the reason is that machine translation doesn't particularly emphasize strict word order, so the positional information provided by Position Embedding is already sufficient; moreover, BLEU, the evaluation metric for translation, doesn't heavily penalize word-order deviations either.
4. Not every problem requires long-range, global dependencies — many problems depend only on local structure, and in these cases pure attention isn't a great fit either. Google seems to have been aware of this issue too, since the paper mentions a restricted version of self-attention (though it doesn't appear to be used in the main experiments), which assumes that the current word is only connected to the $r$ words before and after it, so attention only occurs among these $2r+1$ words, making the computational cost $\mathcal{O}(nr)$. This also lets the model capture local structure in the sequence. But clearly, this is just the notion of a convolutional window from CNNs!
Through this discussion, we can appreciate that treating attention as a standalone layer, to be mixed with CNN and RNN structures, would likely make fuller use of the respective strengths of each — rather than what Google's paper claims with "Attention is All You Need," which is, frankly, a bit of an overcorrection (a rather bold claim), and one that, in practice, doesn't quite hold up. Given the actual scope of the work, perhaps a more modest title like "Attention is All Seq2Seq Needs" (which, admittedly, is still a fairly bold claim) would have earned more universal agreement.
Code Implementation
Finally, to give this post some practical value, I've tried to provide an implementation of the paper's Multi-Head Attention. Readers who need it are welcome to use it directly, or adapt it as a reference.
Note that although the idea of "multi-head" is simple — repeat something several times and concatenate — you shouldn't actually write your code that way, because it will be very slow. TensorFlow does not automatically parallelize operations, for example:
a = tf.zeros((10, 10))
b = a + 1
c = a + 2
Here, the computations for b and c are executed sequentially even though b and c don't depend on each other. So we must merge the multi-head operations into a single tensor computation, since multiplication within a single tensor is automatically parallelized internally.
Additionally, we need to mask the sequence to ignore the effect of padding. Typically masking involves zeroing out the padded region, but in attention, masking must be applied before the softmax, by subtracting a large number from the padded positions (so that after softmax they become essentially 0). All of this is implemented in the accompanying code.
TensorFlow version
Here is the TensorFlow implementation:
https://github.com/bojone/attention/blob/master/attention_tf.py
Keras version
Keras is still one of my favorite deep learning frameworks, so of course I had to write one for Keras as well:
https://github.com/bojone/attention/blob/master/attention_keras.py
Code test
A quick test on IMDB using Keras (without masking):
from __future__ import print_function
from keras.preprocessing import sequence
from keras.datasets import imdb
from attention_keras import *
max_features = 20000
maxlen = 80
batch_size = 32
print('Loading data...')
(x_train, y_train), (x_test, y_test) = imdb.load_data(num_words=max_features)
print(len(x_train), 'train sequences')
print(len(x_test), 'test sequences')
print('Pad sequences (samples x time)')
x_train = sequence.pad_sequences(x_train, maxlen=maxlen)
x_test = sequence.pad_sequences(x_test, maxlen=maxlen)
print('x_train shape:', x_train.shape)
print('x_test shape:', x_test.shape)
from keras.models import Model
from keras.layers import *
S_inputs = Input(shape=(None,), dtype='int32')
embeddings = Embedding(max_features, 128)(S_inputs)
# embeddings = SinCosPositionEmbedding(128)(embeddings) # 增加Position_Embedding能轻微提高准确率
O_seq = Attention(8,16)([embeddings,embeddings,embeddings])
O_seq = GlobalAveragePooling1D()(O_seq)
O_seq = Dropout(0.5)(O_seq)
outputs = Dense(1, activation='sigmoid')(O_seq)
model = Model(inputs=S_inputs, outputs=outputs)
# try using different optimizers and different optimizer configs
model.compile(loss='binary_crossentropy',
optimizer='adam',
metrics=['accuracy'])
print('Train...')
model.fit(x_train, y_train,
batch_size=batch_size,
epochs=5,
validation_data=(x_test, y_test))
Results without Position Embedding:
Train on 25000 samples, validate on 25000 samples
25000/25000 [==============================] - 9s - loss: 0.4090 - acc: 0.8126 - val_loss: 0.3541 - val_acc: 0.8430
Epoch 2/5
25000/25000 [==============================] - 9s - loss: 0.2528 - acc: 0.8976 - val_loss: 0.3962 - val_acc: 0.8284
Epoch 3/5
25000/25000 [==============================] - 9s - loss: 0.1731 - acc: 0.9335 - val_loss: 0.5172 - val_acc: 0.8137
Epoch 4/5
25000/25000 [==============================] - 9s - loss: 0.1172 - acc: 0.9568 - val_loss: 0.6185 - val_acc: 0.8009
Epoch 5/5
25000/25000 [==============================] - 9s - loss: 0.0750 - acc: 0.9730 - val_loss: 0.9310 - val_acc: 0.7925
Results with Position Embedding:
Train on 25000 samples, validate on 25000 samples
Epoch 1/5
25000/25000 [==============================] - 9s - loss: 0.5179 - acc: 0.7183 - val_loss: 0.3540 - val_acc: 0.8413
Epoch 2/5
25000/25000 [==============================] - 9s - loss: 0.2880 - acc: 0.8786 - val_loss: 0.3464 - val_acc: 0.8447
Epoch 3/5
25000/25000 [==============================] - 9s - loss: 0.1584 - acc: 0.9404 - val_loss: 0.4398 - val_acc: 0.8313
Epoch 4/5
25000/25000 [==============================] - 9s - loss: 0.0588 - acc: 0.9803 - val_loss: 0.5836 - val_acc: 0.8243
Epoch 5/5
25000/25000 [==============================] - 9s - loss: 0.0182 - acc: 0.9947 - val_loss: 0.8095 - val_acc: 0.8178
It seems the peak accuracy is even a bit higher than that of a single-layer LSTM, and we can also see that Position Embedding both improves accuracy and reduces overfitting.
Computational Cost Analysis
As we can see, the computational cost of attention is actually not low. For example, in self-attention, we first need to apply three linear projections to $\boldsymbol{X}$, which already amounts to the cost of a 1D convolution with kernel size 3 — though this part of the cost is still only $\mathcal{O}(n)$. Then there are two matrix multiplications involving the sequence with itself, each of which has a computational cost of $\mathcal{O}(n^2)$ — if the sequence is sufficiently long, this cost becomes quite difficult to accept.
This also suggests that the restricted version of attention is a promising direction for future research, and that combining attention with CNNs and RNNs may well be the more moderate path forward.
Conclusion
Thanks to Google for this brilliant showcase, which has broadened our horizons and deepened our understanding of attention. Google's achievement here, in a sense, embodies the philosophy that "great simplicity comes from great sophistication" — it truly is a rare gem in NLP research. This post, centered around Google's masterwork, is a modest attempt of my own, and I hope it will help interested readers better understand attention. Finally, I sincerely welcome any suggestions and criticism.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.