Module 3 · Reading 3.3 · read before S16: Transformer I

The PDF posted in Canvas is the assigned version of this reading. It is built from this page, so the text is identical; anything here marked as deep reading is optional and is not in the PDF.

By 2016, the recipe for sequence tasks was settled: recurrent encoder, recurrent decoder, attention connecting them. Attention was the add-on; recurrence was the backbone. The 2017 paper “Attention Is All You Need,” by Vaswani and seven colleagues at Google, inverted that hierarchy with the bluntness its title advertises: delete the recurrence, keep the attention, and add enough structure around it (multiple heads, positional encodings, residual connections, normalization, and per-position feedforward layers) to make a deep stack of attention trainable. The result, the Transformer, translated better than the recurrent state of the art while training in a fraction of the time, because every position in a sequence could now be processed in parallel.

Nearly a decade later, this is the architecture of essentially every large language model, and of much of vision and speech besides. The remarkable part is how little has changed: the block you will build in this section (self-attention, then a feedforward network, each wrapped in a residual connection with layer normalization) is, with minor variations, the repeating unit of GPT-style models today. This section takes the block apart, then reassembles it in PyTorch and verifies the properties everyone claims for it: the shapes, and the causal masking that makes autoregressive training possible.

How to read this

From bolt-on to backbone

The Transformer keeps the encoder–decoder outline of its recurrent predecessors but replaces every recurrent layer with self-attention: attention where the queries, keys, and values all come from the same sequence. Each token constructs a query from its own representation and uses it to gather information from every other token — directly, in one step. The relay that plagued RNNs is gone: any token can read from any other at constant path length, and because no position’s computation waits on another’s, the whole layer is one batch of matrix multiplies. The serial bottleneck disappears with it.

The original model was sized modestly by today’s standards: the base configuration in the paper uses 6 encoder and 6 decoder layers, model width d_model = 512, 8 attention heads, and feedforward width d_ff = 2048, about 65M parameters in total. The paper reports state-of-the-art BLEU scores on English–German and English–French translation at a fraction of the training cost of prior systems. But the architecture’s importance is not that result; it is that the same block, scaled up and stacked deeper, kept working — a property explored in 3.5 Pretraining and Scaling Laws.

People. “Attention Is All You Need” appeared at NeurIPS 2017 with eight authors (Vaswani, Shazeer, Parmar, Uszkoreit, Jones, Gomez, Kaiser, and Polosukhin) and a footnote declaring the listing order random and the contributions equal, an unusual gesture that has become part of the paper’s lore. Several of the authors went on to found prominent AI companies. The paper itself was aimed squarely at machine translation; its abstract makes no claims about general language understanding, let alone the LLM era it enabled.

Multi-head attention

A single attention head computes one weighted average per token — one “kind” of lookup. But a verb may need to find its subject, its object, and its adverbs at once. Multi-head attention runs \(h\) attention functions in parallel, each in its own learned subspace, and concatenates the results.

Concretely, with d_model = 64 and n_heads = 4: project the input three ways with learned matrices to get queries, keys, and values, but split each 64-dimensional projection into 4 chunks of d_head = 16. Each head performs scaled dot-product attention (from the previous section) entirely within its 16-dimensional slice, producing its own weight pattern over tokens. The 4 head outputs are concatenated back to 64 dimensions and passed through a final learned projection \(W_O\). The cost is roughly that of one full-width head (the splitting is free), but the model gets \(h\) independent places to look at once. The paper’s stated motivation is exactly this: averaging with a single head washes out distinct relationships that separate heads can keep separate.

What do trained heads actually do with their independence? Analyses of trained models have repeatedly found a mix of the interpretable and the mysterious. Some heads adopt recognizably syntactic roles: attending to the previous token, to the verb of a clause, to matching brackets or coreferent mentions; others show diffuse patterns that resist naming. The specialization is also unevenly load-bearing: Michel, Levy, and Neubig’s pointedly titled NeurIPS 2019 study “Are Sixteen Heads Really Better than One?” reported that a substantial fraction of heads in trained translation and language models can be pruned after training with little performance loss, though a core of heads is indispensable. The practical moral for your own implementations is to treat n_heads as a mild hyperparameter rather than a sacred constant — what matters most is that d_head = d_model / n_heads stays large enough (typically 64-plus in production models) for each head’s dot products to be informative.

Positional encoding: giving order back

