Transformers & Self-Attention: The Architecture That Ate Deep Learning
Canonical lesson. This is the single source of truth for self-attention used by both the Deep Learning track and the NLP & Transformers track (Track 5). If you arrived here from a Track 5 prerequisite link (e.g., from Multi-Head Attention, Positional Encoding, or BERT), you are in the right place — work through this lesson, then return to the NLP-track sequel.
LSTMs read a sentence one word at a time, accumulating memory in a hidden state. Smart, but sequential — you cannot compute step 50 until you have finished steps 1 through 49. So in 2017 eight Google Brain researchers asked a simple, heretical question: what if we throw away the recurrence entirely and let every position look at every other position in one parallel shot? They called the paper "Attention Is All You Need." It worked so well that within five years it had replaced RNNs in NLP, conquered vision (ViT), folded proteins (AlphaFold's Evoformer), and produced ChatGPT, Claude, and Gemini. This lesson is the on-ramp to that architecture from a deep-learning angle — we will earn every piece of the equation, then put them together into a working transformer block in numpy.
Learning Objectives
After this lesson, you will be able to:
Explain why RNNs and LSTMs hit a ceiling — sequential compute, fading long-range memory, no parallel training — and why dropping recurrence was the obvious next move once attention existed
Derive the scaled-dot-product attention formula Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V from first principles, and explain what each of Query, Key, and Value represents semantically
Justify the sqrt(d_k) scaling factor by tracing what happens to softmax when dot products grow with dimension, and the causal mask by tracing what GPT must NOT see during training
Compose a full transformer block out of multi-head self-attention, residual connections, LayerNorm, and a position-wise feed-forward network, and distinguish pre-norm from post-norm
Tell apart the three transformer families — encoder-only (BERT), decoder-only (GPT, Claude, LLaMA), encoder-decoder (T5) — and pick the right one for a given task
Reason about the O(N^2) cost of attention in sequence length and name the modern fixes (FlashAttention, sliding window, Mamba) without diving into them
The transformer is not magic. It is three ideas welded together: (1) self-attention as a learned soft lookup over the sequence, (2) residual streams that make depth cheap, and (3) a feed-forward sublayer that does the actual "thinking" in each block. Once you see those three pieces clearly, every modern LLM is just "what if we did that 96 times and trained on the entire internet?"
Active Recall
Before we tear down RNNs and put up transformers, recall from the LSTM lesson: (a) what specific problem did the LSTM cell-state highway solve, and (b) what was Bahdanau attention's role in the seq2seq pipeline? Write both before reading on. Why this matters: every architectural choice below — QKV, masking, residual streams — is a direct response to a limitation you already know but may not have on the tip of your tongue.
Write your answer in your own words — don't look back at the lesson. This is the most effective way to remember what you just learned.
Recall where we left off in the previous lesson. The LSTM solved the vanishing-gradient problem with a cell-state highway, and Bahdanau attention solved the fixed-context-vector bottleneck by letting the decoder look back at every encoder state. But seq2seq with attention still had two fundamental problems that scaling could not fix:
Sequential compute. To compute hidden state h_50, you must first compute h_1, h_2, ..., h_49 in order. There is no way to parallelize this in the time dimension. A modern GPU has thousands of cores sitting idle while the LSTM crawls forward one token at a time.
Effective context still degrades. Even with the cell-state highway, an LSTM's memory is a fixed-size vector that gets overwritten at every step. Information from 500 tokens ago has been multiplied by 500 forget gates, each slightly less than 1. The signal is technically present but practically faded. Long-document understanding (legal contracts, code repositories, scientific papers) demands that every token can reach every other token equally easily.
Bahdanau already showed us that letting the decoder attend to every encoder state worked beautifully. So Vaswani and his co-authors asked the obvious follow-up: what if attention is enough on its own? What if we just drop the recurrence and let attention do all the work — within a sequence (self-attention), across sequences (cross-attention), and even from a token to its own past (causal self-attention)?
The seq2seq attention you saw in the LSTM lesson was cross-attention — the decoder (one sequence) attended to the encoder's hidden states (a different sequence). The transformer's primary innovation is to apply the same operation within a single sequence: every token attends to every other token in the same input. This is self-attention.
Operation
Queries come from
Keys/Values come from
Used in
Self-attention
Sequence A
Sequence A
Encoder layers, decoder self-attention sublayer
Cross-attention
Sequence A (e.g., decoder)
Sequence B (e.g., encoder)
Encoder-decoder bridge (T5, original transformer)
Causal self-attention
Sequence A
Sequence A, but masked so position i sees only j <= i
Decoder-only models (GPT, Claude, LLaMA)
Self-attention is what makes a transformer feel "contextual": the same word ("bank") gets different output representations depending on its neighbors ("river bank" vs "investment bank") because its query vector matches different keys in each context. RNNs eventually achieved this through their hidden state too, but only after sequentially processing the whole prefix. Self-attention does it in one shot, in parallel, with O(1) hops between any two positions.
Mathematically, every token's embedding x (a vector of dimension d_model) gets projected into three vectors:
Q=XWQ,K=XWK,V=XWV
Why bother with three separate projections instead of just dotting embeddings with each other directly? Because the same word plays very different roles depending on whether it is asking for context or offering it. The word "bank" looking for context (Q) needs a different representation than "bank" offering itself as context (K). The split lets each role be optimized independently during training.
Loading visualization...
The visualization above shows attention scores between every pair of tokens. Each cell (i, j) is how much token i attends to token j. Notice three things: (1) the diagonal is usually bright — tokens often attend to themselves; (2) some columns light up across many rows — those are "popular" tokens that many positions want context from (often nouns or verbs); (3) the pattern is asymmetric — attention from "it" to "animal" is not the same as attention from "animal" to "it".
Here is the equation that does the actual work — the one you will see written on every transformer slide, every system-design whiteboard, every blog post about LLMs:
Attention(Q,K,V)=softmax(dkQK⊤)V
Let us unpack the four operations.
Step 1 — QK^T: the matrix multiplication produces an N-by-N matrix of compatibility scores. Entry (i, j) is Q_i · K_j — the dot product of token i's query with token j's key. A large positive value means "these vectors point in the same direction in the learned space" → "token j is highly relevant to whatever token i is looking for."
Step 2 — scale by sqrt(d_k): without this scaling, dot products grow as d_k grows (random vectors of dimension d have expected dot-product magnitude proportional to sqrt(d)). Large logits push softmax into saturation — one entry near 1, the rest near 0 — and gradients through that softmax collapse to near zero. Dividing by sqrt(d_k) keeps the variance of the scores roughly constant regardless of d_k.
Step 3 — row-wise softmax: softmax turns the N scores in row i into a probability distribution over the N positions. Row i now says "token i puts 30% of its attention on token j=2, 15% on j=7, 5% on j=0, ..." and the row sums to 1.
Step 4 — multiply by V: each output row is a convex combination of the value vectors, weighted by the softmax probabilities. The output for token i is sum_j alpha_{i,j} * V_j. That is the new representation of token i — a mix of relevant value vectors from across the entire sequence.
What Do You Think?
You have a sequence of length 1000. You double it to length 2000. By what factor do the memory and compute requirements of one self-attention layer increase?
One attention operation gives the model one lookup pattern per layer. But a sentence usually needs many simultaneous lookups: who is the subject, who is the verb, what does the pronoun refer to, what is the tense, where does this clause start. The transformer's solution is dead simple: just run several attention operations in parallel with different W^Q, W^K, W^V matrices, then concatenate the results.
Why do heads end up specializing? Because they are initialized with different random weights and trained jointly with the rest of the network. Stochastic gradient descent breaks symmetry — each head finds a different niche in the loss landscape. Interpretability work has found heads that specialize in:
Positional heads. Attend mostly to immediate neighbors (left or right).
Syntactic heads. Attend from verbs to their subjects, or from pronouns to their referents.
Induction heads. Attend from "X" to the position right after a previous occurrence of "X" — the mechanism behind in-context learning in LLMs.
Punctuation-tracking heads. Attend to matching parentheses, quotes, or section boundaries.
Modern LLMs use 8–96 heads per layer. The original transformer used 8. Claude Sonnet and GPT-4 likely use 64+. The exact count is a hyperparameter; you trade head count for per-head dimension at fixed compute budget.
Quick check
Why does multi-head attention with H heads of dimension d_model/H typically outperform a single attention head of dimension d_model, even though the parameter count and compute cost are roughly equal?
#Positional Encoding: Attention Is Permutation-Invariant
There is a subtle issue with self-attention: the equation softmax(QK^T / sqrt(d_k)) V is permutation-equivariant — if you shuffle the rows of X, the output rows shuffle the same way, but otherwise nothing changes. The attention operation has no idea what order the tokens are in.
This is a disaster for language. "Dog bites man" and "Man bites dog" have identical token bags but very different meanings. Without injecting position information somewhere, a transformer would treat them as the same input.
What Do You Think?
If we removed positional encoding from a transformer language model entirely, what would happen to its ability to distinguish 'dog bites man' from 'man bites dog'?
The original transformer added a fixed sinusoidal pattern to each position's embedding:
The visualization shows the sinusoidal pattern as a heatmap — each row is a position, each column is an embedding dimension. Low-index columns oscillate fast (wavelength ~ 2pi) and high-index columns oscillate slowly (wavelength ~ 100002*pi). The slow-frequency dimensions give the model a coarse "where am I in the document" signal, while the fast-frequency dimensions encode fine-grained local position.
Modern alternatives:
Learned absolute (BERT, GPT-2): just learn an embedding per position index. Simple. Fails to extrapolate beyond training length.
RoPE — Rotary Position Embedding (LLaMA, GPT-J, Mistral, modern Claude): rotates Q and K vectors by an angle proportional to position. Has nice properties (relative-position-aware, extrapolates better, easy to extend context). The current default in open-source LLMs.
ALiBi (Press et al. 2022): no PE at all — instead, add a linear bias to attention scores that penalizes distant tokens. Very strong extrapolation properties.
The NLP track goes deeper on position encoding. The DL takeaway: attention is permutation-invariant, so position must be injected, and exactly how you inject it matters for long-context behavior.
#Rotary Position Embedding (RoPE): the actual math
RoPE is named in the previous section but the derivation is worth seeing because it explains every "long-context" trick in the modern LLM stack — NTK scaling, YaRN, LongRoPE — none of which make sense as black-box knobs.
The core idea: instead of adding a position vector to the token embedding (sinusoidal / learned absolute), RoPE rotates the query and key vectors in-place by an angle proportional to position. Take a 2D pair of features (x_1, x_2) from inside Q (or K) at position m. Apply the standard 2D rotation matrix by angle m * theta:
Apply RoPE to Q and K (never to V — only the score matters). Now do the dot product. The key property of rotation matrices is R(a)^T R(b) = R(b - a), so the inner product between a query at position m and a key at position n reduces to:
⟨RoPE(q,m),RoPE(k,n)⟩=qTR((n−m)θ)k
This is why RoPE generalizes to lengths the model never saw at training. Sinusoidal-PE-plus-absolute-attention would have to interpolate the unseen positions into the input embedding; RoPE never embeds an absolute position at all — only the relative dependence is exposed to the model, and the relative dependence is the same at sequence position 1000 as at position 100,000.
NTK-aware scaling for context extension. A model trained on context length 4K with base theta_i = 10000^(-2i/d) saturates its high-frequency RoPE channels once positions exceed 4K — they wrap past one full rotation and the model has never seen the wrap-around. NTK-aware scaling rescales the base so that low-frequency channels are stretched (carrying the long-range signal) while high-frequency channels are barely touched (preserving local resolution):
θi′=θi⋅s−2i/(d−2),s=LtrainLnew
#Grouped-Query Attention (GQA): the inference-cost optimization
Multi-head attention as defined above uses H independent heads, each with its own Q, K, V projections. For H = 64 heads of dimension d_h = 128, that is 64 * 3 * d_model * d_h parameters and — more importantly — a KV cache of size 2 * H * d_h = 16384 floats per token per layer. At 80 layers and 100K context, the KV cache alone eats hundreds of gigabytes. Inference-time bandwidth becomes the bottleneck.
The fix: share K and V projections across groups of heads. Multi-Query Attention (MQA, Shazeer 2019) goes to the extreme — all H query heads share one K head and one V head. Grouped-Query Attention (GQA, Ainslie et al. 2023) is the middle ground: H/G KV heads, where each KV head is shared by G query heads.
A 70B-param decoder uses 64 attention heads with d_h = 128 and d_model = 8192. The team switches from full MHA to GQA with G = 8. What happens to the inference-time KV cache size and the model's quality?
#Toy Implementation: Scaled-Dot-Product Attention from Scratch
Time to make this concrete. The MathPlayground below implements Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V end-to-end in numpy on a 5-word toy sequence, then visualizes the attention weights as a heatmap. Run it. Modify the input embeddings. Watch the attention pattern change.
Loading visualization...
What did the heatmap show? With random weights you should see roughly uniform attention with a few "hot" cells. In a trained transformer the same plot lights up in semantically meaningful ways — pronouns point at their antecedents, verbs point at their subjects, end-of-sentence tokens broadcast over the whole sequence. That is what AttentionHeatmapViz visualizes from a real trained model:
Loading visualization...
The heatmap above shows learned attention patterns from a trained transformer on the sentence "The animal didn't cross the street because it was too tired." Look at the row labeled "it" — you should see a bright cell on the column "animal", not on "street". The model has learned (through pure self-supervision on text) that "it" co-refers with "animal". No one annotated coreference for it. The attention pattern fell out of the language-modeling loss.
So far we have described attention where every token sees every other token. That works for encoders (like BERT) where you have the full input available. But for decoders that generate one token at a time (like GPT), there is a problem: at training time you want to compute the loss on the whole sequence in one parallel forward pass. But during generation, when predicting token t+1, the model should only have access to tokens 1..t — not the future tokens 1..N. If we let the model peek at future tokens during training, it would just learn to copy them at test time.
The fix is a causal mask: set all attention scores above the diagonal to negative infinity before the softmax. After softmax, those positions become exactly zero — the model cannot attend to the future.
Self-attention alone is not enough to make a transformer. The full transformer block wraps the attention sublayer with three more pieces that turn it from a single operation into a reusable, stackable unit:
Residual connection (a.k.a. skip connection): the input to each sublayer is added back to the sublayer's output. This is the same idea as ResNets — it gives gradients a direct path back to earlier layers and lets you stack many blocks (32, 64, 96+) without vanishing gradients.
LayerNorm: normalizes each token's representation across the feature dimension. Stabilizes training; especially important for deep transformer stacks.
Feed-forward network (FFN): a two-layer MLP applied independently to each token (no token-mixing — that already happened in attention). Usually expands to 4 * d_model in the hidden layer, then projects back. This is where most of the "thinking" happens — attention moves information between positions, the FFN refines what each position contains.
The block, in modern pre-norm form:
def transformer_block(x):
# Pre-norm: normalize first, then sublayer, then add residual
x = x + multi_head_self_attention(layer_norm(x)) # sublayer 1
x = x + ffn(layer_norm(x)) # sublayer 2
return x
def ffn(x):
# Position-wise FFN: same weights applied to every token independently
return W_2 @ gelu(W_1 @ x + b_1) + b_2
That is it. The entire transformer is a stack of these blocks — 6 in the original paper, 12 in GPT-1, 24 in BERT-large, 96 in GPT-3, 80+ in GPT-4-scale models. Each block has the same shape; only the learned weights differ.
Loading visualization...
The architecture diagram above shows the full data flow through one block: embeddings + positional encoding → multi-head self-attention → residual + LayerNorm → FFN → residual + LayerNorm → output. The decoder variant inserts an additional cross-attention sublayer between self-attention and FFN to attend to encoder outputs.
#Pre-norm vs Post-norm: A Tiny Reordering with Huge Consequences
The original 2017 transformer used post-norm: x_out = LayerNorm(x + Sublayer(x)). Modern LLMs (GPT-2 onward, LLaMA, Claude, every recent transformer) use pre-norm: x_out = x + Sublayer(LayerNorm(x)).
The difference looks trivial — normalize before vs after the sublayer — but it dramatically changes training stability.
Quick check
Pre-norm transformers (LayerNorm before each sublayer, residual added at the end) train more stably at depth than post-norm transformers. Why?
The second MathPlayground builds the full transformer encoder block — multi-head self-attention with sqrt(d_k) scaling, residual + LayerNorm, then an FFN with GELU, then another residual + LayerNorm. Then we run a toy sequence-classification experiment: the network learns to classify whether a length-6 sequence of random tokens has more 1s than 0s. Watch the loss drop as the (admittedly small) transformer learns to count.
Loading visualization...
A real transformer would use autograd (PyTorch's torch.autograd or JAX's grad), not the crude SPSA estimator above — but the architecture is identical. Replace the gradient estimator and you have a working transformer encoder, in a few hundred lines of numpy, that you understand top to bottom.
Encoder is bidirectional; decoder is causal + cross-attends to encoder
Original transformer (2017), T5, BART, Whisper, AlphaFold's Evoformer (in spirit)
Translation, summarization, speech-to-text, any task with a clean input → output mapping
The encoder-decoder is the most general form and was the original 2017 design. Decoder-only models eventually won the LLM race because: (1) you only need to maintain one set of weights, simpler training; (2) decoder-only handles "input + output" by concatenating them and using prompt engineering, which turns out to be insanely flexible; (3) scaling laws favor a single deep stack over a split encoder-decoder at fixed parameter count.
Encoder-only models still dominate embedding tasks (semantic search, RAG retrieval, classification) because bidirectional context produces much stronger fixed representations than causal models.
Self-attention has one well-known weakness: its compute and memory both scale as O(N^2) in sequence length. The QK^T matrix is N-by-N; for a 100k-token context, that is 10 billion entries per head per layer. At 32 layers and 32 heads, you are looking at 10 trillion attention entries to materialize during the forward pass — completely infeasible without engineering tricks.
The modern fixes (covered in depth in the NLP track):
FlashAttention (Dao et al. 2022, 2023) — still O(N^2) in compute, but tiled into GPU SRAM so you never materialize the full N-by-N matrix. 2-4x speedup, 10x+ memory reduction. Now the default attention kernel in PyTorch.
Sliding-window / local attention (Mistral, Longformer) — each token only attends to a window of W neighbors. O(N * W) instead of O(N^2). Cheap but loses long-range information; usually combined with a few "global" attention heads.
Linear / Performer / Linformer. Approximate the softmax attention with a linear-cost surrogate. Faster but consistently slightly worse quality.
Mamba and state-space models (Gu & Dao 2023) — an entirely different architecture with O(N) inference cost, competitive quality. The first serious challenger to the transformer in years.
Long-context LLMs (1M+ tokens) are a major frontier; expect this list to keep growing.
Stepping back: why did this one architecture eat all of deep learning?
Parallel training. Every position is computed in one shot, every layer is the same shape. GPUs ride this beautifully. RNNs cannot be parallelized in time; CNNs are local. Transformers are the first architecture that lets you spend a million dollars of compute on one training run productively.
Long-range dependencies are O(1) hops. Token 1 can attend directly to token 100,000 in one layer. In an LSTM that would require 100,000 sequential steps of memory propagation. In a CNN with kernel size 3, that would require log_3(100,000) ≈ 11 layers.
Scales beautifully. Empirically, loss falls as a smooth power law in compute, data, and parameters (Kaplan 2020, Chinchilla / Hoffmann 2022) — but only for transformers, at least so far. Other architectures plateau.
Universal. The same block works for text (GPT, BERT, T5), images (ViT), audio (Whisper), proteins (AlphaFold Evoformer, ESM-2), code (Codex), video (Sora). One architecture, many modalities. This was not obvious in 2017.
Interpretability surface. Attention weights are explicitly visible — you can stare at them, find heads that specialize in coreference or induction, and start to reverse-engineer what the model "does." Anthropic's mechanistic interpretability program lives on top of this fact.
Self-attention is a learned, soft, parallel lookup over the sequence. Each token produces a Query (what I want), Key (what I have), and Value (what to bring back). The score softmax(QK^T / sqrt(d_k)) is the lookup; the multiplication by V is the retrieval. The whole operation is differentiable end-to-end.
The sqrt(d_k) scaling is not optional. Without it, dot-product magnitudes scale with sqrt(d_k), softmax saturates, gradients vanish, and training collapses. This is the single most common implementation bug in hand-rolled attention.
Multi-head attention lets different heads specialize. Splitting d_model across H heads gives the model H independent lookup patterns per layer. Interpretability work has found heads that handle coreference, induction, syntax, and punctuation — all emerging from random initialization plus gradient descent.
Positional encoding is mandatory because attention is permutation-invariant. Without injecting position info (sinusoidal, learned, RoPE, or ALiBi), 'dog bites man' and 'man bites dog' would be literally identical inputs to the network.
The transformer block = MHA + residual + LayerNorm + FFN + residual + LayerNorm. Pre-norm (LayerNorm before the sublayer) is the modern default; it gives gradients a clean residual highway and lets you train 96+ layer stacks stably.
Three families, one block. Encoder-only (BERT) uses bidirectional attention. Decoder-only (GPT, Claude, LLaMA) uses causal-masked attention and dominates generative LLMs. Encoder-decoder (T5, original transformer) uses a bidirectional encoder plus a causal decoder with cross-attention.
Attention is O(N^2) in sequence length. The central engineering challenge of long-context LLMs. FlashAttention, sliding-window attention, and Mamba are the major modern responses.
Why is the scaled-dot-product attention formula divided by sqrt(d_k) before the softmax?
Recurrence gave us sequential memory; attention gave us global parallel context. The transformer block — multi-head self-attention plus a position-wise FFN, wrapped in residuals and LayerNorms — is the single architectural unit underneath every modern foundation model. Track 5 (NLP & Transformers) picks up from here: tokenization, BERT, GPT, scaling laws, RLHF, and the modern frontier of long-context LLMs.