Goodbye to recurrence

The previous block stripped attention down to its bones: three roles (query, key and value), a function that compares, and a weighted sum, with not a single recurrence inside. What it left untouched was the machine attention sits in, and that machine is still recurrent from top to bottom: the keys and values are states of a recurrent encoder, and the queries come from a recurrent decoder. Inside, attention waits for nothing. Outside, every call still has to wait for the decoder to finish its previous step and hand over the next query.

This lesson builds no new piece, and it helps to say up front what it removes and what it keeps. It removes the recurrence, not the encoder or the decoder. The encoder still reads the source sentence and the decoder still writes the target one, but neither keeps a state that is passed from one position to the next. Each layer handles every position at once, and the only thing that puts a token in contact with all the others is attention. Building that version takes the rest of the block; this lesson first counts what you gain and what you lose. That is not a matter of taste, because recurrence charges a price you can count. Take a news story of $800$ tokens and any recurrent network: reading it takes $800$ steps, one after another, and buying a machine with $800$ processors changes nothing, because step $401$ needs a number that step $400$ has not finished computing. That is the first bill. The second took the whole of Block 3 to pay: whatever position $1$ has to tell position $800$ travels through $799$ multiplications by the same matrix, and by the time it arrives there is almost nothing left of it. Both bills have the same source, and removing it clears both.

The diagram below is where all of this ends up: the complete architecture the block builds, which is Figure 1 of the paper Attention Is All You Need, box by box. At this point half the names mean nothing to you, and that is fine, because each box says which lesson builds it. What you can read today is the one claim this lesson makes. Turn on the marker and, of the fifteen boxes, three light up. The other twelve work on one row at a time and have no way of knowing what is in the next one.

The chain of waits no computer can shorten

The recurrence from the previous block fits on one line, and that line says it all:

with $\mathbf{h}_t \in \mathbb{R}^{d_h}$ and $\mathbf{h}_0 = \mathbf{0}$. Of the two terms, the second waits for nobody. $\mathbf{W}_{xh}\mathbf{x}_t$ looks only at the token in its own position, so all $T$ products can be computed at once as soon as the text is tokenised. The first term does wait. $\mathbf{h}_{t-1}$ is the result of the previous step, and there is no way to bring it forward without computing it. One half can be shared out among processors and the other cannot, and the half that cannot is the one that matters: the depth of the chain, meaning how many operations have to happen one after another, is $T$.

Set that against the attention call as the previous block left it. The score of the pair $(i,j)$ is

and no $e$ from any other pair appears in it. The wait the previous block ended on, each row of $\mathbf{Q}$ waiting for the one before, is not in this formula. It lives in the decoder that makes the queries, which needs the context from call $i$ to produce query $i+1$. That is why the call was made once per step, row by row. Take that decoder away and suppose, for now, that $\mathbf{Q}$, $\mathbf{K}$ and $\mathbf{V}$ arrive whole, one row per position of the sentence. Who makes them without a recurrence is the question this lesson ends on. With all three in hand, each of the $T^{2}$ scores comes from two vectors that are already there, so all $T^{2}$ can be computed at once. That is what it means for them to fit in a single matrix product, $\mathbf{Q}\mathbf{K}^{\top}$. The softmax normalises each row using only its own entries, and no row looks at another; the mixture is a second product. That makes three operations in a row, and the three does not change when the text gets longer.

One caveat, so that you do not read more into this than it says. When generating text, the decoder still writes one token per step, because it cannot look at a token it has not written yet. What disappears is the wait when processing a sentence that is already complete: the source sentence always, and the target sentence during training, when the correct translation is already known. The recurrence could not take advantage of that even in training: with the translation right there, the state at step $t$ still needed the one at $t-1$. How the decoder is trained in a single pass without being allowed to see the answer is the subject of the lesson on masksComing up · Block 5 · Lesson 7 · In SpanishEscribir no es leer: el decoder y su máscaraTodo lo que lleva construido el bloque supone la frase entera delante y las T posiciones calculadas a la vez, y hay media arquitectura para la que eso es falso: un modelo que escribe no tiene el resultado, lo va produciendo. Entrenarlo de una sola pasada pone la respuesta al lado de la posición que tiene que adivinarla, y la solución no es un bucle sino un sumando: menos infinito en la puntuación, antes del softmax, para que la exponencial lo borre del numerador y del denominador a la vez. Eso deja cada fila quedándose con su propio principio y volviéndolo a normalizar, así que la primera posición se lleva la unidad entera y ninguna fila deja de sumar 1. Con eso montado, la segunda subcapa del decoder es la atención del bloque anterior sin un solo cambio: consultas de lo que se escribe, claves y valores de lo que se leyó, mapa rectangular y ninguna máscara..