Self-attention is permutation-equivariant: shuffle the input tokens and the outputs shuffle identically, because attention treats its keys as an unordered set. For language this is disqualifying — “dog bites man” and “man bites dog” would get identical (reordered) representations. The Transformer’s fix is to add a position-dependent vector to each token embedding before the first layer, so that identical words at different positions become different vectors.

The original paper uses fixed sinusoidal encodings, with position \(\mathrm{pos}\), dimension pair \(i\) (dimensions \(2i\) and \(2i+1\)), and wavelengths forming a geometric progression from \(2\pi\) up to \(10000 \cdot 2\pi\):

\[PE_{\mathrm{pos},\, 2i} = \sin\bigl(\mathrm{pos} / 10000^{2i / d_{\mathrm{model}}}\bigr), \qquad PE_{\mathrm{pos},\, 2i+1} = \cos\bigl(\mathrm{pos} / 10000^{2i / d_{\mathrm{model}}}\bigr)\]

Each position gets a unique fingerprint; nearby positions get similar fingerprints; and the paper notes that learned position embeddings performed nearly identically in their experiments (GPT-style models chose learned embeddings; many modern models use rotary embeddings instead, derived just below). The claim that attention is order-blind, and that positional encoding fixes it, is testable in a dozen lines, so test it:

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)

# Self-attention alone is permutation-equivariant: shuffle the input
# tokens and the outputs shuffle the same way -- word order is invisible.
# Adding positional encodings breaks the symmetry.

d_model, T = 32, 6

