The Transformer: From Recurrence Failure to Attention

In twenty seventeen, a team at Google published a paper called Attention Is All You Need. The architecture inside it, the Transformer, now sits underneath essentially every large language model you have used. Today we are going to build that architecture out of the problem it was invented to solve. So here is the problem. A sequence goes in and a sequence comes out: an English sentence, and its German translation. For years the standard machinery for that was a recurrent network. It reads the input one position at a time, and it drags a hidden state along with it, updating that state at every word. Written down, the whole idea is one line. The hidden state at step t is some function of the hidden state at step t minus one, and the input at position t. Feed in a word, update the memory, move on. Now look at what that line commits you to. To compute the state at position five you need position four, and for that, three, and two, and one. The positions are chained together. You cannot compute them at the same time however many processors you own, and that is what makes these models slow to train. There is a second cost, and this one is about learning rather than speed. Suppose the first word of a sentence carries something you only need at the very last word. Watch what that information has to do to get there. It goes through every state in between, one hop at a time, and it has to survive all of them. Six words here, five hops. In a forty word sentence, thirty nine hops. Every one of those is another chance for the signal to fade away before it arrives. Convolutional models fix half of it. They do compute the positions in parallel, but the number of operations relating two positions still grows with the distance between them: linearly for one family, logarithmically for another. What we want is a layer where any two positions are a single step apart, and where every position is computed at the same time. And here is what does that. This is scaled dot product attention, and this is the whole architecture built out of it. You are not meant to understand either of them yet. I am showing them now because everything that follows is a walk toward this diagram, and by the end you will know where every box in it comes from.

Here is a sentence. The river bank flooded. Look at the third word, bank. On its own that word is ambiguous. It could be the side of a river, or it could be a building that keeps your money. Nothing in the letters settles it, and the word that does settle it is river, sitting one position to the left. So whatever we build has to let the representation of bank reach out, find river, and pull information from it. The paper borrows its vocabulary for that from databases. Every position produces three vectors. A query, saying what this position is looking for. A key, advertising what this position holds. And a value, which is what this position hands over if it gets chosen. Now, how does a query decide which keys it likes? By the plainest measure of similarity there is, the dot product. Here is a query in yellow, and here are two keys, one belonging to river and one belonging to the. The score for a key is the query dotted with that key, and nothing else. Watch that score as I turn the query. This yellow segment is the shadow the query casts on the river key, and its length is the dot product. Point the query along that key and the shadow is as long as it gets. Turn the query away and the shadow shrinks. Turn it until the two are nearly at right angles, and the score is almost nothing. So each key gets a number, and a large number means this key matches what the query is hunting for. That is the whole of the first step. Back to our sentence, with all four words in it. The query belongs to bank, and it scores every key in the sentence, including its own. Four words, four scores. Against river the score comes out high, because that is the key this query was built to match. But scores are not weights. They can be any size at all, and some of them will be negative, and we are about to take an average with them. So push them through a softmax. Exponentiate each score, divide by the sum of all of them, and out come four numbers that are positive and add up to one. There they are, drawn over the words they belong to. River takes the largest share, flooded takes a decent one, and the takes almost nothing. The output of the layer for bank is the weighted sum of the four value vectors, using exactly those weights. Mostly river, a little flooded, barely any of the. That is scaled dot product attention with the scaling left out, and the scaling is the one thing left to explain. A dot product in d k dimensions is a sum of d k separate products. If you assume the components are independent, with mean zero and unit variance, then the sum has mean zero and variance d k. So its typical magnitude grows like the square root of the dimension, and in the paper the dimension per head is sixty four. Watch what that does to the softmax. I am going to grow the scores while leaving their pattern alone. The largest one runs away with everything, and the others are crushed flat against the floor. That is the failure the paper is worried about. A softmax pinned that hard against one and zero has almost no gradient left, so the layer stops learning. The fix is one division. Divide every score by the square root of d k, and the scores come back to a size that does not depend on how wide the vectors happen to be. Now write it once for every position at the same time. Stack all the queries as the rows of a matrix Q, all the keys as the rows of K, all the values as the rows of V. Then Q times K transpose is every score against every key in a single matrix multiply. Q and K have n rows of width d k, so that product is n by n: one score for every ordered pair of positions in the sentence. Take a softmax along each row, multiply by V, and you have the output for all n positions. There is no loop anywhere in that expression, which is exactly the point. And here is the paper's own drawing of it. Read the left column of boxes upward. Q and K go into a matrix multiply, the result is scaled by that square root, there is an optional mask we will need later, then a softmax, then a second matrix multiply against V. Six boxes, and you have just derived every one of them. Compare that with the chain we started on. Any position can read any other position with one multiply and one softmax, whether they are neighbours or forty words apart. Constant path length, and every position computed at once. That was the whole wish list.

