Area 01 // LLMsLesson 14 / 14015 min read
Self-attention
Self-attention lets every token look at the other tokens in the text and pull in information from the ones that matter to it, which is how a model works out that "it" refers to "the animal" and not "the street".
Think of it like this
Analogy
You're at a noisy party and someone says "it was delicious". To understand "it", you glance around the conversation: was someone talking about cake, the weather, a film? You pay attention to the earlier words that help explain "it" and mostly ignore the rest.
Self-attention gives every word that ability. Each word asks a question ("what am I referring to?"), every other word advertises what it's about, and each word blends in information from the ones that answer its question best.
The problem it solves
Words get their meaning from other words. "Bank" means something different in "river bank" and "bank account". "It" means nothing until you know what it refers to, which may be many words earlier.
Older models read text one word at a time and struggled to remember things far back. They were also slow to train because each step waited for the previous one. Attention lets every word connect to every other word directly, in parallel. It's the core idea behind every modern LLM.
How it works, step by step
Take "The animal was tired". Each token starts as a vector (previous lesson).
- Make three vectors per token by multiplying it with three learned matrices. These are its query, key and value:
- Query (Q): "what am I looking for?"
- Key (K): "what do I contain?" (like a label on a box)
- Value (V): "what will I hand over if you pick me?" (the box contents)
- Score every pair. For each token, compare its query with every key using a dot product. A high score means "you're relevant to me".
- Scale the scores by dividing by √(vector size), which keeps the numbers in a range where softmax behaves well.
- Apply the causal mask. In a language model, a token may only look at itself and earlier tokens; later ones are blocked (their scores set to minus infinity).
- Softmax each row so the scores become attention weights that add up to 1 (softmax).
- Blend the values. Each token's output is the weighted sum of the value vectors: mostly from the tokens it attends to most.
- Do it many times in parallel. Each attention head has its own Q, K, V matrices and learns different relationships: one tracks which noun a pronoun refers to, another watches nearby words, another follows syntax. Their outputs are joined together.
The formula you'll see in papers is exactly these steps: Attention(Q, K, V) = softmax(Q·Kᵀ / √d) · V.
Why long contexts are expensive
Every token compares itself with every earlier token, so the work grows with the square of the length: 1,000 tokens means about 500,000 comparisons per head per layer; 100,000 tokens means about 5 billion. Providers use tricks (efficient attention kernels, caching, KV cache) but long prompts still cost more time and money.
See it in code
- Install
pip install numpy- Needs
- Python 3.10+
Scaled dot-product attention with a causal mask, in numpy. The weight matrices are random here (a trained model learns them), so which word attends to which is just an illustration.
import numpy as np
tokens = ["The", "animal", "was", "tired"]
rng = np.random.default_rng(1)
d = 8
X = rng.normal(size=(4, d)) # one vector per token (from embeddings)
# Learned weight matrices (random here) turn each token into a query, key and value.
Wq, Wk, Wv = (rng.normal(size=(d, d)) for _ in range(3))
Q, K, V = X @ Wq, X @ Wk, X @ Wv
scores = Q @ K.T / np.sqrt(d) # how well each query matches each key
mask = np.triu(np.ones((4, 4)), k=1) * -1e9 # causal mask: no peeking at later tokens
scores = scores + mask
weights = np.exp(scores - scores.max(axis=1, keepdims=True))
weights = weights / weights.sum(axis=1, keepdims=True) # softmax per row
out = weights @ V # each token: weighted mix of values
print("attention weights (row = token looking, column = token looked at):")
for t, row in zip(tokens, weights):
print(f"{t:>7}", " ".join(f"{w:.2f}" for w in row))
print("output shape:", out.shape, "(same as input: 4 tokens x 8 numbers)")attention weights (row = token looking, column = token looked at):
The 1.00 0.00 0.00 0.00
animal 0.68 0.32 0.00 0.00
was 0.13 0.81 0.06 0.00
tired 0.00 1.00 0.00 0.00
output shape: (4, 8) (same as input: 4 tokens x 8 numbers)Read it row by row: everything above the diagonal is 0 (the mask), and each row sums to 1. Here "tired" puts all its weight on "animal", the kind of link a trained model learns on purpose.
A real example
Example
Contract review: linking a clause to a definition 30 pages earlier
A legal assistant reads a 40-page contract and is asked "can the supplier end the deal early?". The answer depends on a clause on page 31 that uses the term "Material Breach", which is defined on page 2.
Attention lets the tokens of "Material Breach" on page 31 attend directly to the definition on page 2, however far apart they are, as long as both are inside the context window. Older sequential models often lost information over that distance.
The easy way
Start here
Hold on to the party picture: every word asks a question, every word advertises an answer, and each word blends in the words that answer it best. Then read the snippet's output row by row; that table is attention.
When you grow: the production setup
In real models, attention runs with dozens of heads in each of dozens of layers, using optimised GPU kernels (FlashAttention) and variants that share keys and values across heads (grouped-query attention) to save memory. Those details drive the memory and speed numbers you'll meet in AI system design.
Trade-offs
What you gain
- Any token can use information from any earlier token, however far back.
- All positions are computed in parallel, which made training huge models practical.
- Multiple heads capture many kinds of relationships at once.
What it costs you
- Cost grows with the square of the sequence length.
- Memory for keys and values grows with every token (the KV cache).
- Attention patterns are hard to interpret reliably, even if you can visualise them.
When not to use it
- You don't choose attention as an app builder; but knowing its cost tells you to avoid stuffing huge irrelevant context into prompts.
Common mistakes
- Thinking attention means the model "focuses" like a human, with understanding. It's weighted averaging learned from data.
- Believing more context is free. Long prompts are slower and more expensive, and relevant facts can get lost in the middle.
- Reading too much into a single attention heat-map as an explanation of a decision.
- Forgetting the causal mask exists, then being confused why models generate strictly left to right.
Interview questions
Try answering out loud first, then open the model answer.
Explain self-attention.
Model answer
Each token's vector is projected into a query, key and value with learned matrices. For each token, its query is dot-producted with the keys of all (allowed) tokens, scaled by √d_k, masked for causality, and softmaxed into weights; the output is the weighted sum of value vectors. This lets each token gather context from the most relevant other tokens. Multi-head attention runs several of these in parallel with separate projections and concatenates the results.
Why divide by √d in attention?
What's the computational complexity of self-attention and why does it matter?
Model answer
O(n²·d) time and O(n²) memory for the attention matrix per head and layer, where n is sequence length. It makes long contexts expensive and slow, motivating FlashAttention (memory-efficient exact attention), KV caching during generation, grouped-query attention, sparse or sliding-window attention, and system-level choices like retrieving only relevant context instead of sending everything.
Quiz
Pick an answer to see whether it’s right and why.
0 of 6 answered
Key terms
- Self-attention
- The transformer operation where each token compares itself with other tokens and takes in information from the most relevant ones.
- Query, key and value
- The three vectors each token makes in attention. Queries are compared with keys to get weights; the output is a weighted mix of values.
- Attention head
- One independent attention computation. Models run many heads in parallel so each can learn a different kind of relationship.
- Causal mask
- A rule in attention that lets each token look only at itself and earlier tokens, matching left-to-right generation.
- Context window
- The maximum number of tokens a model can consider in one request, counting both the input and the generated output.
Before and after
Progress is saved in this browser only.