Attention and transformers
Let every position look at every other with attention, and assemble the transformer block behind modern AI.
- Compute scaled dot-product attention from queries, keys and values
- Apply causal masks and explain multi-head attention and positional encodings
- Describe a transformer block: attention, feed-forward, residuals and layer norm
“Neuro chased the ball because it kept bouncing.” To understand it, you look back at ball. Attention lets every position in a sequence look at every other position directly, and decide how much each one matters.
Each token’s vector is projected three ways:
- a query - what am I looking for?
- a key - what do I contain?
- a value - what do I pass on if you attend to me?
A token’s attention to another is the dot product of its query with the other’s key: similar means relevant. Softmax turns those scores into weights, and the output is the weighted average of the values:
Dividing by (the key size) keeps the dot products from growing with the dimension, which would push softmax into flat, gradient-free regions.
Try it
Where does each word look?
Click a word to see its attention over the sentence. Click it - which word does it attend to most? Then raise the temperature to see the attention spread out.
Click a word to see where it looks
“it” pays the most attention to “ball” (66%). The weights add up to 100%.
- Neuro8.9%
- chased2.0%
- the2.0%
- ball65.6%
- because5.4%
- it8.9%
- kept5.4%
- bouncing2.0%
1import numpy as np
2
3def softmax(z):
4 e = np.exp(z - z.max(axis=-1, keepdims=True))
5 return e / e.sum(axis=-1, keepdims=True)
6
7def attention(Q, K, V, causal=False):
8 scores = Q @ K.T / np.sqrt(K.shape[1]) # (seq, seq)
9 if causal:
10 future = np.triu(np.ones_like(scores), k=1).astype(bool)
11 scores = np.where(future, -np.inf, scores)
12 weights = softmax(scores)
13 return weights @ V, weights
14
15rng = np.random.default_rng(0)
16X = rng.normal(size=(4, 8)) # 4 tokens, 8 features each
17W_q, W_k, W_v = (rng.normal(size=(8, 8)) for _ in range(3))
18out, weights = attention(X @ W_q, X @ W_k, X @ W_v, causal=True)
19print(out.shape)
20print(np.round(weights, 2))(4, 8) [[1. 0. 0. 0. ] [1. 0. 0. 0. ] [0.17 0.36 0.47 0. ] [0.05 0.35 0.04 0.56]]
The causal mask sets scores for future positions to −∞ before the softmax, so each token can only attend to itself and earlier tokens. Language models need it: when predicting the next word, peeking ahead would be cheating. (The upper triangle of the weights is all zeros.)
Multi-head attention runs several attentions in parallel with different projections - one head might track grammar, another which noun a pronoun refers to - and concatenates the results. And because attention treats the input as a set, transformers add positional encodings (fixed sine waves or learned vectors) to the token embeddings so word order isn’t lost.
The transformer block
The 2017 paper “Attention Is All You Need” built a whole architecture from attention. A transformer block is:
x = x + MultiHeadAttention(LayerNorm(x)) # mix information across positions
x = x + FeedForward(LayerNorm(x)) # process each position: Linear → GELU → LinearStack dozens of these. Everything from this track is in there: linear layers, a GELU activation, residual connections, layer normalization, softmax. A language model like Claude is a large stack of these blocks trained with cross-entropy to predict the next token - with AdamW, warmup and gradient clipping. You now know every piece.
Key takeaways
Attention: softmax(QKᵀ/√dₖ)V - each position takes a weighted average of all values, weighted by query-key similarity.
A causal mask hides future positions; multi-head attention runs several attentions in parallel.
Positional encodings give attention a sense of order.
A transformer block = attention + feed-forward, each with layer norm and a residual connection.
Lesson quiz
7 questions · pass with 5 correct · up to 50 XP
Passing this quiz completes the lesson and keeps your streak going. Questions you miss come back in review sessions later.
Practice: write Python
Write Python in the editor and run it against sample inputs. Python runs locally in your browser using a WebAssembly runtime.
Causal self-attention
The input is tokens dim seed causal (causal is 1 or 0). The starter draws X and the projection weights from np.random.default_rng(seed). Implement scaled dot-product attention with an optional causal mask, then print the attention weights (each row to 2 decimals) and the sum of each output row to 3 decimals.
- Causal
- Not causal
Python runs in a sandboxed browser worker with a 60 second time limit. Its runtime loads from the Pyodide CDN; your code stays in this browser.
Sinusoidal positional encodings
The original transformer encodes position in a vector of size with sines and cosines of different frequencies:
The input is positions d. Print the encoding for positions 0 to positions−1, one row each, values to 3 decimals.
- 3 positions, d = 4
Python runs in a sandboxed browser worker with a 60 second time limit. Its runtime loads from the Pyodide CDN; your code stays in this browser.
Questions about this lesson
Stuck? Ask. Figured something out? Share it. Explaining is one of the best ways to learn.
Loading posts…