Let us look hard at what we just built. One query, one set of scores, one softmax, one weighted average. That average is a single answer to a single question, and a word in a sentence is usually in more than one relationship at a time. Take bank again. One thing you might want to know is what describes it, and the answer to that is river. The weights for that question look like this, with almost everything on one word. But you might equally want to know what happened to it, and the answer to that is flooded, which is a completely different word. Those weights look like this instead, with almost everything on flooded. A single head cannot deliver both. It has one distribution to spend, so the best it can do is put weight on both words and hand back a blend of the two values. The paper's phrase for that is that averaging inhibits it, and blending two answers is often worse than either answer alone. The fix is not to make attention cleverer. It is to run several copies of it side by side and let each copy specialise. Each copy is called a head, and each head gets its own three learned matrices. Head number i takes the queries, the keys and the values, and multiplies each of them by a projection matrix of its own. There is one matrix for the queries, a second one for the keys, and a third for the values. All three are narrow: they squeeze the full model width down to d k, and each head then works inside its own small subspace. Every head returns a short output vector, one per position. Lay all the heads' outputs side by side, which is what concatenation means here, and multiply the result by one final matrix W O. That is multi head attention, and it is the whole definition. So the model is free to give one head the describes it relationship and another the happened to it relationship, and W O decides how to combine what they found. Nothing is averaged away before the model has had a chance to use it. Now, is that eight times the work? It is not, and the reason is the squeezing. The paper uses eight heads with a model width of five hundred and twelve, and it sets the width of each head to the model width divided by the number of heads. Five hundred and twelve over eight is sixty four. The dominant cost inside one head is the score matrix, which is n by n entries, each of them a dot product in d k dimensions. So one head costs something of order n squared times d k. Multiply that by the eight heads and the d k cancels against the eight. You are back to n squared times the full model width, which is what a single head at full width would have cost. Eight heads, eight separate subspaces, and essentially the same arithmetic bill at the end of it. Here is the right half of the paper's figure, which is precisely that picture. Along the bottom, V, K and Q each go through their own linear projection. The middle block is the scaled dot product attention we built, drawn once but stacked h deep. At the top, concatenate and pass through one more linear layer. And the ablation in the paper says this matters. Cut it down to a single head, holding the total computation fixed, and translation quality drops by nine tenths of a BLEU point. Push it up to thirty two heads and the quality falls again, because each head has become too narrow to say anything useful.

There is something badly wrong with what we have built, and it is easiest to see by breaking it. Here is our sentence in the order it was written. And here it is shuffled into nonsense. Now ask what the attention layer would compute for each of them. Every score in that layer is a dot product between two vectors. Shuffling the sentence does not change any vector, so it does not change any dot product. It gives you the same scores in a different order, and the same outputs in a different order. Which means the layer cannot tell these two apart. A recurrent network never had this problem, because reading one word at a time is itself a statement about order. We threw the reading order out along with the recurrence, so now we have to put it back in by hand. The trick is disarmingly simple. Give every position its own vector, built from its position number alone, and add that vector to the word embedding before anything else happens. It has to be the same width as the embedding, five hundred and twelve, so that the two can be added at all. And here is how that vector is built. Give each pair of coordinates a frequency of its own, running from one all the way down to one over ten thousand. Then the even coordinate of the vector is the sine of that frequency times the position, and the odd coordinate is the cosine. Plot three of those coordinates against position and you can see what they do. The blue one has a wavelength of about six positions, so it cycles fast. The green one takes about twenty positions to come round. The magenta one takes over sixty and has barely started. Now read one position off all three at once. The dashed line marks the position, and the three yellow dots are its three coordinates. Slide the position along, and watch the fast coordinate race while the slow one barely moves. The wavelengths form a geometric progression, from two pi at one end to ten thousand times two pi at the other. The fast coordinates pin down exactly where you are locally, and the slow ones say roughly where you are in the sentence as a whole. Together, they are a fingerprint for the position. Sinusoids buy one more thing, and it is the reason they were chosen. Move eight positions along from here, and then eight more. The encoding turns through the same angle on every coordinate both times, so a shift of eight is one fixed linear map, wherever in the sentence you apply it. And they never run out. Suppose the longest sentence in training was forty tokens, this red line here. A table of learned position vectors simply has no entry past that line. The sinusoids carry straight on through it, because they are formulas rather than a lookup table. The paper is honest about the trade. It tried learned position embeddings as well and got nearly identical translation scores. It kept the sinusoids because of that extrapolation, on the hunch that the model might then handle sequences longer than anything it was trained on.

