11. Attention Mechanisms

A convolution decides in advance which inputs may influence which outputs: neighbours, always, with the same weights. Attention makes the opposite choice. It computes the connection strengths from the data, at run time, for every pair of positions.

That is the whole mechanism, and it is worth holding onto as a one-line definition: attention is a dense layer whose weights are a function of the input. Everything else in this chapter — the \(\sqrt{d_k}\), the masks, the multiple heads, the cache that dominates your inference bill — falls out of that sentence and the cost of taking it literally.

Historically it arrived as a patch. Sequence-to-sequence translation compressed the whole source sentence into a single vector, and long sentences got worse the longer they got; Bahdanau et al. let the decoder look back at every source position instead¹. Seven years later the patch had eaten the architecture: "Attention Is All You Need" removed the recurrence and kept only the looking back².

What came before: recurrence, and how far the gradient reaches

Before attention, a sequence was read one step at a time. An RNN carries a state \(h_t\) and rewrites it at every token:

Think of the state as a fixed-size notepad you rewrite after every word. To remember the first word at the end of a paragraph, its note has to survive every rewrite along the way.

The LSTM¹⁴ adds to that a cell \(c_t\) and three gates the network itself learns to open and close — forget, input and output:

The cell is the whole trick. A plain RNN multiplies the state by a matrix at every step; the LSTM multiplies by \(f_t\), a number between 0 and 1 the network picks. With \(f_t\) near 1, information crosses many steps without being rewritten — the same idea as the residual shortcut of chapter 9, nineteen years earlier.

But near 1 is not 1, and the gradient travelling from the last token to the first is still a product of \(T\) factors, and a product of many numbers below 1 shrinks fast: \(0.9^{64} \approx 10^{-3}\). The panel measures all three paths on one scale:

At 64 tokens with weights at scale 1, the gradient reaches the first token at \(5 \times 10^{-4}\) of what left the last one in the plain RNN, and at \(4 \times 10^{-2}\) along the LSTM's cell path. Go to 256 tokens and those become \(10^{-14}\) and \(2 \times 10^{-6}\). Both curves are straight lines on a logarithmic plot, the signature of exponential decay: the LSTM did not solve the problem, it bought a better constant.

Attention has no such product. The output reads every position at once, so the path to any token is one step, and the gradient arrives divided by an attention weight rather than by distance. Add to that recurrence's second cost, which the plot does not show: the \(T\) steps are sequential, and a whole GPU waits on one state at a time. It was that combination — reach and parallelism — that retired recurrence, not an isolated gain in accuracy.

The idea: every word re-reads the others

Compare two sentences:

• "I sat on the bank of the river."
• "I went to the bank to withdraw cash."

The model's first layer (the embedding, just below) hands "bank" the same vector in both: it looks only at the token, not at its neighbours. Attention is what resolves the ambiguity. It lets every word build a new vector by mixing the vectors of the others, with more weight on the ones that matter: "bank" pulls from "sat" and "river" in the first sentence and from "withdraw" and "cash" in the second. That is what an attention layer outputs — one vector per token, now with context.

How do you decide how much each neighbour weighs? The answer comes from a structure you already know, the dictionary:

stock = {"apple": 3, "pear": 7}
stock["pear"]       # the key matches exactly one entry and returns its value: 7

Attention is a soft dictionary. Instead of demanding an exact match, it compares the question against every key, gives each one a similarity score, and returns a weighted average of all the values. Because the average is a smooth function of the scores, it can be differentiated — and the network learns, by gradient, what to ask and what to advertise.

From tokens to vectors: the embedding layer

Before there is anything to compare, the text has to become vectors. It reaches the model as integer indices, and the first layer is a table \(E \in \mathbb{R}^{V \times d}\) with one row per vocabulary token. The forward pass is a lookup:

Stacking the sentence's \(n\) vectors \(y_t\) gives the matrix \(X \in \mathbb{R}^{n \times d}\) that attention receives. Both forms give the same vector, but the second explains the backward pass. Multiplying by a vector that is zero everywhere except at one position produces a gradient that is likewise zero everywhere except in that row, with repetitions summed:

Three consequences:

• The update is sparse. In a step over a few thousand tokens, most rows of \(E\) receive exactly zero gradient. That is why frameworks keep a separate path for embeddings (sparse=True, IndexSelect) rather than treating \(E\) as just another dense matrix.
• A rare token learns slowly. The panel draws tokens from a Zipf law, as real language does: the first few appear several times per sequence, the long tail appears almost never, and its rows barely move.
• The table is enormous, and usually reused. At 128k rows and 4096 columns that is 525 million parameters in the input alone. Tying that matrix to the output layer, which projects back to the vocabulary, saves half of it — and it works because both speak about the same space.

Query, Key, Value

The soft dictionary needs three things from each token, and attention names each one:

• Query — what this position is looking for. It is the question.
• Key — what each position advertises about itself. It is the label compared with the question.
• Value — what each position actually delivers if selected. It is the content.

Score every key against the query with a dot product, turn the scores into a distribution with the softmax, and return the corresponding mixture of values:

Q, K and V are three different linear projections of the same input, \(Q = XW_Q\), \(K = XW_K\), \(V = XW_V\). The matrices \(W_Q\), \(W_K\) and \(W_V\) are the layer's only learned parameters. The three projections are why a token can ask for one thing, advertise another, and deliver a third — a pronoun can query for "the noun I refer to" without itself being a noun.

flowchart LR
  X["X<br/>(n × d)"] --> Q["Q = X·W_Q"]
  X --> K["K = X·W_K"]
  X --> V["V = X·W_V"]
  Q --> S["S = Q·Kᵀ / √d_k<br/>(n × n)"]
  K --> S
  S --> A["A = softmax(S)<br/>each row sums to 1"]
  A --> O["O = A·V<br/>(n × d_v)"]
  V --> O

A worked example, by hand

Take the first sentence and three of its words: "sat · bank · river" (the others are left out only so the arithmetic fits). We will compute the output for "bank". The vectors have two dimensions and were picked by hand; the values have a readable meaning, \([\text{river},\ \text{money}]\):

The query of "bank" is \(q = [1,\ 2]\) and \(d_k = 2\).

1. Scores. The dot product of \(q\) with each key: \(q \cdot k_\text{sat} = 1\), \(q \cdot k_\text{bank} = 1 + 2 = 3\), \(q \cdot k_\text{river} = 4\).
2. Scale. Dividing by \(\sqrt{2} \approx 1.41\): \([0.71,\ 2.12,\ 2.83]\).
3. Softmax. \(e^{0.71} \approx 2.0\), \(e^{2.12} \approx 8.3\), \(e^{2.83} \approx 16.9\); they sum to \(27.3\). Dividing each by the sum, the weights are \(a = [0.07,\ 0.31,\ 0.62]\). They are positive and sum to 1.
4. Mix. The output is the average of the values with those weights:

The "bank" that came in on the fence, \([0.5,\ 0.5]\), leaves the layer as \([0.85,\ 0.15]\): the river bank. No weight was fixed by position. The choice came from how much the query of "bank" resembles each neighbour's key.

The panel starts at exactly these numbers. Redo the sum following steps ①–④, then try:

• change the value of "river" to \([0,\ 1]\) — as if the sentence said "…bank… cash" — and watch the verdict flip without a single attention weight moving;
• move the query to \([1,\ 0]\) and watch the weight leave "river" and split between "sat" and "bank";
• set the query to zero: every score becomes 0 and attention becomes a plain average.

The shape of every matrix

For a sentence of \(n\) tokens, each with \(d\) dimensions, this is how the shapes chain together. Half the bugs in a new implementation show up as a shape error in this table:

The same thing, in code

