Self-Attention

What a word is actually doing

Every position emits a query, a key and a value. Each query is scored against every key — including its own — giving an n × n grid of scores. Softmax each row, and you have the matrix above: row i is how word i distributes its attention across the whole sentence.

The output for that word is the weighted blend of all the values. So after one self-attention layer, every word's representation has been rewritten as a mixture of the words it found relevant. The word "it" stops being a generic pronoun vector and becomes something closer to whatever it refers to.

The example worth staring at

"The animal didn't cross the street because it was too tired." What does "it" refer to?

Change one word — "because it was too wide" — and the answer flips from the animal to the street. Nothing about the grammar changed. The only way to resolve it is to notice which noun makes sense with the adjective, which means "it" has to be able to look at both nouns and at the adjective, all at once.

An RNN reading left to right has to carry that decision in its hidden state across six words. Self-attention just looks — but not, it turns out, in one go.

Layer 2 is where it resolves. By then "it" has absorbed a little of everything it attended to in layer 1, including the adjective, so the query it sends the second time is no longer generic — and that is enough to separate the two nouns. Switch to layer 2 and the reference follows the adjective: "tired" keeps it on the animal, "wide" moves it to the street.

This is the real reason transformers are deep. A single attention layer mixes information; it takes another to use what was mixed.

A sequence relating itself

Self-attention is attention where the queries, keys and values all come from the same sequence. Every token asks every token — including itself — how relevant it is, and builds a new representation from the answer.

That is what turns a bag of independent word vectors into contextual ones. Before self-attention, "bank" has one vector. After a layer of self-attention, the vector for "bank" in "river bank" has absorbed information from "river", and differs from the same word in "savings bank".

The mechanics for one token:

1. Project its input into a query.
2. Project every token, including itself, into keys and values.
3. Dot the query with each key, scale by √dₖ, softmax to get weights.
4. Output the weighted sum of the values.

Every token does this simultaneously, which is one matrix multiplication rather than a loop.

Why it needs positional information

Self-attention has no notion of order. Shuffle the input tokens and the set of outputs is the same set, just permuted — the mechanism sees a bag, not a sequence.

For language that is fatal: "the dog bit the man" and "the man bit the dog" would be indistinguishable.

Positional encodings supply the missing information by adding (or otherwise injecting) a position-dependent signal into the token representations before attention. The original transformer used fixed sinusoidal patterns; modern models mostly use rotary embeddings (RoPE), which rotate the query and key vectors by an angle proportional to position and generalise better to lengths not seen in training.

This is worth remembering as a design fact rather than a detail: attention supplies the relationships, positional encoding supplies the order, and neither works alone.

Multiple heads

One set of attention weights can only express one kind of relationship at a time. Multi-head attention runs several attention computations in parallel, each with its own learned projections, and concatenates the results.

With 8 heads and a model dimension of 512, each head works in 64 dimensions — the total cost is roughly unchanged, and the model gains several independent views of the sequence.

What the heads specialise in, when inspected, is genuinely varied: some track syntactic dependencies such as subject to verb, some attend to the previous token, some to matching brackets or quotation marks, some appear to do little at all. Studies that prune heads find many can be removed with little loss, which suggests a good deal of redundancy.

The practical reading: heads give the layer capacity to represent several relationships at once, and the specialisations that emerge are a side effect rather than something designed.

Attention, computed by hand on one sentence

Every step of scaled dot-product attention on a real six-word sentence, with the weights printed -- so you can see which word attends to which, and why the scaling factor is there.

import numpy as np

rng = np.random.default_rng(0)
words = ["the", "cat", "sat", "on", "the", "mat"]
T, D = len(words), 8

# stand-in embeddings: "cat" and "mat" deliberately point in similar
# directions, so we can check attention finds that
E = rng.normal(0, 1.0, (T, D))
E[5] = E[1] * 0.8 + rng.normal(0, 0.3, D)      # mat resembles cat
E[4] = E[0]                                     # the two "the"s are identical

Wq = rng.normal(0, 0.5, (D, D))
Wk = rng.normal(0, 0.5, (D, D))
Wv = rng.normal(0, 0.5, (D, D))

Q, K, V = E @ Wq, E @ Wk, E @ Wv
print("one sentence, %d tokens, %d dimensions each." % (T, D))
print("three projections of the same input: Q, K and V, each %s." % (Q.shape,))
print()

