Attention, explained for engineers

In late 2011 I spent a lot of evenings in Octave doing the programming exercises for Andrew Ng’s online machine learning class, the first run of it, the one that kicked off the whole MOOC thing. The exercise that stuck with me was the one where you implement backpropagation by hand for a small neural network and train it to recognize handwritten digits from 20x20 pixel images. I remember the moment it went from 10% accuracy (random) to about 95% because I’d finally gotten a matrix transposed the right way round. I also remember thinking that neural networks were a neat idea that mostly didn’t work in practice, which is what a lot of people thought in 2011.

Six years later I’m reading a paper where eight researchers, nearly all at Google, train a neural network to translate English to German better than anything before it, and I can follow most of it, which I’m a bit proud of. The paper is Attention Is All You Need, by Ashish Vaswani and seven co-authors from Google Brain and Google Research (plus one University of Toronto student who did the work at Google Brain), posted to arXiv on June 12th. The architecture is called the Transformer. I’ve read it three times now and I think the core idea is simpler than the paper makes it look, so here’s my attempt at explaining it to someone who writes software but hasn’t lived in this field.

The way translation models worked until now

The standard neural translation setup for the last few years has been an encoder and a decoder, both recurrent neural networks (usually LSTMs). The encoder reads the English sentence one word at a time, updating a hidden state vector after each word. The decoder then generates the German one word at a time, using its own hidden state and the encoder’s output.

The original version squeezed the whole input sentence into one fixed-size vector, which went badly for long sentences, for fairly obvious reasons. In 2014 Bahdanau, Cho and Bengio added attention: at each step the decoder gets to look back at all of the encoder’s hidden states and take a weighted average of them, with the weights computed on the fly based on what it’s currently trying to produce. When it’s generating the German word for “cat”, it can put most of its weight on the encoder state for “cat”. That worked really well, and it’s in basically every serious system since, including the Google Translate system Google described last September, which uses deep stacks of LSTMs with attention between them.

The problem with RNNs is the R. To compute the hidden state for word 50 you need the hidden state for word 49, which needs word 48, and so on. That’s inherently sequential. You can’t spread one sentence across a GPU’s thousands of cores, so training is slow, and information from the start of a long sentence has to survive dozens of steps of being squashed through the same function to reach the end. Facebook published a paper last month replacing the RNNs with convolutions for exactly this reason, which gets the parallelism back but still needs several layers before distant words can see each other.

What the Transformer does instead

It throws out the recurrence and the convolutions and keeps only the attention. Every word in the sentence attends to every other word directly, in one step, and all of those computations happen at once.

What attention actually computes

This is the part that I think gets lost in the diagrams of boxes and arrows. Attention is a soft dictionary lookup.

A normal dictionary has keys and values. You give it a query, it finds the key that matches exactly, and it gives you that key’s value. Attention does the same thing, except that queries, keys and values are all vectors, and instead of returning the one value whose key matches, it returns a weighted average of all the values, weighted by how well each key matches the query.

“How well it matches” is the dot product of the query with each key. Those scores go through a softmax, which turns them into positive weights that add up to 1, and then you take that weighted sum of the values. The paper’s whole formula is one line:

Attention(Q, K, V) = softmax(Q·Kᵀ / √d_k) · V

where Q, K and V are matrices with one row per word, and d_k is the length of the key vectors. Here’s a tiny worked example with 2-dimensional vectors and three words, so you can do it on paper:

query            q  = [1, 0]

keys             k1 = [1, 0]   k2 = [0, 1]   k3 = [1, 1]
values           v1 = [10, 0]  v2 = [0, 10]  v3 = [5, 5]

scores q·k / √2  0.707         0             0.707
softmax          0.40          0.20          0.40

output = 0.40·v1 + 0.20·v2 + 0.40·v3 = [6.0, 4.0]

The query matched keys 1 and 3 equally well and key 2 not at all, so the output is mostly a blend of values 1 and 3, with a bit of value 2 because softmax never gives anything a weight of exactly zero. (The exact numbers are 0.401, 0.198, 0.401 and an output of [6.02, 3.98], if you’re checking.)

