Self-attention: the sequence looks at itself

The previous lesson left the architecture without a single recurrent piece and with one operation at its centre, attention, which has to be handed three things: queries, keys and values. Nothing is left to produce any of them, because both of the networks that did were recurrent. The only thing that still has a row for every position is the sentence itself, already turned into vectors, so there is one candidate for all three roles and no second one.

Using it as it is does not work, though, and the reason fits in a line: a vector resembles itself more than anything else. Every position would give itself the highest weight, in every row at once, and the layer would hand each one back what it already had, with no parameter to do anything else. And it does need to do something else: in the keys to this car are here, getting are right means having looked at keys, four positions back, and not at car, which sits right next to it and is singular. This lesson builds the layer that can do that, and the price is three learnt matrices, one per role. Multiplying a vector by one of them is projecting it: out comes another vector, usually shorter, whose coordinates are combinations of the original's, and the matrix decides which combinations. Each position is projected three times, giving its query, its key and its value.

The grid below sets a sentence against itself: the row is the position asking, the column is the position answering, and each cell shows how much of that column goes into what that row reads. Type any sentence you like and, above all, switch the projections off and on. With them off, the queries, keys and values are the sentence's vectors as they are, and the highest cell in each row is the one that position gives itself, in all seven rows at once. With them on, the row for are goes looking for keys above everything else.

Three readings of the same matrix

Stack the sentence by rows, the way the previous block stacked the encoder's states:

where row $t$ is $\mathbf{x}_t^{\top}$, the vector of the token at position $t$ (transposed, because in this course vectors are columns). That is everything that goes into the layer: one matrix, and no second network.

The queries, keys and values all come out of it, each through its own matrix:

with $\mathbf{W}^Q, \mathbf{W}^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$ and $\mathbf{W}^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$, so that $\mathbf{Q}, \mathbf{K} \in \mathbb{R}^{T \times d_k}$ and $\mathbf{V} \in \mathbb{R}^{T \times d_v}$. The superscripts on these three matrices are labels for the role each one plays, not layer indices. This is the only place in the course where a superscript does not count layers, and it comes from the paper's notation. There is a single set of these three projections here; a real layer carries several in parallel, which is the subject of a later lesson on multi-head attentionComing up · Block 5 · Lesson 4 · In SpanishVarias cabezas, varios repartosUna posición tiene una fila y en esa fila cabe una sola manera de mirar la frase. La atención multi-head monta h atenciones en paralelo sobre la misma X, cada una con sus tres proyecciones y su propio mapa, y las junta concatenando y proyectando. Lo que cambia no son las puntuaciones sino en qué momento entra el softmax: cada grupo de coordenadas con el suyo, en vez de todas sumadas antes de uno solo, que es lo que deja a una regla tapar a otra sin dejar rastro. Y de repartir la anchura, d_k = d_model/h, sale que el coste no depende de h: los mismos parámetros y las mismas multiplicaciones con una cabeza que con dieciséis..

That is where the self in the name comes from, and it is what separates this lesson from the lesson on attention as query, key and valueBlock 4 · Lesson 6Attention as query, key and valueThe block's mechanism is written in the vocabulary of one particular machine (an encoder that reads, a decoder that writes, states with a bar and states without one), and it needs almost none of that. Here the mechanism is rewritten in full around three roles: a query, a list of keys it is compared against, and a list of values that get mixed. Not one equation changes, and three things come out of the rewrite: the key and the value are the same vector by habit, not by necessity; the gradient reaches each role along one route rather than two; and the whole grid fits in a product of three matrices, with the names the next block uses.. There, the query was the state of one network and the keys and values were the states of another: two sources for three roles. Here all three come out of the same $\mathbf{X}$. The operation that receives them is the same as in that lesson, with one addition: the divisor the paper puts in the score.

with the softmax taken by rows. That divisor, $\sqrt{d_k}$, is the only thing in the line nobody has justified yet. For now, take note of it; I will come back to it at the end.

That $\mathbf{Q}$ comes whole out of $\mathbf{X}$ does not contradict the previous lesson's caveat that, when generating text, the decoder still writes one token per step. Here $\mathbf{X}$ is a sentence that is already complete (the source sentence, or the target one during training), so all the queries come out at once. When generating, the decoder's $\mathbf{X}$ is what it has written so far and gains a row at each step, and $\mathbf{Q}$ gains one with it. But the new query is that row multiplied by $\mathbf{W}^Q$, not the state of a network that needs the previous step's: the only thing forcing the decoder to go one token at a time is that the next token has not been written yet. How that decoder is put together 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..

Now look at the shapes, which is where you can see that the input and the output are the same sentence. $\mathbf{Q}\mathbf{K}^{\top}$ multiplies a $T \times d_k$ matrix by a $d_k \times T$ one, so the product is square, $T \times T$: one row and one column per position. The softmax does not change that shape, so the weight map $\mathbf{A}$ that comes out of it is $T \times T$ as well, with the weight $\alpha_{ij}$ in cell $(i, j)$: how much of position $j$ goes into what position $i$ reads. This is the grid from the start of the lesson, and the literature calls it the attention weights or the attention map. In the previous block it was $T_y \times T_x$ and its two axes were two different sentences; here both axes are the same one, which is why the cell $\alpha_{ii}$ exists and means something: how much of itself a position keeps. Each row still sums to $1$, because the softmax still normalises one row at a time, and no position is left out of any row: every position looks at every position.

