Attention &
The Transformer
The mechanism that let sequences attend to any other part of themselves simultaneously — no sequential bottleneck, no vanishing gradient through time. You'll derive self-attention from a worked numerical example, then extend it to multi-head attention and positional encoding.
The Limitation That Attention Solves
The seq2seq model from Module 6 has a fundamental problem: no matter how long the input sequence is, everything must be compressed into a single fixed-size context vector — the encoder's final hidden state. For a 200-word sentence, that vector must somehow encode every relevant detail the decoder might need. That's too much to ask of a single vector.
In 2015, Bahdanau et al. proposed a solution: instead of compressing everything into one vector upfront, let the decoder look directly at all encoder hidden states and decide, at each decoding step, which encoder positions are most relevant to the current output. This is attention.
- 1Encoder: Instead of keeping only the final hidden state, keep all T hidden states h₁, h₂, ..., hₜ — one per input token.
- 2Alignment scores: At each decoder step, compute a score between the current decoder state and each encoder hidden state: eᵢⱼ = score(sᵢ, hⱼ). High score = "pay more attention to encoder position j at decoder step i".
- 3Attention weights: Normalize the scores with softmax: αᵢⱼ = softmax(eᵢⱼ). These sum to 1 and act as a probability distribution over encoder positions.
- 4Context vector: Weighted sum of encoder hidden states: cᵢ = Σⱼ αᵢⱼ hⱼ. This is a dynamic, input-dependent summary — different at each decoder step.
- 5Decoder step: Decode using both the decoder's hidden state and this context vector. The decoder can now "look at" any part of the input as needed.
RNN bottleneck (before)
All input information must flow sequentially. Path length between word 1 and word 100 is 100 steps. Gradients from distant positions vanish. Can't parallelize across time steps — slow training.
Attention solution (after)
Any position can attend to any other position directly. Path length between word 1 and word 100 is 1 step. No sequential bottleneck. Can process all positions in parallel. Training on GPUs is extremely fast.
The cost of attention
Attention computes similarity between every pair of positions. For sequence length T, this is O(T²) operations. For T=512 tokens, fine. For T=100,000 tokens (books, long documents), this is expensive — active research area.
Interpretability bonus
Attention weights are readable. You can visualize which input tokens the model attended to when producing each output token. This interpretability was a major advantage over the black box of RNN hidden states.
Self-Attention, Step by Step — With a Worked Numerical Example
Self-attention is attention applied within a single sequence — each token attending to every other token in the same sequence. It's the core operation of the transformer. We derive it completely, then trace through a numerical example so the formulas have concrete meaning.
Self-attention uses three linear projections of each input token — Query (Q), Key (K), and Value (V). The intuition is a soft retrieval system: each token's Query asks "what am I looking for?", each Key says "here's what I contain", and the Value is what actually gets retrieved and aggregated.
Attention(Q, K, V) = softmax(QKᵀ / √dₖ) · V
We have 3 tokens (e.g. "I", "love", "math") each embedded in 4 dimensions, and projection dimension dₖ = 3.
🔢 Step-by-step with numbers
[0.3, 1.0, 0.7, 0.1], ← "love"
[0.6, 0.4, 1.0, 0.2]] ← "math"
K = X @ W_K → shape (3, 3)
V = X @ Wᵥ → shape (3, 3)
S[i,j] = dot product of token i's Query with token j's Key
= how much token i "wants to attend to" token j
Example S (before scaling):
[[2.1, 1.4, 0.8], ← "I" attends most to itself (2.1)
[0.9, 3.2, 1.1], ← "love" attends most to itself (3.2)
[1.3, 1.7, 2.9]] ← "math" attends most to itself (2.9)
S_scaled = S / 1.73
= [[1.21, 0.81, 0.46],
[0.52, 1.85, 0.64],
[0.75, 0.98, 1.68]]
Without scaling: for large dₖ, dot products grow large,
softmax saturates (one value dominates), gradients vanish.
Row 0 ("I"): [0.44, 0.29, 0.27] ← attends 44% to "I", 29% to "love"
Row 1 ("love"): [0.21, 0.80, 0.26] ← strongly attends to itself (80%!)
Row 2 ("math"): [0.27, 0.34, 0.69] ← attends most to itself, some to "love"
Each output token is a weighted average of all Value vectors,
weighted by how much it attended to each position.
"math"'s output vector = 0.27 × V["I"] + 0.34 × V["love"] + 0.69 × V["math"]
It has folded in some information from every token, weighted by relevance.
import torch
import torch.nn as nn
import math
class SelfAttention(nn.Module):
"""Single-head scaled dot-product attention."""
def __init__(self, d_model, d_k):
super().__init__()
self.d_k = d_k
self.W_Q = nn.Linear(d_model, d_k, bias=False)
self.W_K = nn.Linear(d_model, d_k, bias=False)
self.W_V = nn.Linear(d_model, d_k, bias=False)
def forward(self, x, mask=None):
# x: (batch, seq_len, d_model)
Q = self.W_Q(x) # (batch, seq_len, d_k)
K = self.W_K(x)
V = self.W_V(x)
# Scaled dot-product attention
scores = Q @ K.transpose(-2, -1) / math.sqrt(self.d_k)
# scores: (batch, seq_len, seq_len)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9) # causal mask
attn_weights = torch.nn.functional.softmax(scores, dim=-1)
output = attn_weights @ V # (batch, seq_len, d_k)
return output, attn_weights
# Test
attn = SelfAttention(d_model=64, d_k=32)
x = torch.randn(4, 10, 64) # batch=4, seq_len=10, d_model=64
out, weights = attn(x)
print("Output shape: ", out.shape) # (4, 10, 32)
print("Weights shape: ", weights.shape) # (4, 10, 10) — attention matrix
print("Weights sum: ", weights.sum(dim=-1).mean().item()) # ≈1.0 (softmax)Multi-Head Attention
Single-head attention computes one attention distribution — one way of weighting which tokens to attend to. But sequences have multiple kinds of relationships simultaneously: syntactic (subject-verb), semantic (coreference), positional (nearby tokens). Multi-head attention runs several attention mechanisms in parallel, each free to learn different relationship types.
MultiHead(Q, K, V) = Concat(head₁, ..., head_h) · W_O
Each head has dimension dₖ = d_model / h
Total parameters ≈ same as single-head with d_model
Each head sees differently
One head might learn to attend to syntactic structure (verb → object). Another might track coreference ("it" → "the cat"). Another might encode positional proximity. The model discovers what's useful — you don't specify.
Subspace projection
Each head projects to a lower-dimensional subspace (d_model / h). Multiple heads give multiple low-dimensional perspectives, which the output projection W_O combines back into the full space.
Parallelism
All heads compute simultaneously — no sequential dependency. On a GPU, computing 8 attention heads in parallel is nearly as fast as computing 1, because the operations are batch matrix multiplications.
Choosing h
Common values: 8 (original transformer), 12 (BERT-base), 16 (BERT-large), 32 (GPT-3 large). Must divide evenly into d_model. More heads = more diverse perspectives but not always better — there's a sweet spot.
import torch
import torch.nn as nn
import math
class MultiHeadAttention(nn.Module):
def __init__(self, d_model=512, n_heads=8):
super().__init__()
assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
self.d_model = d_model
self.n_heads = n_heads
self.d_k = d_model // n_heads
# Single linear projections — efficient: all heads computed together
self.W_Q = nn.Linear(d_model, d_model, bias=False)
self.W_K = nn.Linear(d_model, d_model, bias=False)
self.W_V = nn.Linear(d_model, d_model, bias=False)
self.W_O = nn.Linear(d_model, d_model, bias=False) # output projection
def split_heads(self, x):
# x: (batch, seq, d_model) → (batch, heads, seq, d_k)
B, T, _ = x.shape
x = x.view(B, T, self.n_heads, self.d_k)
return x.transpose(1, 2) # (B, heads, T, d_k)
def forward(self, x, mask=None):
B, T, _ = x.shape
Q = self.split_heads(self.W_Q(x)) # (B, heads, T, d_k)
K = self.split_heads(self.W_K(x))
V = self.split_heads(self.W_V(x))
# Scaled dot-product: all heads at once via batch matmul
scores = Q @ K.transpose(-2, -1) / math.sqrt(self.d_k)
# scores: (B, heads, T, T)
if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
attn = torch.nn.functional.softmax(scores, dim=-1)
out = attn @ V # (B, heads, T, d_k)
# Concatenate heads: (B, heads, T, d_k) → (B, T, d_model)
out = out.transpose(1, 2).contiguous().view(B, T, self.d_model)
return self.W_O(out), attn # output projection + attention weights
# Test
mha = MultiHeadAttention(d_model=512, n_heads=8)
x = torch.randn(2, 16, 512)
out, attn_weights = mha(x)
print("Output:", out.shape) # (2, 16, 512)
print("Attention:", attn_weights.shape) # (2, 8, 16, 16) — 8 heads, 16×16 matrixPositional Encoding
Self-attention is permutation-invariant: if you shuffle the tokens, the attention output is shuffled the same way but otherwise unchanged. The model has no built-in sense of order. "The cat sat on the mat" and "mat the on sat cat the" would produce identical attention patterns (just reordered). This is a problem — word order matters enormously in language.
Positional encoding injects order information by adding a position-dependent vector to each token's embedding before it enters the transformer. The original paper used fixed sinusoidal encodings; modern models learn positional embeddings.
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
Input to transformer = token_embedding + positional_encoding
The elegance of sinusoidal encoding: the model can extrapolate to sequence lengths not seen during training (because the pattern is defined mathematically, not memorized). The relative position between tokens can be computed as a linear function of the encodings — which attention can learn to use.
Sinusoidal (original)
Fixed, not learned. Different frequencies for different dimensions — low dimensions encode coarse position, high dimensions encode fine position. Allows extrapolation to unseen lengths. Used in the original transformer.
Learned absolute (BERT, GPT)
A lookup table of position embeddings, one per position up to max_length. Learned during training alongside token embeddings. Simple and effective. Can't generalize beyond max_length seen during training.
Relative position (RoPE, ALiBi)
Encode relative distances between tokens rather than absolute positions. RoPE (used in LLaMA, GPT-NeoX) rotates Q and K vectors by position — attention scores naturally encode relative distance. Better length generalization.
Why not use index directly?
A raw integer (position 50) vs position (position 0) has a 50× magnitude difference. This would overwhelm the token embeddings. Encodings must have the same magnitude range as embeddings to be additive.
import torch
import torch.nn as nn
import math
class SinusoidalPE(nn.Module):
def __init__(self, d_model=512, max_len=5000, dropout=0.1):
super().__init__()
self.dropout = nn.Dropout(dropout)
# Compute fixed PE matrix: (max_len, d_model)
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len).unsqueeze(1) # (max_len, 1)
div_term = torch.exp(
torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)
)
pe[:, 0::2] = torch.sin(position * div_term) # even dims
pe[:, 1::2] = torch.cos(position * div_term) # odd dims
pe = pe.unsqueeze(0) # (1, max_len, d_model) — broadcast over batch
# Register as buffer: part of state_dict, but not a parameter (not trained)
self.register_buffer('pe', pe)
def forward(self, x):
# x: (batch, seq_len, d_model)
x = x + self.pe[:, :x.size(1)] # add position encoding
return self.dropout(x)
# Learned positional embedding (simpler, used in BERT/GPT)
class LearnedPE(nn.Module):
def __init__(self, d_model=512, max_len=512):
super().__init__()
self.pe = nn.Embedding(max_len, d_model)
def forward(self, x):
B, T, _ = x.shape
positions = torch.arange(T, device=x.device).unsqueeze(0) # (1, T)
return x + self.pe(positions) # broadcasts over batch
# Test sinusoidal PE
pe_layer = SinusoidalPE(d_model=128)
x = torch.randn(4, 20, 128) # batch=4, seq=20, dim=128
out = pe_layer(x)
print(out.shape) # (4, 20, 128) — same shape, position info added