The other distance, the one Block 3 couldn't shorten

The depth of the chain decides how long the computation takes. There is a second distance, and it decides something worse.

Ask yourself which route whatever position $1$ contributes to position $T$ has to travel. In the recurrence it travels inside the state, and the state is rewritten at every step, so that contribution passes through $T-1$ products by $\mathbf{W}_{hh}$ before it arrives. This is exactly the count from the lesson on the vanishing gradientBlock 3 · Lesson 4The vanishing gradientMultiplying by the same recurrent matrix at every step makes the gradient exponential in the distance it travels. What guarantees that it vanishes is γ·σ_max < 1, the steepest tanh slope times the largest singular value of W_hh; the spectral radius only sets the pace of the matrix on its own. Gradient clipping stops the explosion but can't rescue a gradient that has already vanished, and that's what will force a change of architecture.: multiplying $T-1$ times by the same matrix switches the signal off if its spectral radius is below $1$ and blows it up if it is above. The LSTM and the GRU did not shorten that path; they opened an additive path beside it so that the trip could be survived.

In attention, the path between those two positions is a single weight:

What position $1$ contributes is its own term, and in front of it stands a single factor, $\alpha_{T1}$. One, not $T-1$. And since the backward pass is the forward pass travelled in reverse, the gradient that reaches position $1$ from the loss at step $T$ does not go through anything in between either. The distance between two positions stops depending on how many lie between them, which is more than Block 3 managed: the gate made the trip survivable, not short.

What it costs to swap the steps for a grid

None of this is free, and the price is counted in multiplications. Start with the recurrence. A step multiplies $\mathbf{W}_{hh} \in \mathbb{R}^{d_h \times d_h}$ by one vector and $\mathbf{W}_{xh} \in \mathbb{R}^{d_h \times d_{\text{model}}}$ by another. Multiplying a matrix by a vector costs one multiplication per entry, so a step costs one product for each entry of its matrices. That is the count from the lesson on the vanilla RNNBlock 3 · Lesson 2The vanilla RNNOne matrix applied at every position, and a state carried from each step to the next. That's enough to keep the length of the text out of every weight's shape: the same network reads two tokens or forty without padding or truncating anything, what it learns at position 2 works unchanged at position 8, and two sentences with the same tokens in a different order are no longer the same input., read as a cost rather than as memory. Over the whole sentence:

Now the grid. $\mathbf{Q}\mathbf{K}^{\top}$ multiplies a $T \times d_k$ matrix by a $d_k \times T$ one. Each of the $T^{2}$ entries of the result is a dot product of $d_k$ terms, which makes $T^{2} \cdot d_k$ multiplications. The mixture $\mathbf{A}\mathbf{V}$ has $T \cdot d_v$ entries and each one sums $T$ terms, which adds another $T^{2} \cdot d_v$. In total:

Seen pair by pair: scoring entry $(i,j)$ of the grid costs $d_k$ multiplications, and adding its share to the mixture, $\alpha_{ij}\mathbf{v}_j$, costs $d_v$ more. The grid has $T^{2}$ entries.

The two counts compare most easily when all four widths are equal. Suppose they are, and call the common width $d$. That leaves $2Td^{2}$ against $2T^{2}d$, and the ratio between them is $T/d$. The grid costs less than the recurrence as long as the sentence is shorter than the model is wide, and more once it is longer; the two cross at $T = d$, and not before. In the paper these counts are written $O(T^{2}d)$ against $O(Td^{2})$, which is the same comparison with the constants dropped, and that is how you will find it in its Table 1.

What this buys is worth spelling out, because the gain is not only speed. It is a change of axis. The total work grows (with $T$ squared, on long texts), but in exchange that work can be shared out, whereas the recurrence's $T$ steps cannot be shared out at all. A large computation that fits in a single matrix product finishes before a small one broken into $T$ waits.

The price that isn't paid in multiplications

