Unit 08.04: Attention as weighted context

Attention replaced recurrence by removing the sequential bottleneck entirely. The mechanism is a weighted average, and you can write it in three lines.

Every token looks at every token

For each token, compute how similar it is to every other token, turn those similarities into weights that sum to 1, then take the weighted average of all tokens. That average becomes the token's new representation.

Here queries, keys and values are the same vectors, which is self-attention in its simplest form:

import torch
import torch.nn.functional as F

torch.manual_seed(0)
# Four tokens, each a 4-dimensional vector.
tokens = torch.tensor([
    [1.0, 0.0, 0.0, 0.0],    # "the"
    [0.0, 1.0, 0.0, 0.0],    # "cat"
    [0.0, 0.0, 1.0, 0.0],    # "sat"
    [0.0, 1.0, 0.1, 0.0],    # "it"   -- similar to "cat"
])

# Scaled dot-product attention, written out. Q, K, V are the same here.
scores = tokens @ tokens.T / (tokens.shape[1] ** 0.5)
weights = F.softmax(scores, dim=1)
context = weights @ tokens

print("attention weights (rows = query token):")
for name, row in zip(["the", "cat", "sat", "it"], weights):
    print(f"  {name:4} -> {[round(v, 3) for v in row.tolist()]}")

print("\nevery row sums to 1:", [round(v, 3) for v in weights.sum(dim=1).tolist()])
print("'it' attends most to  :", ["the", "cat", "sat", "it"][weights[3].argmax().item()])

# That is the whole idea: each token's new representation is a weighted average
# of every token, with weights decided by similarity. No recurrence, and every
# position is computed in parallel.
print("context shape:", tuple(context.shape))

Every row of the weight matrix sums to 1 — that is the softmax guarantee. Read the row for "it": the largest weight goes to "cat", because their vectors are similar, so "it" is represented partly by "cat".

The scaling by the square root of the dimension keeps the dot products from growing with vector size, which would otherwise push softmax into a near-one-hot regime with vanishing gradients.

The mistake this prevents

Treating attention weights as an explanation of the model's decision. They show what was averaged together at one layer, not why the final answer came out as it did.

Takeaway

Attention is a weighted average whose weights come from similarity. No recurrence, and every position computed in parallel.

Mark complete