class SelfAttention(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.q = nn.Linear(d, d); self.k = nn.Linear(d, d); self.v = nn.Linear(d, d)
    def forward(self, x):
        s = self.q(x) @ self.k(x).transpose(-2, -1) / math.sqrt(x.shape[-1])
        return F.softmax(s, dim=-1) @ self.v(x)

def sinusoidal_pe(T, d):
    pos = torch.arange(T, dtype=torch.float32).unsqueeze(1)
    i = torch.arange(0, d, 2, dtype=torch.float32)
    angles = pos / (10000 ** (i / d))
    pe = torch.zeros(T, d)
    pe[:, 0::2] = torch.sin(angles)
    pe[:, 1::2] = torch.cos(angles)
    return pe

attn = SelfAttention(d_model)
x = torch.randn(1, T, d_model)          # 6 "tokens"
perm = torch.tensor([3, 1, 5, 0, 2, 4]) # a shuffle of positions

# no positional encoding: output of shuffled input == shuffled output
y = attn(x)
y_perm = attn(x[:, perm, :])
print(f"no PE:  max |attn(shuffled x) - shuffle(attn(x))| = "
      f"{(y_perm - y[:, perm, :]).abs().max().item():.8f}")

# with positional encoding: the symmetry is broken
pe = sinusoidal_pe(T, d_model).unsqueeze(0)
z = attn(x + pe)
z_perm = attn(x[:, perm, :] + pe)
print(f"with PE: max |attn(shuffled x+PE) - shuffle(attn(x+PE))| = "
      f"{(z_perm - z[:, perm, :]).abs().max().item():.8f}")

# The sinusoidal table itself: each row is a unique "timestamp" vector.
pe_t = sinusoidal_pe(T, d_model)
print(f"\nPE row 0, first 6 dims: {[f'{v:.4f}' for v in pe_t[0, :6]]}")
print(f"PE row 3, first 6 dims: {[f'{v:.4f}' for v in pe_t[3, :6]]}")

Running this prints:

no PE:  max |attn(shuffled x) - shuffle(attn(x))| = 0.00000012
with PE: max |attn(shuffled x+PE) - shuffle(attn(x+PE))| = 0.14426041

PE row 0, first 6 dims: ['0.0000', '1.0000', '0.0000', '1.0000', '0.0000', '1.0000']
PE row 3, first 6 dims: ['0.1411', '-0.9900', '0.9933', '-0.1160', '0.8126', '0.5828']

Without positional encoding, shuffling inputs and shuffling outputs agree to floating-point noise (1.2e-7): the layer genuinely cannot see order. With the encoding added, the discrepancy jumps six orders of magnitude — position now matters. The printed rows show why: position 0 and position 3 receive visibly different vectors along the three printed dimension pairs; the slow, high-index pairs differ by less than 0.01.

Three generations of positional encoding

Because position is bolted on rather than built in, the encoding scheme is a genuinely open design choice, and it has evolved through three recognizable generations. Knowing all three matters practically: which one a model uses determines how it behaves past its training length, and it is one of the first lines you will read in any model card.

Fixed sinusoids (the 2017 original) need no parameters and are defined for any position, so nothing in the formula stops you from running past the training length, but the model has never seen those positions’ fingerprints interacting, and quality typically degrades beyond the trained range. Learned position embeddings (used by BERT and the GPT-1/2/3 line) replace the formula with an ordinary trainable table, buying flexibility at two costs: a hard ceiling at the table’s size, and no notion that position 7 relates to position 8 the way 107 relates to 108. Both schemes share a deeper awkwardness: they encode absolute position, while language mostly cares about relative offsets — “the adjective two tokens back” is a useful pattern; “the adjective at absolute position 41” is not.

The third generation makes relative offsets the primitive. Rotary position embedding (RoPE), proposed by Su et al. in 2021, rotates query and key dimension pairs by position-dependent angles. For fixed underlying content vectors \(q\) and \(k\), the identity \(R(m)^\top R(n) = R(n-m)\) makes their rotated dot product depend on the relative offset, not the two absolute positions separately. In a real sequence the content vectors also vary by token, so the full score depends on both content and relative position. The script isolates and tests the positional identity:

import numpy as np

np.random.seed(0)

# Rotary position embedding (RoPE). Instead of ADDING a position vector
# to the embedding, RoPE ROTATES each (even, odd) pair of query/key
# dimensions by an angle proportional to the token's position, with a
# different base frequency per pair. The payoff is exact relative
# position dependence: the dot product of a query at position m with a
# key at position n depends only on m - n, never on m and n themselves.

d = 8  # must be even

def rope(x, pos):
    i = np.arange(d // 2)
    theta = 10000.0 ** (-2.0 * i / d)      # one frequency per dimension pair
    ang = pos * theta
    cos, sin = np.cos(ang), np.sin(ang)
    x2 = x.reshape(d // 2, 2)
    out = np.empty_like(x2)
    out[:, 0] = x2[:, 0] * cos - x2[:, 1] * sin   # 2-D rotation of each pair
    out[:, 1] = x2[:, 0] * sin + x2[:, 1] * cos
    return out.reshape(d)

q = np.random.randn(d)
k = np.random.randn(d)

print("score q.k with RoPE applied, at various (query pos m, key pos n):")
for m, n in [(3, 1), (13, 11), (103, 101), (5, 1), (54, 50), (9, 1)]:
    s = rope(q, m) @ rope(k, n)
    print(f"  m = {m:>3}, n = {n:>3}  (m - n = {m - n}):  score = {s:.10f}")

# Contrast: ADDITIVE sinusoidal encodings do not have this property.
def sin_pe(pos):
    i = np.arange(d // 2)
    ang = pos * 10000.0 ** (-2.0 * i / d)
    pe = np.empty(d)
    pe[0::2] = np.sin(ang); pe[1::2] = np.cos(ang)
    return pe

print("\nscore (q + PE[m]).(k + PE[n]) with additive sinusoidal encodings:")
for m, n in [(3, 1), (13, 11), (103, 101)]:
    s = (q + sin_pe(m)) @ (k + sin_pe(n))
    print(f"  m = {m:>3}, n = {n:>3}  (m - n = {m - n}):  score = {s:.10f}")

Running this prints:

score q.k with RoPE applied, at various (query pos m, key pos n):
  m =   3, n =   1  (m - n = 2):  score = 5.9478093347
  m =  13, n =  11  (m - n = 2):  score = 5.9478093347
  m = 103, n = 101  (m - n = 2):  score = 5.9478093347
  m =   5, n =   1  (m - n = 4):  score = 4.7051479015
  m =  54, n =  50  (m - n = 4):  score = 4.7051479015
  m =   9, n =   1  (m - n = 8):  score = 5.6683275369

score (q + PE[m]).(k + PE[n]) with additive sinusoidal encodings:
  m =   3, n =   1  (m - n = 2):  score = 12.0287056485
  m =  13, n =  11  (m - n = 2):  score = 8.2543303859
  m = 103, n = 101  (m - n = 2):  score = 7.0843151351

The verification is exact to the printed ten decimals: every (query, key) pair at offset 2 produces the same score, near the start of a sequence, in the middle, or a hundred positions out, while offsets 4 and 8 each get their own consistent score. The additive sinusoidal contrast shows the property is not automatic: with the same q and k, the offset-2 score drifts from 12.03 to 7.08 depending on where the pair sits. Note also when RoPE acts: it is applied to queries and keys inside every attention layer (never to values, and nothing is added to the token embedding), so position information arrives exactly where scores are computed instead of being diluted through the residual stream. RoPE is the encoding used by the Llama family, GPT-NeoX, and most current open-weight models; alternatives with a similar relative flavor include ALiBi (Press et al., ICLR 2022), which skips encodings entirely and subtracts a distance-proportional penalty from attention scores.

Historical note. RoPE reached the mainstream by an unusual route: Jianlin Su first developed and published the idea on his Chinese-language research blog before the RoFormer preprint appeared in April 2021, and adoption ran far ahead of formal publication; EleutherAI’s GPT-NeoX models and then Meta’s Llama picked it up while the paper was still unrefereed (it eventually appeared in the journal Neurocomputing in 2024). It is a reminder that by the 2020s, the field’s ideas often propagated through blogs, preprints, and open-source code faster than through any reviewing process.

The block: residuals, LayerNorm, and the FFN

For fixed attention weights, the output is a weighted linear combination of values. But the weights themselves come from data-dependent query–key products and a softmax, so self-attention as a function of its input is nonlinear. The Transformer block wraps it with three further ingredients that make deep stacks trainable and expressive.

Residual connections. Every sublayer computes \(x + \operatorname{sublayer}(x)\), not \(\operatorname{sublayer}(x)\). As you saw with ResNets in Chapter 2, the identity path gives gradients a direct route to early layers; He et al. introduced the trick for vision in 2016 and the Transformer adopted it wholesale.

Layer normalization (Ba, Kiros & Hinton, 2016) standardizes each token’s feature vector, then rescales with learned parameters. The 2017 Transformer is post-norm: \(\operatorname{LayerNorm}(x + \operatorname{sublayer}(x))\). Many later models use pre-norm: \(x + \operatorname{sublayer}(\operatorname{LayerNorm}(x))\), which gives a cleaner identity gradient path and often improves deep-stack stability. Pre-norm is a sensible first implementation, but it does not eliminate the need for learning-rate warmup, careful initialization, residual scaling, or gradient clipping in large recipes; those remain model- and scale-dependent choices.

Historical note. Both normalization choices are younger than they look in hindsight. BatchNorm arrived in 2015 for CNNs, LayerNorm in 2016 partly because BatchNorm fit sequence models poorly, and the Transformer adopted LayerNorm the following year. The pre-norm rearrangement then spread through 2018–2019 codebases (including GPT-2) ahead of the papers that analyzed why it helps — a common pattern in this field: practice first, theory later.

The feedforward network (FFN). After attention mixes information across tokens, a two-layer MLP processes each token independently: Linear(d_model, d_ff), a nonlinearity (ReLU originally; GELU in most modern models), Linear(d_ff, d_model), with d_ff conventionally 4 * d_model. The division of labor is worth internalizing: attention moves information between positions; the FFN transforms it in place. Roughly two-thirds of a standard block’s parameters live in the FFN.

The width ratio deserves a closer look, because it is a convention, not a law. The original paper’s 512 -> 2048 expansion set the 4x default, and GPT-2 and BERT kept it, but nothing in the mathematics requires it; what matters empirically is total capacity, and the expand-then-contract shape that gives the nonlinearity a wide space to work in. Modern models bend the number for a specific reason: many (the Llama family among them) replace the plain two-layer MLP with a gated variant, SwiGLU (Shazeer, 2020), which computes \((W_1 x) \odot \operatorname{swish}(W_{\mathrm{gate}} x)\) and then projects back down (three weight matrices instead of two). To keep the parameter budget comparable to a 4x plain FFN, these models shrink the hidden width to roughly two-thirds of 4 * d_model (you will see numbers like d_ff = 11008 against d_model = 4096, rather than 16384). So when a model card’s FFN width looks “wrong,” check the activation function first. A second modern reading of the FFN is worth planting now: because the first layer’s rows are dotted against the token and the second layer’s columns are summed up weighted by the results, the FFN behaves like a large key–value memory — patterns matched by the first layer retrieve stored directions from the second (an interpretation developed by Geva et al., EMNLP 2021). Much of what a Transformer knows, as opposed to how it routes, appears to live in these layers, consistent with where the parameters are.

A minimal block in PyTorch, verified

The script below assembles the whole thing (multi-head causal self-attention plus FFN, pre-norm residuals) and then verifies the property that makes decoder-style Transformers trainable as language models. A causal (autoregressive) model must not let position \(t\) see positions after \(t\); the mask enforces this by setting future scores to -inf before the softmax. If the mask is right, then editing the token at position 5 must leave outputs at positions 0–4 exactly unchanged. That is a hard, testable claim, so the script tests it, and then removes the mask to show the contrast.

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)

class MultiHeadSelfAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model)   # Q, K, V in one matmul
        self.proj = nn.Linear(d_model, d_model)      # output projection W_O

    def forward(self, x, causal=True):
        B, T, D = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        # (B, T, D) -> (B, n_heads, T, d_head)
        def split(t):
            return t.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        q, k, v = split(q), split(k), split(v)
        scores = q @ k.transpose(-2, -1) / math.sqrt(self.d_head)  # (B, H, T, T)
        if causal:
            mask = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
            scores = scores.masked_fill(mask, float("-inf"))
        attn = F.softmax(scores, dim=-1)
        out = attn @ v                                             # (B, H, T, d_head)
        out = out.transpose(1, 2).contiguous().view(B, T, D)       # concat heads
        return self.proj(out)

class TransformerBlock(nn.Module):
    """Pre-norm block: x + Attn(LN(x)), then x + FFN(LN(x))."""
    def __init__(self, d_model, n_heads, d_ff):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = MultiHeadSelfAttention(d_model, n_heads)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model))

    def forward(self, x, causal=True):
        x = x + self.attn(self.ln1(x), causal=causal)
        x = x + self.ffn(self.ln2(x))
        return x

