Attention is the single most important innovation in modern LLMs. This week we derive it from first principles, implement it in PyTorch, and develop intuitions for what attention heads learn to do.
In early sequence-to-sequence models (encoder-decoder RNNs for machine translation), the entire source sentence was compressed into a single fixed-size vector. This bottleneck caused quality to degrade for long sentences.
Attention (Bahdanau et al., 2015) was introduced as a solution: instead of compressing to a single vector, let the decoder look back at all encoder states and weight them dynamically based on relevance.
The modern formulation (Vaswani et al., 2017) is clean and fully parallelisable.
Given a sequence of n tokens, we compute three matrices from the input embeddings X ∈ R^{n × d}:
Q = X W_Q (Queries, n × d_k)
K = X W_K (Keys, n × d_k)
V = X W_V (Values, n × d_v)
W_Q, W_K, W_V are learned weight matrices.
Attention(Q, K, V) = softmax(Q K^T / sqrt(d_k)) V
For token i, the query vector asks a question; the key vectors of all tokens answer that question; the resulting attention weights determine how much of each value to mix in.
In a decoder-only LLM (GPT family), a token must not attend to future tokens (that would be cheating during training and generation).
We apply a causal mask: set Q K^T entries for j > i to -∞ before the softmax.
import torch
import torch.nn.functional as F
import math
def causal_self_attention(Q, K, V):
d_k = Q.size(-1)
scores = Q @ K.transpose(-2, -1) / math.sqrt(d_k)
n = scores.size(-1)
mask = torch.triu(torch.ones(n, n), diagonal=1).bool()
scores = scores.masked_fill(mask, float('-inf'))
weights = F.softmax(scores, dim=-1)
return weights @ V, weights
A single attention head can only ask one “question” at a time. Multi-head attention runs h parallel attention heads, each with its own projection matrices.
MultiHead(Q, K, V) = Concat(head_1, ..., head_h) W_O
head_i = Attention(Q W_Q_i, K W_K_i, V W_V_i)
The output is projected back to d_model with W_O.
Why multiple heads? Different heads learn to capture different types of relationships:
In encoder-decoder transformers (T5, original Transformer), the decoder uses cross-attention to attend to the encoder’s output.
The queries come from the decoder; the keys and values come from the encoder. This is how the model conditions generation on the source sequence.
Standard self-attention is O(n²·d) in time and O(n²) in memory, where n is the sequence length. This is why long contexts are expensive.
For context windows of 128k tokens, this would require ~16GB of memory just for attention weights. Efficient variants (FlashAttention, sparse attention, linear attention) address this.
Computes attention in tiles, never materialising the full n×n matrix. Provides:
Used by default in modern LLM implementations.
Attention weights are interpretable (to some degree). Tools like BertViz allow you to interactively visualise which tokens attend to which.
# Conceptual example of extracting attention weights
from transformers import AutoTokenizer, AutoModel
import torch
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModel.from_pretrained("bert-base-uncased", output_attentions=True)
inputs = tokenizer("The dog chased the cat", return_tensors="pt")
outputs = model(**inputs)
# outputs.attentions: tuple of (batch, heads, seq_len, seq_len) per layer
attn = outputs.attentions[0][0] # layer 0, batch 0: shape (heads, seq_len, seq_len)
See practicals/week04_practical.py:
Lesson_3-selfattention.ipynb