Look at the score again, and this time look closely at the indices. In $e_{ij} = a(\mathbf{q}_i, \mathbf{k}_j)$ the subscripts only point at which two vectors go in; neither $i$ nor $j$ appears inside the computation. The same goes for the mixture, which is a sum, and a sum does not care about the order of its terms. The consequence is awkward: if you shuffle the rows of $\mathbf{K}$ and $\mathbf{V}$, the outputs are shuffled the same way, and not one number changes.

So the dog bites the postman and the postman bites the dog reach the layer as the same set of vectors. The recurrence did not have this problem, though not through any merit of its own: it was solved from birth, because it read in order and the reading order was the sentence's order. Nor did the concatenated MLP from the lesson on why the MLP fails on sequencesBlock 3 · Lesson 1Why the MLP fails on sequencesLining up the tokens' vectors builds order into the input, but at a price: a fixed length that forces you to truncate or pad, and a first layer carved into one block per position, each one trained only on what lands in its own position. A network that learns "not" in the first four positions scores 0.500 when "not" shows up in the fifth., for the opposite reason: each position multiplied its own block of columns, which is why it could tell positions apart and could not generalise across them. Attention has neither of those protections.

The concession is a big one, and this lesson does not fix it: removing the recurrence takes away the only thing that knew the order. The fix is to add to each row something that depends on where it sits, and that is the subject of a later lesson on positional encodingComing up · Block 5 · Lesson 5 · In SpanishCada posición, un puñado de relojesLa capa de atención está terminada y sigue leyendo la frase como un montón de vectores sin orden, así que lo que hay que arreglar es la entrada: cada fila tiene que llegar sabiendo dónde estaba. Meter el número de la posición rompe la escala, dividirlo entre la longitud de la frase hace que la misma coordenada signifique cosas distintas en textos distintos, y aprender un vector por posición deja la 7 y la 8 sin relación ninguna. Los senos y cosenos del artículo resuelven las tres: acotados por construcción, con una longitud de onda por pareja de coordenadas, y con una propiedad que se demuestra en tres líneas —moverse k posiciones es girar un ángulo fijo dentro de cada pareja, con la misma matriz en toda la frase, y el producto escalar entre dos posiciones sólo depende de la distancia que las separa.. Until then, everything that follows treats the input as a bag of vectors.

Check it by hand on a sentence of 50 tokens

Do both counts with numbers to see that the change is not only theoretical. Take $d = 512$ (the width of the paper's model) and a sentence of $T = 50$ tokens.

The recurrence costs $2 \cdot 50 \cdot 512^{2} = 26\,214\,400$ multiplications, spread over $50$ steps that come one after another. The grid costs $2 \cdot 50^{2} \cdot 512 = 2\,560\,000$, and not in $50$ waits but in two matrix products. That is ten times less work, and the depth of the chain has gone from $50$ to a number that does not even mention $50$.

Now go up to $T = 1\,000$, a long text. The grid needs $2 \cdot 1\,000^{2} \cdot 512 = 1\,024\,000\,000$ multiplications and the recurrence $2 \cdot 1\,000 \cdot 512^{2} = 524\,288\,000$. The grid now costs almost double, so this time the newcomer is the expensive one. The ratio is $1\,000/512 \approx 1.95$, exactly as $T/d$ predicted. Try $T = 2\,000$ yourself and check that you get $3.9$. Notice also that the depth of the chain still has not moved, while the recurrence's has gone to $2\,000$.

That is the deal, with both columns in view: you pay in work, which can be shared out between processors, and you buy back waiting, which cannot be shared out at all.

Test your intuition

Four questions: what cannot overlap, where the two counts cross, what happens if you shuffle the rows, and what becomes of the path between two distant positions.

That leaves a gap in the most visible place. The attention call needs queries, keys and values, and throughout the previous block they were made by two recurrent networks: the queries were the decoder's states, one per step, and the keys and values were the encoder's. Those states have gone with the recurrence, and without them the operation this lesson has put at the centre of the architecture has nothing to start from.

The answer is the one that gives the next lesson, on self-attention, its name: let the sequence ask itself. If the queries, the keys and the values are three readings of the same vectors, one per position, no second network is needed to produce anything, and every position gets to look at all the others, itself included. It is the same attention as in the previous block, with all three roles in the hands of a single sequence, and it is what turns the diagram above into something that can be computed.