Part 4: Transformers and Foundation Models¶

Lesson 44: The Attention Mechanism¶

Every model in Part 3 processed spatial structure through convolution: a fixed, local neighborhood pattern, applied the same way everywhere (Lesson 10). Attention is a different way to mix information across a sequence or an image: instead of a fixed neighborhood, every position decides, dynamically, which other positions are relevant to it right now, and mixes them in proportion to that relevance. This lesson builds attention from scratch, validates it against PyTorch's own implementation, and shows concretely what it's doing: soft, content-based retrieval.

In [1]:

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

Queries, keys, and values¶

Every position in a sequence produces three vectors: a query (what it's looking for), a key (what it advertises about itself, for others to match against), and a value (what it actually contributes if selected). Attention compares every query against every key via a dot product — a similarity score — turns those scores into weights with softmax, and returns a weighted sum of values:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

The $\sqrt{d_k}$ scaling (dividing by the square root of the key dimension) keeps the dot products from growing too large as dimensionality increases, which would otherwise push softmax into a regime with vanishing gradients (Lesson 35) — the same instinct as Lesson 32's weight-init scaling, applied to attention scores instead of layer activations.

In [2]:

def attention(Q, K, V, mask=None):
    d_k = Q.shape[-1]
    scores = Q @ K.transpose(-2, -1) / np.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    weights = F.softmax(scores, dim=-1)
    return weights @ V, weights

torch.manual_seed(0)
B, T, D = 2, 5, 8  # batch, sequence length, dimension
Q = torch.randn(B, T, D)
K = torch.randn(B, T, D)
V = torch.randn(B, T, D)

out, weights = attention(Q, K, V)
out_torch = F.scaled_dot_product_attention(Q, K, V)

print(f'output shape: {tuple(out.shape)}')
print(f'each attention-weight row sums to 1: {torch.allclose(weights.sum(-1), torch.ones(B, T))}')
print(f'max abs diff vs F.scaled_dot_product_attention: {(out - out_torch).abs().max().item():.2e}')

output shape: (2, 5, 8)
each attention-weight row sums to 1: True
max abs diff vs F.scaled_dot_product_attention: 2.38e-07

Attention as soft, content-based retrieval¶

Abstract formula aside, here's concretely what attention does: given a query, it looks up the most similar keys in the sequence and returns a blend of their values, weighted by similarity. Build a toy sequence of 5 distinct "tokens" (one-hot vectors, scaled up to make similarity differences sharper) and query for the token that matches position 2 exactly.

In [3]:

vocab = torch.eye(6) * 5.0  # 6 distinct tokens, scaled to sharpen dot-product contrast
seq = vocab[[0, 1, 2, 3, 4]].unsqueeze(0)     # a 5-token sequence: tokens 0,1,2,3,4
query = vocab[2].unsqueeze(0).unsqueeze(0)    # a query vector identical to token 2

out_retrieve, w_retrieve = attention(query, seq, seq)
retrieved_token = out_retrieve[0, 0].argmax().item()

print(f'attention weights over the 5 sequence positions: {w_retrieve[0, 0].round(decimals=3).tolist()}')
print(f'retrieved token id: {retrieved_token}  (queried for token 2)')

plt.figure(figsize=(5, 1.5))
plt.imshow(w_retrieve[0].numpy(), cmap='viridis', aspect='auto')
plt.xlabel('sequence position (key)'); plt.yticks([])
plt.title('Attention weights: query = token 2')
plt.colorbar(fraction=0.02)
plt.show()

attention weights over the 5 sequence positions: [0.0, 0.0, 1.0, 0.0, 0.0]
retrieved token id: 2  (queried for token 2)

No description has been provided for this image

Attention correctly pulls out token 2, weighting position 2 far above the others (visualized as the bright cell in the heatmap). This is the mechanism, stripped down: no convolution kernel, no fixed neighborhood — just "find what's similar, weighted-average it in."

Causal masking¶

Some tasks (predicting the next token in a sequence) require that a position can only attend to positions at or before itself — attending to the future would let the model "cheat" by looking at the answer. This is enforced with a mask: scores for disallowed (future) positions are set to $-\infty$ before the softmax, so they get exactly zero weight.

In [4]:

causal_mask = torch.tril(torch.ones(T, T))  # 1 where position i may attend to position j <= i
out_causal, w_causal = attention(Q, K, V, mask=causal_mask)
out_torch_causal = F.scaled_dot_product_attention(Q, K, V, is_causal=True)

future_weight = w_causal[:, 0, 1:].abs().sum().item()  # position 0 attending to positions 1..T-1
print(f'total attention weight from position 0 onto FUTURE positions: {future_weight:.6f}  (should be exactly 0)')
print(f'max abs diff vs F.scaled_dot_product_attention(..., is_causal=True): {(out_causal - out_torch_causal).abs().max().item():.2e}')

plt.figure(figsize=(4, 3.5))
plt.imshow(w_causal[0].numpy(), cmap='viridis')
plt.xlabel('key position'); plt.ylabel('query position')
plt.title('Causal attention mask\n(every row only attends to itself and the past)')
plt.colorbar(fraction=0.046)
plt.show()

total attention weight from position 0 onto FUTURE positions: 0.000000  (should be exactly 0)
max abs diff vs F.scaled_dot_product_attention(..., is_causal=True): 1.19e-07

Causal masking is exactly why GPT-style models are called "autoregressive": at generation time, each new token's query can only see keys and values from tokens already generated, never ones that don't exist yet. Bidirectional models (like BERT, or every vision transformer in the next lesson) skip this mask entirely — an image doesn't have a "future," so every patch is free to attend to every other patch.

Exercise¶

In the retrieval demo, change query to vocab[2] + vocab[4] (a blend of two tokens) instead of a pure match. Does attention now split its weight between positions 2 and 4? Does the split look roughly even, and can you explain why from the dot-product formula?
Remove the / np.sqrt(d_k) scaling from attention and rerun the retrieval demo with D increased from 8 to 128. Do the attention weights become more or less "peaked" (closer to one-hot) as dimension grows without scaling — and does adding the scaling back restore the original behavior?
Build a mask that allows each position to attend only to itself and its immediate neighbors (a "local attention" window, width 3) instead of the full causal triangle. How does the resulting weight pattern compare to a Lesson 10 convolution kernel applied along the sequence?

← Lesson 42: Instance Segmentation

Lesson 45: The Transformer Architecture →