The Attention Mechanism

The fix

Attention throws away the assumption that the decoder needs one summary. Instead, keep every encoder state, and at each output step let the decoder build a fresh context vector by weighting them:

1. Score. Compare the decoder's current state against every encoder state. One number per input word — how relevant is this word, right now?
2. Normalise. Push the scores through softmax so they are positive and sum to 1. These are the attention weights.
3. Blend. Take the weighted average of the encoder states. That is the context vector for this step, and it is different for every output word.

Nothing is discarded and nothing is compressed. The decoder reaches back into the full input every single time it produces a word.

Why the weights are interpretable

Because they sum to 1 and are all positive, the weights read as a distribution of interest: "to produce this word, I looked 70% at cat and 20% at the". Plotting them for every output word against every input word gives the alignment matrix in the second panel, and on a translation task it recovers word alignment without ever having been told what alignment is.

That interpretability is genuinely useful, but treat it carefully — attention weights show where information was drawn from, which is not quite the same as showing why the model decided what it decided.

Looking at everything at once

A recurrent network reads a sentence one word at a time, carrying a fixed-size hidden state. By word fifty, everything about word one has been compressed into that state, competing with forty-nine other words for room. Long-range dependencies are lost by construction.

Attention removes the bottleneck. Every position looks directly at every other position and decides, per position, how much each one matters.

Concretely, for each word the mechanism computes a weight against every other word, then produces a weighted average of their representations. "It" in "The cat sat on the mat because it was warm" can attend strongly to "mat" — a direct connection, not a signal relayed through five intermediate states.

Two consequences follow immediately, and they are why transformers replaced recurrent networks:

Path length is constant. Any two positions are one step apart, regardless of distance. In an RNN the path is proportional to the gap, and gradients decay along it.

It parallelises. Every position's attention can be computed simultaneously, as one matrix multiplication. An RNN must process position t before t+1, so it cannot use a GPU's parallelism across the sequence.

Queries, keys and values

The mechanism is best read as a soft dictionary lookup.

• The query is what this position is looking for.
• The key is what each position offers as a label.
• The value is what each position actually contributes if selected.

Compare the query against every key to get a score, turn the scores into weights, and take the weighted sum of the values:

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

Reading it left to right: QKᵀ is every query dotted with every key, giving a score matrix. Dividing by √dₖ keeps those scores in a range where softmax is not saturated — without it, large dimensions produce huge dot products, the softmax becomes nearly one-hot, and gradients vanish. The softmax turns scores into weights summing to 1. Multiplying by V produces the output.

All three — Q, K and V — are linear projections of the input, with learned weight matrices. That is the only place parameters enter the mechanism.

A worked example, small enough to follow

Three tokens, and suppose token 2's query produces these dot products against the three keys: 4.0, 1.0, 2.0. Divide by √dₖ = 2: 2.0, 0.5, 1.0.

Softmax: e² = 7.39, e⁰·⁵ = 1.65, e¹ = 2.72; total 11.76. Weights: 0.63, 0.14, 0.23.

The output for token 2 is 0.63 × value₁ + 0.14 × value₂ + 0.23 × value₃ — mostly the first token's contribution, with some of the third.

Those weights are what get visualised as attention heatmaps. They are genuinely informative about what the model connected to what, and they are not an explanation of why — a caution worth keeping, since attention maps are frequently over-interpreted.

The bottleneck it was invented to remove

Attention was not invented for transformers. It was a patch on sequence-to-sequence translation, and understanding the problem it patched makes the mechanism obvious rather than arbitrary.

import numpy as np

rng = np.random.default_rng(0)

print("THE ORIGINAL DESIGN (2014). an encoder RNN reads the whole source")
print("sentence into ONE fixed vector, and a decoder RNN generates from it:")
print("    source -> [encoder RNN] -> context vector -> [decoder RNN] -> target")
print()
H = 8
for n in (5, 15, 40, 100):
    print("   a %3d-word sentence -> the same %d numbers" % (n, H))
print("   that is the bottleneck. every sentence, however long, is squeezed")
print("   through a vector of fixed size, and the first words have been")
print("   overwritten many times by the time the encoder finishes.")
print()

print("here is that overwriting, measured. run an encoder over a sequence")
print("and ask how much of word 0 survives in the final state:")
Wx = rng.normal(0, 0.5, (4, H))
Wh = rng.normal(0, 0.5, (H, H))
def encode(seq):
    h = np.zeros(H)
    states = []
    for x in seq:
        h = np.tanh(x @ Wx + h @ Wh)
        states.append(h.copy())
    return h, states