scores = Q @ K.T
print("STEP 1 -- every query against every key. score[i][j] is how much")
print("token i wants to look at token j:")
print("%8s %s" % ("", "".join("%9s" % w for w in words)))
for i, w in enumerate(words):
    print("%8s %s" % (w, "".join("%9.3f" % s for s in scores[i])))
print()

scaled = scores / np.sqrt(D)
print("STEP 2 -- divide by sqrt(d_k) = %.3f." % np.sqrt(D))
print("   raw score spread    : %.3f to %.3f" % (scores.min(), scores.max()))
print("   after scaling       : %.3f to %.3f" % (scaled.min(), scaled.max()))
print()
print("   why it matters: a dot product of d random terms has variance d, so")
print("   scores grow with dimension. watch what that does to the softmax:")
def peak(z):
    e = np.exp(z - z.max())
    return (e / e.sum()).max()

print("      (16 keys, averaged over 300 draws -- 1/16 = 0.0625 would be")
print("       perfectly uniform attention)")
for d in (8, 64, 512, 4096):
    unscaled, scaled = [], []
    for _ in range(300):
        q = rng.normal(size=d)
        ks = rng.normal(size=(16, d))
        raw = ks @ q
        unscaled.append(peak(raw))
        scaled.append(peak(raw / np.sqrt(d)))
    print("      d=%5d: unscaled peak %.4f, scaled peak %.4f"
          % (d, np.mean(unscaled), np.mean(scaled)))
print("   without the division the softmax saturates -- one token takes")
print("   almost all the weight and the gradient through the others dies.")
print()

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

A = softmax(scaled)
print("STEP 3 -- softmax each ROW, so each token's weights sum to 1:")
print("%8s %s" % ("", "".join("%9s" % w for w in words)))
for i, w in enumerate(words):
    print("%8s %s" % (w, "".join("%9.4f" % a for a in A[i])))
print("   row sums:", np.round(A.sum(1), 6))
print()
print("   note rows 0 and 4 are identical -- both are 'the', and attention")
print("   has no idea they are in different positions. that is exactly the")
print("   gap positional encoding exists to fill.")
print()

out = A @ V
print("STEP 4 -- each output is a weighted average of the value vectors:")
print("   output[1] ('cat') = %s" % np.round(out[1][:4], 4), "...")
manual = sum(A[1][j] * V[j] for j in range(T))
print("   the same thing, summed by hand: %s" % np.round(manual[:4], 4), "...")
print("   identical:", np.allclose(out[1], manual))
print()

print("what each token actually looked at -- its single strongest link:")
for i, w in enumerate(words):
    j = int(A[i].argmax())
    print("   %-5s attended most to %-5s (weight %.4f)" % (w, words[j], A[i][j]))
print()

print("the whole operation is three lines:")
print("   scores = Q @ K.T / sqrt(d_k)")
print("   weights = softmax(scores)")
print("   output = weights @ V")
print()
print("and its cost is the reason context windows are expensive. the score")
print("matrix is T x T:")
for t in (8, 512, 4096, 32768):
    print("   %6d tokens -> %14s entries, %8.1f MB at float32"
          % (t, "{:,}".format(t * t), t * t * 4 / 1024 / 1024))
print("   doubling the context quadruples this. every 'efficient attention'")
print("   paper is an attempt to avoid building that matrix in full.")

Guided tour

1. Watch one layer fail. Set the Layer slider to 1 and the Focus Row to it. Note what it attends to, then switch the Example to the "wide" sentence. The number does not move at all — a single layer's query for "it" cannot depend on a word elsewhere in the sentence.
2. Then watch depth fix it. Set the Layer slider to 2 and switch between the two examples again. Now the reference follows the adjective: "tired" holds it on the animal, "wide" moves it to the street. Layer 2's query was built from what layer 1 collected.
3. Confirm every row is a distribution. Change the Focus Row to any word and check Row Sum. It is always 1 — softmax runs across each row independently.
4. Blur the whole thing. Set the Sharpness slider to 0.2. The matrix goes flat grey: every word attends equally to everything, which is the same as attending to nothing. Set it to 5 and the matrix turns into a near-binary lookup.
5. Cut off the future. Tick Causal Mask. The upper triangle goes black — each position can now only see itself and what came before. This single change is the whole architectural difference between BERT and GPT.
6. Prove it is order-blind. Tick Shuffle Word Order. The words move, but each word keeps exactly the same attention pattern toward the same partners, because the mechanism never consulted position at all. This is precisely why positional encoding has to exist.
7. Read the price. Watch Scores Computed. Ten words need 100 scores; double the sentence and it quadruples. Path Length stays at 1 no matter what — that is the trade self-attention makes.