So here is that architecture again, and this time it is going to make sense. Read it as two towers. The left tower is the encoder: the source sentence goes in at the bottom and comes out as a stack of vectors, one for every input position. The right tower is the decoder, and it emits the translation one token at a time. Take one encoder layer, the grey box on the left, and there are exactly two things inside it. First a multi head self attention sub layer, where the queries, the keys and the values all come from the same place, the output of the layer below. Second, a feed forward network. That is the entire layer, and it is repeated six times. Each of those two is wrapped the same way, and the wrapping is what makes six of them trainable. Whatever the sub layer computes gets added back onto its own input, which is the residual connection, and the sum is then layer normalised. The gradient always has that addition as a shortcut straight back down the stack. And the feed forward network is smaller than it looks. Two linear maps with a rectified linear unit in between, five hundred and twelve up to two thousand and forty eight and back down again. The crucial word is position wise. It runs on each position separately, with no mixing, so all the talking between positions happens in the attention sub layer. One more thing about that width. Every sub layer, and every embedding layer too, produces vectors of the same five hundred and twelve dimensions. That is not tidiness, it is a requirement: the residual connection adds the sub layer's output to its input, and you can only add two vectors of the same length. Now the decoder, which is the same layer with a third piece bolted into the middle of it. Its own self attention comes first, and that one has a problem the encoder never had. At training time the whole target sentence is present at once, so nothing stops position two from simply reading position four and copying the answer. So it is forbidden. Rows here are query positions and columns are key positions. Position one may look only at itself. Position two may look at itself and the one before. Position three sees everything up to and including itself. Every entry above the diagonal is set to minus infinity, and that comes out of a softmax as a weight of exactly zero. Put that together with the output being shifted right by one position, and the prediction for position i can depend only on outputs the model has already committed to. The model stays auto regressive, and it can still be trained on every position of the sentence in parallel. The middle sub layer of the decoder is the third use of attention, and it is the one that does the translating. Its queries come from the decoder, from the sentence being written. Its keys and its values come from the top of the encoder. So every position being generated can look at every position of the source sentence, and choose which of them it needs. And that is the whole model. At the top of the decoder, one linear layer and a softmax turn each output vector into a probability over the vocabulary. The paper ties that matrix to the two embedding matrices, one set of weights doing three jobs, and multiplies the embeddings by the square root of the model width so that the two contributions to the input arrive at a comparable size. Six encoder layers, six decoder layers, three uses of one mechanism, and a feed forward network at every position. Every box on that diagram is now something you have built.

So, is it actually better, and why. The paper answers that with one table, comparing three kinds of layer on three measures. Here are the three layers: self attention, a recurrent layer, and a convolutional one. Start with the last column, because that is what this whole lecture has been about. It is the longest distance a signal has to travel between any two positions in the sentence. For a recurrent layer that distance is n, the whole length of the sentence. For convolutions it is logarithmic in n, better, but still growing with distance. And for self attention it is one, and it does not grow at all, however long the sentence happens to get. The middle column is the one about hardware. A recurrent layer needs n steps that have to happen in order. Both of the others need a constant number, which is to say the whole sentence is computed in one go. That is the difference between training in half a day and training for a week. And there is an honest cost in the first column. Self attention scores every pair of positions, so it is quadratic in the sentence length. It comes out cheaper than recurrence only when the sentence is shorter than the model is wide, and for sentences chopped into word pieces it usually is. Now the part that decided the argument. This is English to German and English to French translation, on the standard twenty fourteen test sets. The top two rows are the best published systems at the time, and both of them are ensembles, several models voting together. Now read the last two rows. The base model already beats everything previously published, at twenty seven point three into German. And the big Transformer, still a single model, scores twenty eight point four, more than two BLEU above every one of those ensembles. Into French it reaches forty one point eight, a new single model record. Then read the last column, and this is the part that changed how people work. The base model cost about three times ten to the eighteen floating point operations. The convolutional ensemble above it cost around a thousand times more. Twelve hours on eight GPUs against weeks of compute, for a better score. There is one more thing worth seeing, because it is the sort of evidence you rarely get. The attention weights are just numbers, so you can draw them. This is layer five of six, and every line leaving the word making is one head deciding what that word should look at. The colours are different heads. And look where several of them land. Not on the neighbours, but far to the right, on more and on difficult, completing the phrase making something more difficult. Nobody told the model about that phrase. It is a long range dependency, found by a mechanism whose whole design was to put long range dependencies one step away. So here is the diagram one last time, and I promised you would know where every box in it comes from. Attention, so that any two positions are one step apart. Several heads, so that a word can be in several relationships at once. Sinusoids, to give back the order that attention throws away. And residuals with layer normalisation, so that six of these will train. Recurrence was the thing everybody assumed a sequence model had to have. It turned out to be the thing holding sequence models back. That is the claim in the title, and it is worth taking literally.