print("   %10s %28s" % ("length", "effect of changing word 0"))
for L in (3, 6, 12, 25, 50):
    shifts = []
    for _ in range(40):
        seq = rng.normal(0, 1.0, (L, 4))
        base, _ = encode(seq)
        alt = seq.copy()
        alt[0] = rng.normal(0, 1.0, 4)               # change ONLY word 0
        changed, _ = encode(alt)
        shifts.append(np.linalg.norm(base - changed))
    print("   %10d %28.5f" % (L, np.mean(shifts)))
print("   by 25 words the first word barely registers. a translation model")
print("   built this way degrades sharply with sentence length, and that is")
print("   exactly what the 2014 papers reported.")
print()

print("THE FIX. keep every encoder state, and let the decoder choose which")
print("to look at as it generates each output word:")
src = ["le", "chat", "noir", "dort"]
tgt = ["the", "black", "cat", "sleeps"]
S, T = len(src), len(tgt)
enc_states = rng.normal(0, 1.0, (S, H))
dec_states = rng.normal(0, 1.0, (T, H))
print("   source: %s" % " ".join(src))
print("   target: %s" % " ".join(tgt))
print()

Wa = rng.normal(0, 0.5, (H, H))
scores = dec_states @ Wa @ enc_states.T
def softmax(z):
    e = np.exp(z - z.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)
A = softmax(scores)
print("   for each output word, a distribution over the source words:")
print("%12s %s" % ("", "".join("%9s" % w for w in src)))
for i, w in enumerate(tgt):
    print("%12s %s" % (w, "".join("%9.4f" % a for a in A[i])))
print("   row sums:", np.round(A.sum(1), 6))
print()
print("   these weights are random here. in a trained model they line up")
print("   with the translation -- and that alignment was the result that")
print("   made the technique famous, because nobody had supervised it. it")
print("   fell out of training the model to translate.")
print()

context = A @ enc_states
print("   each output step gets its OWN context vector, not one shared one:")
for i, w in enumerate(tgt):
    print("      %-8s context = %s ..." % (w, np.round(context[i][:4], 4)))
print("   the bottleneck is gone. the decoder has access to all %d encoder"
      % S)
print("   states at every step, weighted by what it currently needs.")
print()

print("THE THREE SCORING FUNCTIONS people used, before dot-product won:")
q, k = dec_states[0], enc_states[0]
v = rng.normal(0, 0.5, H)
Wc = rng.normal(0, 0.5, (H, H))
Wcat = rng.normal(0, 0.5, (2 * H, H))
print("   dot        q . k                    = %+8.4f" % (q @ k))
print("   general    q @ W @ k                = %+8.4f" % (q @ Wc @ k))
print("   additive   v . tanh(W @ [q; k])     = %+8.4f"
      % (v @ np.tanh(np.concatenate([q, k]) @ Wcat)))
print("   additive (Bahdanau) came first and works with different encoder")
print("   and decoder sizes. dot-product (Luong) is a single matmul with no")
print("   parameters at all, which is why it is what transformers use --")
print("   it maps straight onto hardware built for matrix multiplication.")
print()

print("WHAT CAME NEXT. the 2017 paper's contribution was noticing that if")
print("attention alone can connect any two positions, the RNN underneath is")
print("no longer necessary:")
print("%28s %14s %20s" % ("", "path length", "parallel over t?"))
print("%28s %14s %20s" % ("RNN", "O(n)", "no -- sequential"))
print("%28s %14s %20s" % ("RNN + attention", "O(1) to encoder", "no -- still an RNN"))
print("%28s %14s %20s" % ("transformer", "O(1)", "yes"))
print("   'path length' is how many steps a gradient travels between two")
print("   positions. an RNN's is the distance between them, which is the")
print("   vanishing gradient problem restated. attention's is 1.")
print()
print("so the sequence of ideas is: a fixed context vector is too small ->")
print("let the decoder look at all the encoder states -> if it can do that,")
print("it does not need the recurrence at all.")

Things to try

1. Walk the output. Set the Output Step slider to 1, then step it up one at a time. The bright connection moves across the source sentence as each output word is produced — the decoder is looking somewhere different every time.
2. See the bottleneck it replaced. Tick Bottleneck Mode. Now every output step gets the same flat average of the input, exactly as a fixed context vector would give it. The weights stop moving, focus collapses, and the model has nothing left to distinguish one step from another.
3. Sharpen the focus. Set the Sharpness slider to 4. The weights concentrate almost entirely on one source word and the entropy readout falls toward zero — this is a hard alignment.
4. Blur it. Set the Sharpness slider to 0.2. The weights spread almost evenly and entropy rises toward its maximum. The context vector becomes a bland average, which is the bottleneck failure again, arrived at from a different direction.
5. Watch the whole matrix. Look at the second panel while you change sharpness. The bright diagonal-ish band is the alignment the model has learned; blurring smears it into a uniform grey.