B, T, d_model, n_heads, d_ff = 2, 8, 64, 4, 256
block = TransformerBlock(d_model, n_heads, d_ff)
x = torch.randn(B, T, d_model)
y = block(x)
n_params = sum(p.numel() for p in block.parameters())
print(f"input shape  {tuple(x.shape)}")
print(f"output shape {tuple(y.shape)}")
print(f"parameters   {n_params}")

# Causal-mask property: change the token at position 5; outputs at
# positions 0-4 must be bit-for-bit unchanged, positions 5-7 may change.
x2 = x.clone()
x2[:, 5, :] = torch.randn(B, d_model)
y2 = block(x2)
diff = (y2 - y).abs().amax(dim=(0, 2))  # max abs change per position
for t in range(T):
    print(f"position {t}: max |change| = {diff[t].item():.6f}")

# Same edit without the causal mask: everything changes.
y_nc, y2_nc = block(x, causal=False), block(x2, causal=False)
diff_nc = (y2_nc - y_nc).abs().amax(dim=(0, 2))
print(f"no mask, positions 0-4 max |change| = {diff_nc[:5].amax().item():.6f}")

Running this prints:

input shape  (2, 8, 64)
output shape (2, 8, 64)
parameters   49984
position 0: max |change| = 0.000000
position 1: max |change| = 0.000000
position 2: max |change| = 0.000000
position 3: max |change| = 0.000000
position 4: max |change| = 0.000000
position 5: max |change| = 3.357778
position 6: max |change| = 0.268559
position 7: max |change| = 0.215717
no mask, positions 0-4 max |change| = 0.205075