Why this killed recurrence

The first two rows are why transformers took over. A recurrent network cannot start step 50 until step 49 is done, so training time scales with sentence length and gradients have to survive fifty multiplications. Self-attention computes every position simultaneously, which is what let models grow to the size they are now.

The last two rows are the bill. The quadratic cost is why context windows were small for years, and the order-blindness is a genuine hole that has to be patched.

Traps worth knowing

• Forgetting positional information. Without it, "dog bites man" and "man bites dog" are the same input. Self-attention is permutation-equivariant by construction.
• Assuming the matrix is symmetric. It is not. Word A attending to B says nothing about B attending to A, because queries and keys are different projections.
• Reading heads as syntax. A trained model's heads often attend to punctuation or dump weight on the first token when they have nothing useful to do. Clean linguistic patterns are the exception, not the rule.
• Underestimating n². Doubling the context quadruples both time and memory for attention. Almost every "efficient transformer" paper is an attack on that one term.

Where that leaves you

Self-attention is attention with queries, keys and values all drawn from the same sequence, so every word is rewritten as a blend of the words it finds relevant — which is how a pronoun can resolve to a noun six words away in a single step. Because every position is computed independently, training parallelises completely and the path between any two words is length 1 regardless of distance, and those two properties are what displaced recurrence. One layer only mixes information, though — resolving a pronoun against a distant adjective takes a second layer to use what the first collected, which is why the architecture is stacked. The costs are real and specific: the number of scores grows with the square of the sequence length, and the mechanism is entirely blind to word order until positional information is added.

Where it sits in a transformer block

Self-attention is one of two sublayers, and the surrounding structure matters as much as the mechanism.

x → LayerNorm → Self-attention → (+x) → LayerNorm → Feed-forward → (+x)

The residual connections give gradients a clean path back through dozens of layers. Without them, deep transformers do not train.

LayerNorm keeps activations in a stable range, and placing it before each sublayer (pre-norm) rather than after is what makes very deep stacks trainable.

The feed-forward network is where most of the parameters live — typically expanding to four times the model dimension and back. Attention moves information between positions; the feed-forward layer processes each position independently. Both are needed, and it is a common misconception that attention is doing all the work.

Stack that block 12, 32 or 80 times and you have a transformer.

Complexity, and reading an attention map

Every token attends to every token, so the score matrix is n×n for a sequence of length n. That quadratic cost is the defining constraint of the architecture, and the reason context windows were limited for years.

Attention maps — the n×n weights, visualised as a heatmap — are the most common interpretability tool for transformers, and the most commonly over-read. Three cautions:

• They show what was connected, not what was used. The value vectors and the subsequent layers determine the effect.
• Later layers attend to representations that already mix many positions, so a weight against token 5 does not mean information from token 5 specifically.
• Work on attention as explanation has shown weights can often be substantially altered without changing the prediction.

Useful for intuition and for spotting gross problems; not evidence about reasoning.

Questions people ask

What is the difference between attention and self-attention? Self-attention is the case where queries, keys and values come from the same sequence. Cross-attention draws them from two.

Why does each token attend to itself? Because its own content is usually the most relevant thing available, and excluding it would force the representation to be built entirely from other positions.

How many heads should I use? 8–16 for typical model sizes, with the head dimension usually 64. It is rarely the parameter worth tuning.

Does self-attention replace the feed-forward layer? No — they do different jobs, and most of a transformer's parameters are in the feed-forward layers.

Is self-attention permutation-invariant? Yes, which is precisely why positional encodings exist.

Why is it quadratic? Every position is compared with every other, giving n² pairs.

Recap in one screen

• Self-attention lets each token build a new representation from a content-weighted mix of all tokens.
• That is what makes representations contextual: "bank" differs by sentence.
• It is order-blind on its own; positional encodings supply the sequence.
• Multiple heads give several relationship types at once, at roughly constant cost.
• It lives inside a block with residual connections, LayerNorm and a feed-forward network — all four matter.