Scoring functions

Step 1 needs a way to compare two vectors. Two became standard:

• Additive (Bahdanau, 2014): push both vectors through a small feed-forward layer. Flexible, works when the two have different sizes, but adds parameters and is slower.
• Dot-product (Luong, 2015): just take the dot product. No parameters at all, and it is a matrix multiplication, which is exactly what accelerators are built for.

Dot-product won, and that decision is the reason transformers are fast. Its one wrinkle — that dot products grow with vector size and push softmax into saturation — is fixed by dividing by the square root of the dimension, the "scaled" in scaled dot-product attention.

What came next

Attention was invented as a patch for RNNs, and for three years that is what it was. Then came the observation that gave the paper Attention Is All You Need its title: if the decoder can reach any input position directly, the recurrence is no longer doing anything essential. Remove it, apply attention to the sequence against itself, and you get self-attention — which trains in parallel across all positions instead of stepping through them one at a time.

The next module generalises the three steps above into the query, key and value framing that every transformer is written in.

Where this goes wrong

• Treating attention weights as explanation. They show where information came from, not why the output was chosen. Different weightings can produce the same prediction.
• Forgetting the cost. Every output position scores against every input position, so the work grows with the product of the two lengths. On long sequences this quadratic term, not the model size, is what runs out of memory.
• Expecting clean alignments everywhere. Attention is often diffuse, and heads in a real model frequently attend to punctuation or to the first token for reasons that have nothing to do with meaning.

Summing up

Attention replaces the single fixed context vector with a fresh weighted average of every encoder state, recomputed at each output step: score the decoder state against each input position, softmax the scores into weights that sum to 1, and blend. Because nothing is compressed, long sentences stop degrading, and because the weights are a proper distribution they can be read as an alignment. Dot-product scoring made it fast enough to be worth doing everywhere, and once the decoder could reach any position directly, recurrence turned out to be optional.

Self-attention, cross-attention, and masking

Self-attention has queries, keys and values all derived from the same sequence — a sentence relating its own words to each other. This is what encoder layers do.

Cross-attention takes queries from one sequence and keys and values from another. In translation, the decoder's queries attend to the encoder's representation of the source sentence. In multimodal models, text queries attend to image patches.

Causal (masked) attention prevents a position from seeing anything after it, by setting those scores to negative infinity before the softmax so their weights become zero. This is what makes a decoder able to generate: predicting token t must not depend on token t+1, or the task would be trivial at training time and impossible at inference.

Padding masks serve a different purpose: they zero out the attention paid to padding tokens added for batching. Forgetting them means the model attends to filler as if it were content, which quietly degrades results.

The cost, and what is being done about it

Attention compares every position with every other, so cost grows with the square of the sequence length. Doubling the context quadruples the compute and the memory for the score matrix.

At 512 tokens that is 262,144 pairs — trivial. At 100,000 tokens it is ten billion, which is why long context was the central engineering problem of the last few years.

The main responses:

• FlashAttention computes exact attention without ever materialising the full score matrix, tiling the computation to fit in fast on-chip memory. Same result, several times faster, far less memory — now standard.
• Sparse and local attention restrict each position to a window plus a few global positions (Longformer, BigBird).
• Linear attention approximates the softmax to achieve linear scaling, at some quality cost.
• Multi-query and grouped-query attention share key and value projections across heads, which mainly reduces the memory needed to cache them during generation.
• KV caching stores the keys and values of previous tokens during generation so each new token costs one step rather than a full re-computation.

Questions people ask

Why divide by √dₖ? Dot products grow with dimension, and large values saturate the softmax into a near-one-hot distribution with tiny gradients. The scaling keeps them in a workable range.

Is attention the same as a weighted average? The output is one — but the weights are computed from the content itself, which is what makes it adaptive rather than fixed.

Does attention explain the model? It shows which positions were connected. That is not the same as why the prediction was made, and papers have shown attention weights can be altered substantially without changing the output.

How does attention know word order? It does not — the mechanism is permutation-invariant. Positional encodings are added to supply that information.

Why replace RNNs entirely? Constant path length between any two positions, and full parallelism across the sequence. Together those made much larger models trainable.

Can attention be used outside text? Yes — vision transformers attend across image patches, and it is used in graphs, audio and protein models.

Recap in one screen

• Attention lets every position draw directly from every other, with weights computed from content.
• Query, key, value: compare the query with all keys, softmax the scores, take the weighted sum of values.
• Divide by √dₖ or the softmax saturates and gradients vanish.
• Masking makes generation possible (causal) and batching correct (padding).
• Cost is quadratic in sequence length, which is what FlashAttention, sparse attention and KV caching address.