The whole function fits in three lines. The script redoes the sum above, then runs the same function for all three queries at once — the middle row of \(A\) and \(O\) is the "bank" you computed by hand (the script's comments use the Portuguese example, sentei · banco · praça, with the same numbers):

"""O exemplo feito à mão na página, refeito em código — e depois a mesma conta para todas as linhas.

Contexto: "sentei · banco · praça". A query é a de "banco"; os values têm duas dimensões,
[assento, instituição financeira]. O "banco" sozinho é ambíguo ([0.5, 0.5]); o contexto decide.
"""
import numpy as np

def softmax(s):
    e = np.exp(s - s.max(axis=-1, keepdims=True))     # subtrair o máximo não muda o resultado
    return e / e.sum(axis=-1, keepdims=True)

def attention(Q, K, V):
    S = Q @ K.T / np.sqrt(K.shape[-1])                 # (n, n): score de cada par
    A = softmax(S)                                     # (n, n): cada linha soma 1
    return A @ V, A                                    # (n, d_v): mistura de values

K = np.array([[1.0, 0.0],      # sentei
              [1.0, 1.0],      # banco
              [0.0, 2.0]])     # praça
V = np.array([[1.0, 0.0],      # sentei -> assento
              [0.5, 0.5],      # banco  -> ambíguo
              [1.0, 0.0]])     # praça  -> assento
q = np.array([1.0, 2.0])       # a query de "banco"

np.set_printoptions(precision=2, suppress=True)
s = K @ q
print("q·k         =", s)
print("q·k / √2    =", s / np.sqrt(2))
print("a = softmax =", softmax(s / np.sqrt(2)))
print("Σ a·v       =", softmax(s / np.sqrt(2)) @ V)

# A mesma função, para as três queries ao mesmo tempo: uma linha de A por token.
Q = np.array([[2.0, 0.0],      # sentei
              [1.0, 2.0],      # banco (a mesma query de cima)
              [0.0, 1.0]])     # praça
O, A = attention(Q, K, V)
print("\nA =\n", A)
print("O =\n", O)

q·k         = [1. 3. 4.]
q·k / √2    = [0.71 2.12 2.83]
a = softmax = [0.07 0.31 0.62]
Σ a·v       = [0.85 0.15]

A =
 [[0.45 0.45 0.11]
 [0.07 0.31 0.62]
 [0.14 0.28 0.58]]
O =
 [[0.78 0.22]
 [0.85 0.15]
 [0.86 0.14]]

The whole matrix, and why \(\sqrt{d_k}\) is not cosmetic

One query against three keys is the mechanism. What a layer actually computes is the full \(n \times n\) matrix: every position queries every position. The panel below builds it for a real sentence. Read it by row: a token's row shows where it looks, and every row sums to 1.

Each option under head is a different score function, and all three describe patterns that genuinely occur in trained models — a head that matches by content, a head that just looks one token back, a head that decays with distance. A layer runs many of them at once.

Now turn divide by \(\sqrt{d_k}\) off and drag \(d_k\) up. The matrix goes binary.

The reason is a two-line argument. If \(q\) and \(k\) have independent components with mean 0 and variance 1, then

because it is a sum of \(d_k\) independent terms, each with variance 1. The scores therefore have standard deviation \(\sqrt{d_k}\) — at \(d_k = 128\), about 11, and scores routinely differ by 20. A softmax over scores that differ by 20 is a hard argmax: \(e^{-20} \approx 2 \times 10^{-9}\), so one weight is 1, the rest are \(10^{-9}\), and the gradient through the softmax is zero. Dividing by \(\sqrt{d_k}\) restores unit variance and keeps the layer in the regime where it can still learn.

The backward pass through one head

This subsection is the most technical in the chapter. On a first read, the bold conclusion of the third paragraph is enough.

The sentence above — the gradient through the softmax is zero — falls out of the arithmetic, and writing it costs four lines. With \(S = QK^\top/\sqrt{d_k}\), \(A = \text{softmax}(S)\) and \(O = AV\):

The middle term is the softmax Jacobian, and it is where the lesson sits. Look at the \(A \odot\) factor: if a row of \(A\) is nearly one-hot — one weight at 1 and the rest at \(10^{-9}\) — then almost all of the bracket is multiplied by \(10^{-9}\); and at the one position where \(A\) is not tiny, the bracket itself is close to zero, because the weighted mean it subtracts is that very element. A saturated attention hands back a numerically zero gradient across the whole row. That is what \(\sqrt{d_k}\) protects against.

The same formula gives a cheap invariant for checking an implementation: every row of \(\partial L/\partial S\) sums to exactly zero, because adding a constant to a row of scores does not change its softmax.

"""A passagem reversa de uma cabeça de atenção, conferida por diferenças finitas.

O termo que importa é o do softmax: ele não é um fator por elemento, é uma jacobiana que
subtrai a média ponderada da linha — e é ela que zera o gradiente quando a atenção satura.
"""
import numpy as np

def softmax(s):
    e = np.exp(s - s.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

def forward(Q, K, V):
    d = Q.shape[-1]
    S = Q @ K.T / np.sqrt(d)
    A = softmax(S)
    return A @ V, A

def backward(dO, Q, K, V, A):
    d = Q.shape[-1]
    dV = A.T @ dO
    dA = dO @ V.T
    dS = A * (dA - (dA * A).sum(axis=-1, keepdims=True))     # jacobiana do softmax, linha a linha
    return dS @ K / np.sqrt(d), dS.T @ Q / np.sqrt(d), dV, dS

rng = np.random.default_rng(0)
Q, K, V = (rng.normal(size=(3, 4)) for _ in range(3))
dO = rng.normal(size=(3, 4))

O, A = forward(Q, K, V)
dQ, dK, dV, dS = backward(dO, Q, K, V, A)

def num_grad(M):
    g, h = np.zeros_like(M), 1e-6
    for i in range(M.shape[0]):
        for j in range(M.shape[1]):
            old = M[i, j]
            M[i, j] = old + h; up = (dO * forward(Q, K, V)[0]).sum()
            M[i, j] = old - h; dn = (dO * forward(Q, K, V)[0]).sum()
            M[i, j] = old; g[i, j] = (up - dn) / (2 * h)
    return g

np.set_printoptions(precision=3, suppress=True)
print("pesos de atenção A (cada linha soma 1):\n", A)
print("dQ:\n", dQ)
print("erro máximo contra diferenças finitas — dQ:", f"{np.abs(num_grad(Q) - dQ).max():.1e}",
      " dK:", f"{np.abs(num_grad(K) - dK).max():.1e}",
      " dV:", f"{np.abs(num_grad(V) - dV).max():.1e}")
print("soma de cada linha de dS:", np.round(dS.sum(axis=-1), 12),
      "— sempre zero: somar a mesma constante a uma linha de scores não muda o softmax")

pesos de atenção A (cada linha soma 1):
 [[0.226 0.475 0.298]
 [0.144 0.598 0.259]
 [0.668 0.235 0.097]]
dQ:
 [[0.456 0.354 0.099 0.209]
 [0.407 0.404 0.013 0.129]
 [0.345 0.228 0.109 0.184]]
erro máximo contra diferenças finitas — dQ: 3.2e-10  dK: 6.1e-10  dV: 3.4e-10
soma de cada linha de dS: [ 0. -0. -0.] — sempre zero: somar a mesma constante a uma linha de scores não muda o softmax

Masks: causal, padding, and everything else

Sometimes a token must not look at another: the future, in a model that generates text; the <pad> filler, in a batch. The mask is added to the scores before the softmax, as \(-\infty\) (in practice, a large negative number), and since \(e^{-\infty} = 0\) the forbidden weights become exactly zero:

Turn on causal mask in the panel and the upper triangle goes dark. This one line is the entire difference between an encoder and a decoder, and the reason a decoder can be trained on all positions in parallel: with the mask, position \(i\)'s prediction cannot see its own answer, so one forward pass gives you \(n\) training examples instead of one. For the sentence "the cat climbed the roof":

Four examples, a single pass through the network. Without the mask, position 1 would see "cat" right next to it and learn to copy instead of predict.

Where position comes from

Look at the worked example again: at no point did it use the order of the words. Attention is permutation-equivariant: shuffle the tokens and the outputs shuffle with them, with identical values. On its own it cannot tell "dog bites man" from "man bites dog". Position has to be injected, and how it is injected turned out to matter a great deal for long context.

RoPE is worth understanding rather than memorizing. Picture a clock: each token's hand is turned a little further for every position. The angle between two hands does not depend on what time it is, only on how many positions separate them. In symbols: take Q and K in pairs of dimensions, treat each pair as a complex number, and multiply by \(e^{i\,m\theta}\) where \(m\) is the position. Because rotations compose by adding angles, the dot product between position \(m\) and position \(n\) depends only on \(m - n\):

Absolute rotations in, relative position out — with no extra parameters and no extra tensor to add. Each pair of dimensions turns at a different speed \(\theta\), like the hour, minute and second hands: the fast ones tell close neighbours apart, the slow ones tell long distances apart. That property is also why context windows can be extended after training: rescale the frequencies \(\theta\) and the same weights address a longer sequence (position interpolation, NTK scaling, YaRN¹⁰). Every "we extended it to 1M tokens" announcement is, mechanically, this plus fine-tuning.

Multi-head, and what it turned into

One head produces one distribution per position — one relationship. But a token usually needs several at once: a verb wants to find its subject and look at the previous word; a pronoun wants the noun it refers to and the punctuation that closes the sentence. A single row of weights that sums to 1 cannot point strongly at all of that. Multi-head attention runs \(h\) heads in parallel, each with its own \(W_Q^i, W_K^i, W_V^i\) on \(d/h\)-dimensional slices, and concatenates the outputs:

At \(d = 4096\) and \(h = 32\), each head works with 128-dimensional vectors; the 32 concatenated outputs are 4096 wide again, and \(W^O\) mixes what each head found. The head options in the matrix panel are exactly this kind of specialization.

Note what it does not cost: with \(d/h\) per head, \(h\) heads cost the same FLOPs as one head of full width. Heads are free parallelism, not extra capacity — which is why \(h\) grew and grew.

Then inference happened, and the design changed. The next section is why.

The KV cache is the thing nobody teaches and everybody pays for

A model generates text one token at a time, and every new token needs attention over all the previous ones. What does that require from each old token? Not its query: that only served to compute its own output, back then. But its key and value will be consulted by every token that comes after. Recomputing them at every step would make generating \(n\) tokens cost \(O(n^2)\) full forward passes — so nobody does. Every K and V is computed once and kept.

That cache is now the dominant memory cost of serving a model, and it is the reason "128k context" is a pricing decision rather than an architecture decision:

The 2 counts K and V; \(L\) is the number of layers, \(n\) the context length and \(B\) the batch size. Plug in a real model, Llama 3 8B: \(L = 32\), \(h_{kv} = 8\), \(d_\text{head} = 128\), in 16 bits (2 bytes). Each token costs \(2 \times 32 \times 8 \times 128 \times 2 = 131{,}072\) bytes, or 128 KiB. A 128k-token context therefore takes 16 GiB — as much as the model's own weights — and that is for one user.

Only one term in that product is a free design variable: \(h_{kv}\), the number of key/value heads. Queries can keep all \(h\) heads; K and V can be shared. Llama 3 8B has 32 query heads and only 8 key/value heads — with one per query the cache would be four times larger. This produced the sequence every current model sits somewhere on:

Making \(O(n^2)\) affordable

Attention is \(O(n^2 d)\) in time and, written naively, \(O(n^2)\) in memory: the matrix \(S\) from the shape table has one entry per pair of tokens. At \(n = 128\text{k}\) it has \(1.7 \times 10^{10}\) entries — 34 GB in 16 bits, per head, per layer. Four families of answers exist, and they are not equally successful.

Key takeaways

1. Attention is a dense layer whose weights are computed from the input — the opposite design choice from a convolution's fixed, local, shared weights. It beat recurrence on reach (one step to any token) and parallelism.
2. It works as a soft dictionary: it compares a question with every key and returns a weighted average of the values. The output is one vector with context per token.
3. Q, K, V are three projections of the same input. The separation is what lets a token ask for something it is not.
4. \(\sqrt{d_k}\) exists because score variance grows as \(d_k\). Without it the softmax saturates and the gradient vanishes. Log attention entropy; it is cheap and diagnostic.
5. The mask is one added matrix, and it is the entire encoder/decoder distinction. Forgetting the document mask when packing is a real and common pretraining bug.
6. Attention does not see order. RoPE encodes absolute rotations that yield relative offsets in the dot product — which is both why it works and why context can be extended by rescaling frequencies.
7. Heads are free parallelism, not extra capacity: \(h\) heads of width \(d/h\) cost what one head of width \(d\) costs.
8. The KV cache dominates inference memory (128 KiB per token in Llama 3 8B). \(h_{kv}\) is the one term you can shrink, which is why GQA is universal and MLA is where the frontier went. Inference is bandwidth-bound, not compute-bound: optimize bytes moved.
9. FlashAttention is exact and free. Sparsity, sliding windows and SSMs trade something; in 2026 the winning trades are hybrid stacks and trainable sparse attention.