The verification passes. Shapes are preserved ((2, 8, 64) in, (2, 8, 64) out), which is what lets blocks stack arbitrarily deep. The parameter count checks against hand arithmetic: attention contributes 64*192 + 192 (QKV) plus 64*64 + 64 (output projection), the FFN contributes 64*256 + 256 plus 256*64 + 64, the two LayerNorms 128 each; total 49,984, and two-thirds of it in the FFN as promised. Most importantly, with the causal mask, changing token 5 changes nothing at positions 0–4 (differences exactly 0.0, not merely small) while positions 5–7 shift; without the mask, past positions move too. This “future edits don’t touch the past” property is precisely why one forward pass over a \(T\)-token sequence yields \(T\) valid next-token training examples at once — the training-time parallelism that made large-scale language modeling economical.

Pitfall. An off-by-one in the causal mask (diagonal=0 instead of diagonal=1 in torch.triu, or masking after the softmax instead of before) produces a model that either cannot see itself or can partially see the future. The second bug is the dangerous one: training loss drops implausibly fast (predicting a token is easy when you can peek at it), and the model is useless at generation time. Run the perturbation test above on any attention implementation you write; HW3’s test suite does exactly this to yours.

Counting a real model, by hand

The 49,984-parameter toy block scales to headline model sizes by nothing more than arithmetic, and doing that arithmetic once, in full, permanently demystifies numbers like “124M parameters.” The target: a model shaped exactly like GPT-2 small — vocabulary 50,257, context 1,024, width d_model = 768, 12 blocks, 12 heads, d_ff = 4 * 768 = 3072. The script tallies every piece by formula, then instantiates real PyTorch modules and counts, demanding exact agreement.

