Transformers & Self-Attention, Part 2: Rotary Positions, GQA, the Causal Mask and the Full Transformer Block
Canonical lesson, continued. This is the second half of the single source of truth for self-attention used by the Deep Learning track and the Transformers & LLMs track (Track 5). It assumes Part 1, which builds queries, keys and values, the scaled dot-product formula, multi-head attention and sinusoidal positions. If you arrived from a Track 5 link about the transformer block, the causal mask or RoPE, you are in the right place.
Part 1 built the attention operation. One attention layer is still not a transformer. This part first shows how modern models encode position with rotations (RoPE) and shrink the memory cost of many heads (GQA). It then adds the causal mask that lets a model write text one token at a time, and wraps everything in the residual, normalization and feed-forward pieces that get stacked 12 to 96 times. It ends with the three model families and the O(N^2) wall that limits long contexts.
Learning Objectives
After this lesson, you will be able to:
Derive RoPE as a rotation of query and key pairs, and show that the score between two rotated vectors depends only on how far apart their positions are
Explain grouped-query attention and compute how much it shrinks the KV cache
Justify the causal mask by tracing what GPT must NOT see during training, and apply it with minus infinity before the softmax
Compose a full transformer block out of multi-head self-attention, residual connections, LayerNorm, and a position-wise feed-forward network, count its parameters, 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
In plain words: Part 1 made one layer that lets words read each other. Part 2 makes it usable. Position is given by turning each query and key by an angle that grows with the word's place. A mask stops a word from reading the future so the model can be trained to write. A residual path, a normalization step and a small feed-forward network wrap the attention so many copies can be stacked without the signal breaking. The result is the transformer block.
Words you need10 terms
Rotary position embedding. Each query and key is turned by an angle proportional to its position, so the score between two tokens depends on how far apart they are.
The stored keys and values of tokens already generated, so each new token only has to compute its own row.
Several query heads share one set of keys and values, which makes the KV cache several times smaller.
A triangular mask that sets the scores of future positions to minus infinity, so a token can only read itself and what came before.
A path that adds a sublayer's input back to its output, so the signal and the gradient have a direct road through many layers.
Rescales each token's vector to zero mean and unit spread, which keeps training stable in deep stacks.
A small two-layer MLP applied to each token on its own. It refines each token and does not mix tokens.
Placing LayerNorm before each sublayer, so the residual path stays untouched. The modern default.
A bidirectional encoder reads the input and a causal decoder writes the output while cross-attending to it.
Attention scores every pair of tokens, so doubling the sequence length quadruples the score matrix.
#Rotary Position Embedding (RoPE): the actual math
Part 1 named RoPE as a modern alternative to sinusoidal encodings, 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
In the lab below, move the query position m and the key position n and watch the score depend only on their difference.
Loading visualization...
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 in Part 1 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?
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 LLaMA 2 70B. 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.
Every number in a transformer block is a matrix whose shape you can read off, so its size is a short sum. The multi-head derivation in Part 1 already counted the attention part, so what is left is the feed-forward network and the totals.
Pblock≈attention4d2+FFN8d2=12d2
The model-wide total follows by stacking blocks and adding the embedding table.
#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.
Both orderings keep a residual connection, and the next derivation shows what that connection does to gradients.
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.
A Pyodide-backed scratchpad for math lessons.
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.
RoPE rotates queries and keys by an angle that grows with position. The score between two rotated vectors then depends only on their offset, which is what makes relative position and long-context tricks work.
Grouped-query attention shares keys and values across groups of query heads. With 64 query heads and 8 KV heads the KV cache is 8 times smaller at almost no loss in quality.
The causal mask is what lets a transformer be trained in parallel and still generate left to right. Future scores become minus infinity before the softmax, so their weights are exactly zero.
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.
What is the purpose of the causal mask in a decoder-only transformer like GPT?
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 architectural unit underneath most modern foundation models.