The Attention Mechanism

Quick Context

The encoder–decoder models that came before attention worked like this: an RNN read the whole input sentence and squeezed it into a single fixed-length vector. And a second RNN made the output from that vector alone.

The flaw is visible as soon as you say it out loud. Each sentence, whether five words or fifty, has to fit in the same number of floats. The encoder's final state is ruled by any it read last. And the beginning of a long sentence is simply gone by the time decoding starts. Translation quality fell off a cliff past about twenty words. Attention (Bahdanau et al., 2014) removes the bottleneck fully. Instead of compressing the input into one vector, the decoder looks at all of the encoder's hidden states and computes a weighted average, with weights that depend on the current decoding step. At each output position, the decoder gives a query from its current state, scores it against each encoder state (the keys), and uses the scores to compute a weighted sum (the context vector). The scores go through softmax, so they sum to 1. And the highest-scoring encoder state rules the context. This lets the model "attend" to the apt part of the input for each output word. When translating "the cat sat," the word "chat" attends mostly to "cat." The scoring function can be a dot product (Luong attention), an additive network (Bahdanau attention). Or a scaled dot product (transformer attention). Self-attention applies the same method within a single sequence rather than between encoder and decoder. This is the foundation of the transformer architecture. Multi-head attention runs a few attention steps in parallel to capture different types of links. This page shows the attention weights at each decoding step, highlights. This input words the model focuses on, and covers the three scoring functions.

The fix

Attention throws away the assumption that the decoder needs one summary. Instead, keep each 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 each encoder state. One number per input word — how apt 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 each output word.

Nothing is discarded and nothing is compressed. The decoder reaches back into the full input each single time it gives 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 give this word, I looked 70% at cat and 20% at the". Plotting them for each output word against each 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 truly useful, but treat it carefully. Attention weights show where data was drawn from. This is not quite the same as showing why the model decided what it decided.

Looking at all of it at once

A recurrent network reads a sentence one word at a time, carrying a fixed-size hidden state. By word fifty, all of it 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. Each position looks right at each other position and decides, per position, how much each one matters.

Concretely, for each word the method computes a weight against each other word. Then gives 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 middle states.

Two results follow at once, and they are why transformers replaced recurrent networks:

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

It parallelises. Each position's attention can be computed at once, as one matrix product. 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 method 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 really adds if selected.

Compare the query against each 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 each query dotted with each key. This gives a score matrix. Dividing by √dₖ keeps those scores in a range where softmax is not saturated. Without it, large dimensions give huge dot products, the softmax becomes nearly one-hot, and gradients vanish. The softmax turns scores into weights summing to 1. Multiplying by V gives 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 method.

A worked case, small enough to follow

Three tokens, and suppose token 2's query gives 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 share, with some of the third.

Those weights are what get visualised as attention heatmaps. They are truly informative about what the model connected to what. And they are not an reason of why. A caution worth keeping, since attention maps are often 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 method 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 is squeezed through a")
print("   fixed-size vector, and the first words fade as the encoder runs.")
print()

print("measured: run an encoder and ask how much of word 0 survives:")
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 with sentence length, as the 2014 papers showed.")
print()

The fix: keep each encoder state and let the decoder choose which ones to look at as it makes each output word.

import numpy as np

rng = np.random.default_rng(0)
H = 8

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 align with")
print("   the translation -- that unsupervised alignment made the technique famous.")
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()

Here are the three scoring functions people used before dot-product attention won, and how the transformer removed the RNN fully.

import numpy as np

rng = np.random.default_rng(0)
H = 8
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 gave. The decoder is looking somewhere different each time.
2. See the bottleneck it replaced. Tick Bottleneck Mode. Now each 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 fully 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 product. This 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 right, the recurrence is no longer doing anything key. 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 extends the three steps above into the query, key and value framing that each transformer is written in.

Where this goes wrong

• Treating attention weights as reason. They show where data came from, not why the output was chosen. Different weightings can give the same prediction.
• Forgetting the cost. Each output position scores against each 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 all over. Attention is often diffuse, and heads in a real model often 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 each 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 all over, and once the decoder could reach any position right, 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 not doable 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 each position with each 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 work to fit in fast on-chip memory. Same result, a few 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 reach linear scaling, at some quality cost.
• Multi-query and grouped-query attention share key and value projections across heads. This mainly reduces the memory needed to cache them during generation.
• KV caching stores the keys and values of last tokens during generation so each new token costs one step rather than a full re-work.

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. This 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 much without changing the output.

How does attention know word order? It does not. The method is permutation-invariant. Positional encodings are added to supply that data.

Why replace RNNs fully? 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 each position draw right from each 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 doable (causal) and batching correct (padding).
• Cost is quadratic in sequence length, which is what FlashAttention, sparse attention and KV caching address.