One new ingredient appears at the output end. A language model’s final layer maps d_model = 768 back to 50,257 vocabulary logits, a \(V \times d\) matrix exactly the same shape as the input embedding table. Weight tying (Press & Wolf, EACL 2017) uses one matrix for both jobs: the embedding that turns id 2026 into a vector is also row 2026 of the output projection. The intuition: both matrices pair token identities with directions in representation space, and tying them makes tokens’ input and output representations agree. The original Transformer, GPT-1, and GPT-2 all tied; some larger modern models untie because at their scale the parameter savings stop mattering.

import torch
import torch.nn as nn

torch.manual_seed(0)

# Parameter accounting for a full GPT-2-small-shaped model, twice: once
# by hand with the formulas, once by instantiating real modules and
# counting. The two must agree exactly.

V, T_ctx, d, L, n_heads = 50257, 1024, 768, 12, 12
d_ff = 4 * d

# --- by hand ---
tok_emb = V * d                     # token embedding table
pos_emb = T_ctx * d                 # learned position embeddings
attn = (3 * d * d + 3 * d) + (d * d + d)      # QKV proj + output proj
ffn = (d * d_ff + d_ff) + (d_ff * d + d)      # two linear layers
ln = 2 * (2 * d)                    # two LayerNorms per block (scale + shift)
block = attn + ffn + ln
final_ln = 2 * d
head_untied = V * d                 # output projection to vocab logits

total_tied = tok_emb + pos_emb + L * block + final_ln          # head shares tok_emb
total_untied = total_tied + head_untied

print(f"token embeddings      {tok_emb:>12,}")
print(f"position embeddings   {pos_emb:>12,}")
print(f"one block: attention  {attn:>12,}")
print(f"one block: FFN        {ffn:>12,}")
print(f"one block: LayerNorms {ln:>12,}")
print(f"one block total       {block:>12,}   x {L} layers = {L * block:,}")
print(f"final LayerNorm       {final_ln:>12,}")
print(f"TOTAL (tied head)     {total_tied:>12,}")
print(f"TOTAL (untied head)   {total_untied:>12,}")
print(f"FFN share of block params: {ffn / block:.1%}")

# --- by instantiation ---
class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.qkv = nn.Linear(d, 3 * d)
        self.proj = nn.Linear(d, d)
        self.ffn = nn.Sequential(nn.Linear(d, d_ff), nn.GELU(), nn.Linear(d_ff, d))

class GPT2Small(nn.Module):
    def __init__(self):
        super().__init__()
        self.tok = nn.Embedding(V, d)
        self.pos = nn.Embedding(T_ctx, d)
        self.blocks = nn.ModuleList(Block() for _ in range(L))
        self.ln_f = nn.LayerNorm(d)
        self.head = nn.Linear(d, V, bias=False)
        self.head.weight = self.tok.weight       # weight tying

model = GPT2Small()
counted = sum(p.numel() for p in {id(p): p for p in model.parameters()}.values())
print(f"\ninstantiated and counted (deduplicating tied weights): {counted:,}")
print(f"matches hand arithmetic: {counted == total_tied}")

Running this prints:

token embeddings        38,597,376
position embeddings        786,432
one block: attention     2,362,368
one block: FFN           4,722,432
one block: LayerNorms        3,072
one block total          7,087,872   x 12 layers = 85,054,464
final LayerNorm              1,536
TOTAL (tied head)      124,439,808
TOTAL (untied head)    163,037,184
FFN share of block params: 66.6%

instantiated and counted (deduplicating tied weights): 124,439,808
matches hand arithmetic: True

The total lands at 124.4M (the size reported for GPT-2 small), and every line of the budget teaches something. The embedding table alone is 38.6M parameters, roughly 31% of the model; weight tying is why the total is not 163M, and why the untied-versus-tied distinction matters when you compare model sizes across papers. The twelve blocks contribute 85M, with the promised two-thirds FFN share (66.6%) holding at full width, since both attention and FFN parameters scale as \(d^2\). And notice what is absent: the causal mask, the softmax, the residual connections, and the positional rotations of a RoPE model contribute nothing — most of what makes the architecture work costs no parameters at all. As \(d\) grows, the \(d^2\) block terms outrun the \(V \cdot d\) embedding terms, which is why at multi-billion scale the embedding share dwindles and “parameters” becomes nearly synonymous with “blocks.”