The / √d_k was the bit I didn’t understand on first reading. The paper says only that without it, dot products grow too large for bigger key sizes and push the softmax into regions with tiny gradients. I had to work out why on paper. If the entries of the query and key are random with mean 0 and variance 1, their dot product is a sum of d_k such products and has variance d_k. With 64-dimensional keys, which is what they use, the raw scores have a standard deviation of 8, and when you push numbers that size through a softmax you get one weight near 1 and the rest near 0, where learning stalls. Dividing by √d_k = 8 puts the variance back to 1. That’s the kind of detail you only find by actually trying to train the thing, and I’m glad they mentioned it.

So where do the queries, keys and values come from? In self-attention, all three come from the same sentence. Each word’s current vector is multiplied by three different learned weight matrices to produce its query, its key and its value. So the word “it” produces a query that, after training, presumably means something like “I’m a pronoun, what am I referring to?”, and the word “cat” produces a key that answers that kind of question well. Nobody programs that in. The weight matrices are learned by backpropagation, exactly like the digit-recognition network from that 2011 class, just with a lot more of them.

Multiple heads, positions, and the rest of the model

One attention lookup per word can only pull out one kind of relationship at a time, so they run eight of them in parallel, each with its own weight matrices, and concatenate the results. They call these heads. One head might end up tracking which noun a pronoun refers to and another which verb goes with which subject, though that’s my speculation, the paper only shows a few visualizations.

The obvious problem with throwing out recurrence is that attention doesn’t care about order. “Dog bites man” and “man bites dog” contain the same words, and a weighted average doesn’t know which came first. Their fix is to add a positional encoding to each word’s vector before it goes in: a fixed pattern of sine and cosine waves at different frequencies, so each position gets a unique signature and relative positions can be worked out from the differences. It’s a slightly odd-looking trick. Their reasoning is that for any fixed offset, the encoding of position p+k is a linear function of the encoding of position p, so the model should find it easy to learn things like “look two words back”. It’s a hypothesis rather than something they prove, and Facebook’s convolutional model just learns a vector per position instead. I’d be curious whether it makes much difference.

Each layer of the model is a self-attention step followed by a small ordinary feed-forward network applied to each word separately, with residual connections and layer normalization around both. The encoder stacks six of these. The decoder also has six, with two changes: its self-attention is masked so a word can’t peek at the words after it (otherwise it could cheat by looking at the answer during training), and it has an extra attention step where the queries come from the decoder and the keys and values come from the encoder’s output, which is the 2014-style attention again. The base model uses 512-dimensional word vectors and 8 heads of 64 dimensions each.

Why it matters, maybe

The results are good. On the standard WMT 2014 English-to-German benchmark the big model scores 28.4 BLEU, which the paper says is over 2 points better than the best previous results, including ensembles of several models. It also sets a new single-model record on English-to-French. BLEU numbers mean very little to me on their own, but 2 points on a benchmark that people have been grinding on for years is a lot.

What I find more interesting is the cost. The big model trained in 3.5 days on eight GPUs. The base model, which already beats most previous systems, trained in about 12 hours on eight. The paper has a table comparing layer types that makes the reason obvious: for self-attention, the number of sequential operations per layer is constant, where for an RNN it grows with the length of the sentence, and the longest path between any two words is one step instead of n. The cost is that every word attends to every other word, so the work grows with the square of the sentence length. For sentences that’s fine. For a whole document it might not be, and the paper says as much and suggests restricting attention to a neighborhood for long inputs.

What interests me as an engineer is that this design is shaped around the hardware. RNNs are elegant on paper and miserable on a GPU, because a GPU wants to do the same thing to a lot of data at once. The Transformer is basically a pile of matrix multiplies, which is exactly what GPUs are for. My guess, and it’s only a guess, is that we’ll see more models designed backwards from what a GPU is good at, rather than from what seems most natural for the problem.

Whether it generalizes beyond translation, I don’t know. They show it doing reasonably on English parsing as well, which is a hint. Google has put the code into a new open-source library called Tensor2Tensor, so people outside Google will be able to try it on other problems soon, and we’ll find out. I’m skeptical of any paper title that says you only need one thing. But I said something similar about neural networks in 2011, and I was wrong about that.*

* In fairness to 2011 me, the digits exercise trained on 5,000 examples on a laptop CPU. The Transformer’s training data is about 4.5 million sentence pairs for English-German alone, on hardware that didn’t exist then.