And count the parameters, which is the other property of shape that decides what follows in the block. The three matrices hold between them

and $T$ does not appear in that count. They multiply rows of $\mathbf{X}$, and it makes no difference to them how many rows there are. Compare that with 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.: there, each position had its own block of columns $\mathbf{W}^{(1)}_t$, so lengthening the input meant adding parameters, and each position learnt on its own. These three matrices are the same for every position and every length.

Without projections, the layer copies itself

What remains to be defended is what looks like a detour: if $\mathbf{Q}$, $\mathbf{K}$ and $\mathbf{V}$ are readings of $\mathbf{X}$, why not use it as it is and save the three projections? Set $\mathbf{Q} = \mathbf{K} = \mathbf{V} = \mathbf{X}$ and work it through. The score of the pair $(i,j)$ becomes the dot product of the two token vectors,

and two things follow, neither of them good. First, $\mathbf{x}_i^{\top}\mathbf{x}_j = \mathbf{x}_j^{\top}\mathbf{x}_i$, so the grid of scores is symmetric. The weights are not (each row is divided by its own sum, and those sums have no reason to match), but the ordering within each row is fixed by a number that cannot tell the position asking from the position answering. A language can: are needs keys far more than keys needs are.

The second is worse. Suppose all the token vectors have the same length, $\lVert \mathbf{x}_t \rVert = r$. I am assuming this so that the calculation comes out exact, and the Cauchy–Schwarz inequality does the rest:

with equality only when $\mathbf{x}_j$ points in the direction of $\mathbf{x}_i$. The diagonal cell is the largest in its row, always, in every row at once, so the softmax gives the highest weight to the position itself and the mixture leans towards $\mathbf{x}_i$: the layer hands back a faded version of what came in. How far it leans depends on how long the vectors are: the larger $r$ is, the more tightly the row closes on its diagonal. If the vectors have different lengths, the conclusion keeps its shape: the row then goes to whichever position has the longest vector, which is still a property of the vectors and not of what the sentence needs.

What settles the argument is neither of these, but what is missing from both. That layer has not a single parameter. There is nothing to adjust, so there is no way for it to learn that a plural verb has to go and find a plural noun. The three matrices are what separate the three roles (asking is not being asked, and neither is contributing content) and, on top of that, they are the only thing the layer can learn.

The sequence in all three roles, in NumPy

The whole layer is five lines. Read them with the shapes in front of you, which is where you see the sentence go in once and come out as three different readings.

Change T to $9$ and run it again: the map becomes $9 \times 9$, the output gains three rows and the parameter line does not move. That is the layer's shape property in a nutshell.

Now the check that backs up the previous section. The cell below removes the projections, measures what happens, and then puts them back.

The first two parts are the previous section in numbers. Unprojected, all six rows put their maximum on their own diagonal, and the line with the vectors three times longer shows what decides whether that maximum is a preference or a monopoly: at length $1$ the row shares its weight out almost evenly, and stretching the same vectors makes the weight a position gives itself shoot up. With two different readings of the same matrix, that comes to an end and the positions start to look at one another. The third part is something else, and it is worth reading slowly, because it is what the next lesson is about.

Test your intuition

Five questions: where $\mathbf{Q}$, $\mathbf{K}$ and $\mathbf{V}$ come from, how big the map is, what the layer computes without its projections, what is symmetric and what is not, and what changes when the sentence gets longer.

That leaves the divisor, and it cannot be put off any longer. The third part of the last cell measures it without touching anything else: with queries and keys of $d_k = 8$, the largest weight in a row averages $0.558$, and with random vectors drawn the same way at $d_k = 512$ it climbs to $0.983$. A row like that has stopped being a mixture: it keeps one position and throws away the other five. This has had a name since the lesson on activation functionsBlock 2 · Lesson 2Activation functions and non-linearityThe step passes along the sign of the pre-activation and drops everything else, and its derivative is zero wherever it even exists. Sigmoid, hyperbolic tangent and ReLU keep the magnitude instead and have a real derivative; none of them can be a straight line, because chaining affine functions together only ever gives back another affine function.: the row is saturated. Dividing by $\sqrt{d_k}$, the maximum grows too, but slowly, from $0.330$ to $0.486$, and the row goes on sharing its weight out.

That root is not a constant someone kept tweaking until things came out right: it is calculated, and what has to be calculated is how much a dot product grows as coordinates are added to it. That calculation takes up the whole of the next lesson, on the scaled dot-product, and it deserves to be done slowly, because it is where two things this course left in different blocks meet again: the saturation that switches off a share-out and the gradient that vanishes on its way through one. Neither of them needed a sigmoid or a recurrence to appear.