A budget for compute pairs naturally with the budget for parameters: a useful rule of thumb (popularized in the scaling-law literature you will meet in 3.5) is that a forward-plus-backward pass costs about \(6N\) floating-point operations per token for a model with \(N\) parameters: every parameter is touched once forward, twice backward, each touch a multiply-add. Parameter arithmetic like the above is thus the first step of training-cost arithmetic, and doing both by hand is a skill this book will lean on again.

The full architecture, and what comes next

A complete model is now assembly. Map token ids to vectors with an embedding table, add positional encodings, apply \(N\) identical blocks, and finish with a task head: for language modeling, a final LayerNorm and a linear projection to vocabulary logits. The original Transformer had two towers: an encoder whose self-attention is unmasked (every source token sees every other), and a decoder that uses causal self-attention plus cross-attention, where queries come from decoder tokens but keys and values come from encoder outputs — the direct descendant of 2015-style translation attention.

Later work found that you rarely need both towers at once. Keeping only the encoder yields models built for understanding (BERT); keeping only the decoder yields models built for generation (the GPT line); keeping both suits sequence-to-sequence tasks (T5). Which tower you keep, and what objective you train it on, defines the model families of the next section — along with the question this section quietly skipped: what exactly is a “token”?

Going deeper

The 2017 block, updated part by part. “The architecture barely changed” is true at the block diagram level and false in the fine print, and the fine print is a good index of what a decade of scale taught. Normalization: pre-norm placement (as discussed above), and often RMSNorm (Zhang & Sennrich, NeurIPS 2019) instead of LayerNorm; it drops the mean-centering and the bias, keeping only the scale, and does the job with less compute. Activations: ReLU gave way to GELU, then in many models to the gated SwiGLU with its adjusted width. Positions: additive encodings gave way to RoPE at the attention layers. Attention: multi-head became grouped-query in most serving-oriented models to shrink the KV cache (see 3.2’s Going deeper). Biases: many current models delete the bias terms from their linear layers entirely, having found them dispensable at scale. Each change is small; the striking fact is that the composition (attention, FFN, residuals, normalization, repeated) has absorbed them all without structural revision.

Stretching the context after training. The positional-encoding generations story continues past pretraining. Because RoPE is a formula rather than a table, a trained model’s context can be extended by manipulating the formula: position interpolation (Chen et al., 2023) rescales positions so a longer sequence maps into the rotation range the model saw during training, and a brief fine-tune adapts the model to the compressed spacing — one reason context lengths leapt in 2023–2024 without models being retrained from scratch. ALiBi took the other route: build length-extrapolation into the score penalty from the start. The area remains active; when a model advertises a very long context, the interesting question is what mix of encoding choice, interpolation trick, and long-document fine-tuning produced it.

Scaling the FFN sideways: mixture of experts. Since the FFN holds most parameters, it is the natural place to add more, and mixture-of-experts (MoE) blocks do so without adding proportional compute. Replace one FFN with many parallel “expert” FFNs plus a small router that sends each token to one or a few of them: total parameters multiply, but per-token FLOPs stay near constant. The idea predates the Transformer (Shazeer et al. explored it for LSTMs in 2017), was industrialized for Transformers by the Switch Transformer (Fedus, Zoph & Shazeer, JMLR 2022), and appears in several frontier and open-weight models today (Mixtral, among the documented ones). MoE complicates training (routers need load-balancing losses) and serving (all experts must sit in memory), which is why dense models remain common — but it is the clearest example of the block’s division of labor being exploited deliberately: routing stays dense, knowledge storage gets sharded.

Reading the residual stream. Interpretability researchers describe a Transformer not as a pipeline but as a residual stream: a d_model-wide channel per token that every attention head and FFN reads from and writes small updates into (a framing developed in Elhage et al.’s 2021 “A Mathematical Framework for Transformer Circuits”). The view has real analytic payoff: attention-head compositions implementing identifiable algorithms (“induction heads” that complete repeated patterns are the best-documented case). And it explains at a glance why pre-norm stacks are well-behaved: the stream itself is never squashed, only read through normalization. If you pursue mechanistic interpretability later in the course, this vocabulary is the entry point.

Common pitfalls

  • Training loss drops suspiciously fast, generations are garbage. Cause: causal-mask leakage — the model can see the token it is predicting (off-by-one mask, or targets not shifted by one relative to inputs). Fix: run the future-perturbation test; check that inputs are tokens 0..T-1 and targets are tokens 1..T.
  • Deep post-norm stack diverges in the first few hundred steps. Cause: post-norm’s poorly scaled early gradients combined with a full-size learning rate from step 0. Fix: switch to pre-norm, or keep post-norm with a learning-rate warmup; check gradient norms per layer.
  • Head-split reshape silently scrambles features. Symptom: model trains but plateaus far above a reference implementation. Cause: reshaping (B, T, D) to (B, n_heads, T, d_head) with a single view mixes head and time axes; the correct order is view(B, T, H, d_head) then transpose(1, 2). Fix: unit-test that head 0’s output matches a single-head attention run on the corresponding slice.
  • NaNs appear when a padding row is fully masked. Cause: a query whose keys are all -inf makes softmax compute 0/0. Fix: ensure every query keeps at least one visible key (e.g., self), or zero out fully masked rows after the softmax.

Try it yourself

  1. Stack three TransformerBlocks and rerun the perturbation test. Does the “past is untouched” property survive depth? Why must it?
  2. Compute the parameter count of the paper’s base layer (d_model = 512, n_heads = 8, d_ff = 2048) by hand from the formulas above, then verify by instantiating the block and counting.
  3. Convert the block to post-norm (\(\operatorname{LayerNorm}(x + \operatorname{sublayer}(x))\)) and compare the norm of the output of a 12-block stack (random weights, random input) against the pre-norm version. Which grows with depth, and what does that imply for gradients?
  4. Replace the sinusoidal encoding with a learned nn.Embedding(T, d_model) of positions and rerun the permutation experiment. Does the symmetry still break before any training?
  5. In the RoPE script, the relative-position identity was verified for dot products of rotated vectors. Check what rotation does to norms: print ||rope(q, m)|| for several m. What property of rotation matrices explains the result, and why does it matter that the encoding not change vector lengths?
  6. Redo the parameter count for GPT-2 medium (L = 24, d = 1024, context and vocabulary unchanged). Before running, predict the embedding table’s share of the total, and check how the FFN share of block parameters moves. Verify by instantiation.

Further reading

  • Vaswani et al., “Attention Is All You Need,” NeurIPS 2017. arXiv:1706.03762 — the source; short by modern standards, and worth reading with this section’s vocabulary in hand.
  • Ba, Kiros & Hinton, “Layer Normalization,” 2016. arXiv:1607.06450 — why per-token normalization fits sequence models where BatchNorm does not.
  • Xiong et al., “On Layer Normalization in the Transformer Architecture,” ICML 2020. arXiv:2002.04745 — the formal pre-norm versus post-norm analysis, including why post-norm needs warmup.
  • He et al., “Deep Residual Learning for Image Recognition,” CVPR 2016. arXiv:1512.03385 — where the residual connection came from; Chapter 2 covered it for vision.
  • Su et al., “RoFormer: Enhanced Transformer with Rotary Position Embedding,” 2021 (journal version Neurocomputing, 2024). arXiv:2104.09864 — the RoPE derivation whose central identity you verified numerically above.
  • Press, Smith & Lewis, “Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation,” ICLR 2022. arXiv:2108.12409 — ALiBi: position as a score penalty, and the cleanest study of length extrapolation.
  • Press & Wolf, “Using the Output Embedding to Improve Language Models,” EACL 2017. arXiv:1608.05859 — the weight-tying result behind the 38.6M-parameter saving in the counting example.
  • d2l.ai, The Transformer Architecture — the full encoder–decoder assembled and trained on a real translation task.
  • Harvard NLP, The Annotated Transformer — the paper reimplemented line by line; the standard companion for HW3.
  • PyTorch docs: nn.TransformerEncoderLayer — the built-in block, including its norm_first flag for the pre-norm/post-norm choice.

Before class you should be able to

  • Explain why multiple heads exist and compute d_head from d_model and n_heads
  • State what permutation equivariance means for self-attention and how positional encoding breaks it
  • Name the three additions the block makes around attention and state the division of labor between attention and the FFN

For S17 and HW3 (optional now)

  • Explain the causal mask, and predict what the future-perturbation test must show if it is correct
  • Say what pre-norm changes relative to post-norm and why pre-norm trains more stably
  • Count the parameters of a Transformer block from d_model and d_ff
  • Explain weight tying and why it saves \(V \cdot d\) parameters
  • State where roughly two-thirds of a block’s parameters live, and why

Connects to

  • This session is covered by Quiz 3 (Module 3: S13–S22), given at S23.
  • HW3 (Transformer from scratch, out at S18, due at S22): this section is its core reading; you will implement the attention, multi-head, and block code here almost line for line.
  • Midterm grades are reported this week, based on your work so far.

← S15 · Guided Review · All readings · S17 · Transformer II →