Lesson 2 · The Transformer

The Architecture at a Glance

In Lesson 1, each token was converted to a dense vector by looking up a row in the learned embedding matrix E, then a positional embedding P_t was added to mark its place in the sequence. That operation produces the transformer's actual input — one d-dimensional vector per token, stacked into a matrix:

Embeddings → Transformer input (from Lesson 1)x⁽⁰⁾_t = E[id_t] + P_t // token embedding + positional embedding Stack T tokens → X⁽⁰⁾ ∈ ℝ^T×d // this matrix enters the transformer

The transformer's job is to refine X⁽⁰⁾ into progressively richer representations by passing it through L successive blocks. Each block performs two operations: a self-attention layer (which lets every token gather information from every earlier token and itself) and a feedforward network (which processes each position independently). Every block reads from and writes back to these same vectors — the residual stream. Unlike a plain feedforward network — where each layer's output replaces the previous one — the residual stream works by addition: each block contributes a correction to the existing vectors, so the original token representation is not replaced and stays part of the sum at every depth (although later layers can add directions that cancel parts of it).

For GPT-2 small there are 12 such blocks, each operating on the same 768-dimensional vectors that the embedding lookup produced. The model applies those two operations 12 times before producing output logits. This block count is what practitioners mean by the model's depth: a 12-block transformer has depth 12 by the standard convention, or depth 24 if the attention and feedforward sublayers are each counted separately. When researchers refer to "layer N activations" or compare model depths, they almost always mean the block count.

Figure 1 — From Token IDs to Output Logits

Token IDs → E[id_t] + P_t = X⁽⁰⁾ ∈ ℝ^T×d

Embedding lookup + positional encoding (Lesson 1)

↓

× L blocks

Layer Norm

↓

Multi-Head Self-Attention

↕ residual connection (+ x)

Layer Norm

↓

Feedforward Network (MLP)

↕ residual connection (+ x′)

↓

Final Layer Norm

↓

LayerNorm(X^(L)) → unembedding → output logits over vocabulary

The embedding vectors from Lesson 1 are the literal starting point (X⁽⁰⁾). Each transformer block refines them in place — adding information from earlier tokens via attention, then processing each position independently via the feedforward layer. After the last block, a final layer norm is applied and the result is projected through the unembedding matrix to produce next-token probabilities.

One transformer block (pre-norm convention)x′ = x + MHA(LayerNorm(x)) // attention sub-layer (multi-head attention, tab 4) with residual x_out = x′ + FFN(LayerNorm(x′)) // feedforward sub-layer with residual

Notice what stays constant across all 12 blocks: the dimensionality d. Every block receives a T×d matrix and emits a T×d matrix. The embedding lookup fixes d, and the transformer never changes it. The blocks only refine the content of those vectors — not their shape. This is what makes the residual connection possible: each sublayer adds a correction of the same shape to the stream it reads from.

Key intuition: The embedding vectors from Lesson 1 are a first draft — each token represented in isolation, with no awareness of context. The transformer's blocks are a sequence of revisions: each pass lets every token incorporate information from its neighbors, progressively enriching the representation until it is rich enough to predict the next word.

Model scale: where do all the parameters go?

A 1.5B-parameter model like GPT-2 XL sounds abstract. Most of those parameters live in just two places per block: the attention projection matrices W_Q, W_K, W_V, W_O (each d × d, so 4d² per block), and the feedforward weight matrices W_up and W_down (d × 4d and 4d × d, so another 8d² per block). That gives roughly 12d² parameters per block times L blocks — for GPT-2 XL that's 12 × 1,600² × 48 ≈ 1.5B (GPT-2 small: 12 × 768² × 12 ≈ 85M, plus ~39M of embeddings ≈ 124M).

Parameter budget per blockAttention (W_Q, W_K, W_V, W_O): 4 × d² // four d×d matrices Feedforward (W_up, W_down): 8 × d² // d→4d→d, both directions ────────────────────────────────────────── Total per block: 12 × d² GPT-2 small (d=768, L=12): ~85M + embeddings ≈ 124M params GPT-2 XL (d=1600, L=48): ~1.5B params

Neuroscience parallel: The layered structure mirrors the hierarchical organization of sensory cortex — V1, V2, V4, IT in vision; A1, belt, parabelt in audition. Early layers extract local features; later layers build abstract, task-relevant representations. The difference is that transformer layers are identical in structure — the abstraction hierarchy emerges purely from training, not from anatomical specialization.

The Residual Stream

The residual stream is the sequence of d-dimensional vectors that passes through the entire model — one vector per token position. It starts as the sum of the token embedding and the positional embedding, and each subsequent block reads from it and writes an additive correction back to it.

At no point does any layer replace the stream. Attention adds a correction; the feedforward adds another correction. The original token embedding is still part of the sum at every layer — it has been incrementally refined, though later layers can write directions that partly cancel or crowd out what earlier layers wrote.

Figure 2 — The Residual Stream Through L Blocks

The residual stream (central vertical axis) carries the representation through all L blocks, numbered 0 to L−1. Blocks 0 and 1 show the pattern concretely; block l shows the general form (it reads x^(l) and writes x^(l+1)). Each block's self-attention sublayer (blue) and feedforward sublayer (orange) each compute a correction and write it to the stream (⊕). The stream is never replaced — only added to, so x⁽⁰⁾ stays in the sum at every depth (later terms can still cancel parts of it).

Why residuals matter

Without residual connections, gradients vanish as they backpropagate through many layers — the product of many Jacobians approaches zero. Residuals solve this by providing a gradient highway: the loss gradient can flow back directly to early layers without passing through every transformation. This is the same insight behind ResNets in computer vision.

The residual formulation also has a useful interpretive consequence: a sublayer that learns nothing useful can simply output zero, so the block acts as the identity function and leaves the stream unchanged. This means the model can effectively decide how deep to process each token by tuning how much each block contributes.

Common misconception: It is tempting to think of the residual stream as a "working memory" that gets overwritten layer by layer. It isn't replaced wholesale — every layer's contribution stays in the sum, unless a later layer writes a direction that cancels it (which models do learn, for example to remove information they no longer need). Mechanistic interpretability research often exploits this, tracing how specific circuits write information in early layers and read it back in later ones.

The superposition hypothesis

The residual stream has d dimensions — 768 for GPT-2 small. But the model needs to track far more than 768 independent concepts simultaneously. Research on superposition suggests the model packs many more features into these dimensions by encoding them as near-orthogonal directions, accepting small interference between features as a tolerable cost. This is an active area of mechanistic interpretability research.

Neuroscience parallel: The residual stream resembles the concept of a population code in systems neuroscience — a large set of neurons jointly encoding many variables simultaneously, with each neuron responding to a mixture of features (mixed selectivity). Just as a single prefrontal neuron might encode task context, reward, and stimulus identity all at once, individual dimensions of the residual stream encode overlapping, distributed information.

Neuroscience parallel — network dynamics: The residual stream update rule (x^(l+1) = x^(l) + Δ^(l)) is mathematically identical to an Euler step in a discrete dynamical system. This makes the residual stream a state-space trajectory: depth plays the role of time, and each layer applies a vector field that moves the representation toward something more useful for prediction. Neural ODE models (Chen et al., 2018) formalize this connection explicitly, treating network depth as a continuous time axis. The closer neuroscience analog is the ventral visual stream (V1→V2→V4→IT) — a spatial unfolding of representations across a processing hierarchy — rather than temporal oscillatory dynamics, since a transformer has no recurrence across time. The high-dimensional trajectory framing also connects naturally to superposition: just as neural populations traverse trajectories through high-dimensional activity spaces during a task, the residual stream navigates a high-dimensional representation space where many features are packed as near-orthogonal directions.

Self-Attention

Attention is the mechanism that lets every token position look at every earlier position (and itself) and decide how much information to gather from each. Unlike the feedforward layers, which process each position in isolation, attention computes pairwise relationships across the sequence.

Queries, Keys, and Values

Each position in the sequence generates three vectors from its current residual stream state: a Query Q (what information am I looking for?), a Key K (what information do I advertise?), and a Value V (what information do I actually contain?). These are computed by projecting x through three learned weight matrices.

Figure 3 — Computing Q, K, V

One shared x_t fans out to three independent projections, each with a distinct role. Below the separator: q_t · k_s scores how much position t attends to position s; softmax normalises scores into weights α_t,s; and the weighted sum over v_s produces the output for position t. The matrix operation QK^⊤ computes all T² dot products simultaneously.

How Q, K, and V acquire meaning. The names Query, Key, and Value are borrowed from database retrieval — a useful analogy, but one the architecture does not enforce. The only structural constraint is this: Q and K are compared via dot product to produce attention scores, while V supplies the content that gets weighted and summed. Which features each projection captures — and whether they behave like "questions," "labels," and "answers" in any intuitive sense — is determined entirely by training. Gradient descent finds whatever assignment of roles minimises prediction loss. The semantic descriptions we attach to these projections are patterns observed in trained models, not properties specified in advance.

Attention computationQ = X W_Q, K = X W_K, V = X W_V // project all positions at once S = Q K^⊤ / √d_k + mask // T×T score matrix (S_t,s = q_t·k_s/√d_k) mask_t,s = 0 if s ≤ t, −∞ if s > t // causal mask: no looking at later tokens A = softmax(S) // row-wise softmax → attention weights (e^−∞ = 0) Output = A V // weighted sum of values

The dot product Q K^⊤ produces a score for every (query, key) pair — how well position t's query matches position s's key. Dividing by √d_k prevents the dot products from growing too large (which would push the softmax into regions of near-zero gradient). The softmax then converts scores into weights that sum to one per row.

Causal masking: During training and inference, a token at position t must not attend to future positions. This is enforced by setting the scores above the diagonal to −∞ before the softmax (which maps them to 0 weight). The model therefore sees a lower-triangular attention pattern — position t can only gather information from positions 0 through t (positions are numbered from 0).

The √d_k scaling factor

The scaling by √d_k is subtle but important. If q and k are d_k-dimensional random vectors with unit-variance components, their dot product has variance d_k (it's a sum of d_k random terms). Dividing by √d_k brings the variance back to one. Without this, large dot products would make the softmax extremely peaky — one position would receive nearly all the attention — which creates a gradient bottleneck.

Interactive: Softmax on Attention Scores

Edit the raw attention scores below to see how softmax converts them to probabilities. Notice how even a small score difference produces a large probability concentration.

Neuroscience parallel: The attention pattern — the matrix of α_t,s values that determines how much each position attends to each other — resembles effective connectivity between brain areas: dynamic, context-dependent routing of information that changes trial by trial. The learned projection matrices W_Q, W_K, W_V are the fixed-weight tracts; the attention pattern they produce is recomputed fresh for every new input. This is reminiscent of top-down attentional modulation in the visual cortex, where the strength of signal propagation from one area to another is gated by task demands.

Multi-Head Attention

A single attention head produces one d_k-dimensional output per position. But a sentence has many kinds of structure simultaneously — syntactic dependencies, coreference relations, semantic roles, positional proximity. Running H parallel attention heads allows the model to attend to different types of relationships at once, each head learning its own Q/K/V projection matrices.

Each head reads the full d-dimensional stream and projects it, with its own W_Q, W_K, W_V, into a smaller space of size d_k = d / H. For GPT-2 small with d = 768 and H = 12 heads, each head works in its own 64-dimensional space — a learned projection of all 768 dimensions, not a slice of them. The outputs from all heads are concatenated and projected back to d dimensions.

MHA(X) (multi-head attention) is the whole attention sub-layer of one block written as a single function: it takes the stream X (T × d: one row per token) and returns an update of the same size, which is added to the stream. In GPT-2 its input is the layer-normed stream, so a block computes X′ = X + MHA(LayerNorm(X)) (tab 6).

Figure 4 — Multi-Head Attention

Input X ∈ ℝ^T×d

↓

↓ concatenate

Concat(head₀, …, head_H−1) ∈ ℝ^T×Hd_k = ℝ^T×d

↓

× W_O ∈ ℝ^d×d

↓

MHA output ∈ ℝ^T×d

Each head independently computes attention over the sequence in its own projected 64-dimensional space. Concatenation places the heads' outputs side by side, and W_O maps them back to d dimensions; as the equations below show, this is the same as each head writing its own update and the updates being added.

Multi-head attentionAttention(Q, K, V) = softmax(Q K^⊤ / √d_k + mask) V // one head (tab 3); mask = −∞ above the diagonal head_h = Attention(X W^h_Q, X W^h_K, X W^h_V) // each head h = 0 … H−1; T × d_k MHA(X) = Concat(head₀, …, head_H−1) W_O // join and project: T × d = head₀ W⁰_O + head₁ W¹_O + … + head_H−1 W^H−1_O // the same: a sum of the heads' writes W^h_Q, W^h_K, W^h_V ∈ ℝ^d×d_k where d_k = d/H W_O ∈ ℝ^d×d; W^h_O ∈ ℝ^d_k×d = rows h·d_k … (h+1)·d_k−1 of W_O // biases (b_Q, b_K, b_V, b_O) omitted

The two forms of MHA(X) are equal because multiplying a concatenation by a matrix is the same as multiplying each piece by its own block of rows and adding: head h only ever meets rows h·d_k to (h+1)·d_k−1 of W_O, its own W^h_O. So each head writes its own d-dimensional update, head_h W^h_O, and the attention sub-layer's output is their sum (plus the bias b_O). Heads in the same layer do not exchange information; they interact only through the stream, when a later layer reads what an earlier one wrote. Interpretability work uses this per-head form to ask what each head contributes.

Symbol	Meaning	Size (GPT-2 small)
X	the (layer-normed) residual stream, one row per token	T × 768
T, d, H, d_k	tokens, stream width, heads per layer, head width (d / H)	T, 768, 12, 64
W^h_Q, W^h_K, W^h_V	head h's query, key and value matrices (superscript = which head, subscript = which matrix)	768 × 64 each
Q, K, V	one head's queries, keys and values: what each token looks for, offers, and passes on	T × 64 each
head_h	head h's output: at each position, the attention-weighted average of the values	T × 64
W_O, W^h_O	output matrix of the layer; head h's block of 64 rows of it	768 × 768; 64 × 768
MHA(X)	the attention sub-layer's output, added to the stream	T × 768

What do different heads specialize in?

Interpretability research has found that some attention heads develop recognisable functions. Much of the evidence for the first three types below comes from BERT, a bidirectional model (Clark et al., 2019); GPT-style models show similar heads, plus induction heads. Some well-documented head types include:

Syntactic heads — attend from a word to a grammatically related word, such as a direct object to its verb or a determiner to its noun; in BERT, some heads follow particular dependency relations with fairly high accuracy.

Positional heads — attend to the immediately preceding token regardless of content, effectively passing information one step along the sequence. (BERT also has heads that attend to the following token; a causal mask forbids this in GPT-style models.)

Coreference heads — link a pronoun or repeated mention back to an earlier mention of the same entity ("she" → "the scientist").

Induction heads

›

If the sequence contains "A B … A", two heads work together. First, a previous-token head in an earlier layer writes into B's position that B is preceded by A. Then, at the second A, the induction head looks for a position preceded by A, so it attends to B and copies B into the output — predicting that "A B" will repeat. Olsson et al. (2022) present preliminary evidence that these heads underlie much of in-context learning: the model's ability to pick up on patterns demonstrated earlier in the prompt (see Lesson 7, tab 5).

Importantly, no single head handles everything. The power comes from the ensemble — H heads each contributing a different "read" of the context, each mapped back to d dimensions by its block of W_O, and added together into a single update to the residual stream.

Interactive: Head Specialization

The heatmap below shows attention weights for the same sentence across four hypothetical heads. Row t, column s shows how much position t attends to position s (upper-triangle cells are masked — causal attention only looks back). Switch heads to see the contrast: positional, syntactic, duplicate-token, and global patterns all emerge from the same architecture with different learned weights.

Neuroscience parallel: Multi-head attention is conceptually similar to parallel processing streams in sensory cortex — the ventral (what) and dorsal (where) visual pathways extract different features from the same retinal input simultaneously. In auditory cortex, pitch, timbre, and spatial location are processed in partially distinct pathways before being integrated. The key difference is anatomical: these cortical streams are physically separated, while attention heads are mathematical abstractions running in parallel on the same hardware.

Feedforward Layers and Layer Normalization

The Feedforward Network

After attention gathers information across positions, a position-wise feedforward network processes each token's vector independently. It is a two-layer MLP: expand from d to 4d dimensions, apply a nonlinearity, then project back to d. Despite its simplicity, the feedforward layer holds more parameters than the attention layer for most models (8d² vs 4d² per block).

Division of labor. Each transformer block alternates between two categorically different operations. Attention is the only place tokens communicate — it routes information across positions, assembling context. The feedforward network then operates on each position in isolation, with no awareness of neighboring tokens, transforming what attention assembled into a richer local representation. This alternation — gather, then process; gather, then process — repeats L times and is the core computational motif of the transformer.

Figure 5 — Feedforward Expansion and Compression

The input vector of dimension d is first expanded to 4d by the first linear layer, a nonlinearity is applied element-wise, then it is compressed back to d. This expansion-compression pattern can act as a kind of memory: research suggests each hidden unit works like a key–value pair, with the key responding to particular input patterns and the value promoting likely next tokens (Geva et al., 2021).

Feedforward networkFFN(x) = GELU(x W_up) W_down // x is a row vector, as in attention W_up ∈ ℝ^d×4d // expand: d → 4d W_down ∈ ℝ^4d×d // compress: 4d → d

GPT-2-era models (GPT-2, GPT-3, BERT) use GELU (Gaussian Error Linear Unit) rather than ReLU; most recent models use the gated SwiGLU variant described below. GELU is smoother and approximates ReLU while keeping a nonzero gradient for almost all inputs (unlike ReLU's flat zero for x < 0). It is defined as x · Φ(x) where Φ is the standard normal CDF — in practice approximated as 0.5x(1 + tanh(√(2/π)(x + 0.044715x³))). For a detailed treatment of ReLU, GELU, and activation functions more broadly, see Lesson 1, tab 2.

Gated MLPs and W_gate

Many recent models — LLaMA, PaLM, Mistral — replace the two-matrix FFN with a gated variant that adds a third weight matrix, W_gate. Instead of applying the nonlinearity directly to x W_up, the hidden layer is element-wise multiplied (⊙) by a learned gate:

SwiGLU feedforward (gated)FFN_SwiGLU(x) = ( x W_up ⊙ SiLU(x W_gate) ) W_down W_up, W_gate ∈ ℝ^d×d_ff // both expand d → d_ff W_down ∈ ℝ^d_ff×d // compress d_ff → d d_ff ≈ (8/3)d // LLaMA: 3·d·d_ff ≈ 8d², same as the 4d GELU FFN SiLU(z) = z · σ(z) // "Sigmoid Linear Unit" — smooth gate, ≈ 0 for z ≪ 0

The SiLU gate selectively opens or closes each hidden dimension based on the current input. W_gate learns which intermediate directions are worth passing through; W_up learns what content to project into those directions. The result is that the network has dynamic, input-dependent control over which "neurons" contribute — a richer computational motif than a fixed nonlinearity applied uniformly. Because the gate adds a third matrix, the hidden size is reduced from 4d to about (8/3)d (as in LLaMA) so the FFN keeps roughly 8d² parameters. At that matched parameter count SwiGLU outperforms plain GELU (Shazeer, 2020), which is why it is now the dominant FFN design in frontier models.

FFN as memory: Geva et al. (2021) found that the feedforward layers act like key-value memories. W_up columns are "keys": each fires on particular input patterns (shallow ones such as a shared last word in lower layers, more semantic ones in upper layers). W_down rows are "values": what is written back to the residual stream when a key fires, which raises the probability of tokens that tend to follow that pattern. (An illustrative example: a key that fires on "The Eiffel Tower is in" could write a direction that promotes "Paris".) Later work links factual recall to such mid-layer FFN memories (Lesson 7, tab 7).

Layer Normalization

Before each sub-layer, a copy of the residual stream is normalized across its d dimensions to have zero mean and unit variance, then rescaled by learned parameters γ and β. This gives every attention and feedforward sublayer an input of consistent scale, however large the stream has grown — useful when composing dozens of transformations (24 sublayers in GPT-2 small, 96 in GPT-2 XL). Only the copy fed into the sublayer is normalized; the residual stream itself is not. Many recent models (Llama, Mistral) use RMSNorm, a simpler variant that skips the mean subtraction and the shift β.

Layer normalizationLayerNorm(x) = γ ⊙ (x − μ) / σ + β μ = mean(x) // mean over d dimensions σ = std(x) // std over d dimensions γ, β ∈ ℝ^d // learned scale and shift

Figure 6 — Layer Norm Step by Step

A 5-dimensional token vector walked through each arithmetic step. Note that γ and β restore the network's ability to express any mean and scale — LayerNorm does not lock representations to zero mean, it removes uncontrolled scale before reintroducing learned scale. Layer norm operates independently at each position, normalizing across the d feature dimensions (not across the sequence).

Pre-norm vs post-norm

The original "Attention Is All You Need" paper (Vaswani et al., 2017) placed layer norm after each sublayer (post-norm). Modern models — GPT-2, GPT-3, LLaMA — place it before (pre-norm). Pre-norm is easier to train at scale because it leaves a clean identity path: each sublayer's output is simply added to the stream, nothing normalizes the stream itself, and gradients pass straight back through the additions. Post-norm puts a layer norm on that path after every addition, which makes gradients near the output large at the start of training and calls for a careful learning-rate warm-up (Xiong et al., 2020). The price of pre-norm is that the stream's norm grows with depth, which is one reason a final layer norm is applied before the unembedding.

Neuroscience parallel: The FFN's expand–nonlinearity–compress structure (d → 4d → d) resembles the local circuit motif of a cortical column: inputs arrive at a position, are fanned out across a larger population of neurons (layer 4 → supragranular layers), a nonlinearity shapes how strongly each unit responds, and the result is projected back to a lower-dimensional output. Crucially, this processing is local — each token's FFN computation is independent of every other position, just as a cortical column's recurrent processing is largely confined to its own spatial neighborhood before long-range projections (the analog of attention) carry information across areas. Unlike ReLU, GELU has no hard threshold: negative inputs give small nonzero outputs, so the 4d units are not silenced, and their activity is graded and fairly dense. The closer analog of Olshausen & Field's sparse coding — cortex representing stimuli with a small fraction of active neurons — is found not in single units but in features: directions recovered by sparse autoencoders (Lesson 7, tab 8).

The Full Forward Pass

With the pieces in place — embeddings, residual stream, attention, feedforward, layer norm — we can trace exactly what happens to an input sequence from raw tokens to output logits. Step through each stage below.

Tokenization and Embedding Lookup

›

The input string is tokenized into T integer IDs. Each ID is used to select a row from the embedding matrix E ∈ ℝ^|V|×d. A positional embedding P_t is added at each position.

Input: T token IDs

E ∈ ℝ^|V|×d

→

X⁽⁰⁾ ∈ ℝ^T×d

Each row of X⁽⁰⁾ is the starting vector for that token position. All subsequent computation operates on these T vectors in parallel.

Layer Norm (pre-norm)

›

At each block, the residual stream is first normalized by LayerNorm (layer normalization, tab 5: each row centred, scaled to standard deviation 1, then given a learned scale and shift; not a logarithm). LayerNorm is applied independently at every position (row of X).

X^(l) ∈ ℝ^T×d

→

LayerNorm(X^(l)) ∈ ℝ^T×d

The normalized version is passed into the attention and feedforward sublayers. The original X^(l) is retained for the residual connection.

Multi-Head Self-Attention

›

For each of H heads, Q, K, V are computed by linear projection. Attention weights A are formed as softmax(QK^⊤/√d_k) with causal masking applied. The attended values AV are computed, all H head outputs are concatenated, and projected by W_O: the same as adding each head's own write, head_h W^h_O, plus the bias b_O (tab 4).

LayerNorm(X) → Q, K, V ∈ ℝ^T×d_k per head

A = softmax(QK^⊤/√d_k) ∈ ℝ^T×T

MHA(LayerNorm(X)) ∈ ℝ^T×d

The attention computation has complexity O(T²d) — it scales quadratically with sequence length T, which is one motivation for efficient attention variants in long-context models.

Residual Addition (post-attention)

›

The attention output is added back to the original (pre-norm) residual stream. This is the first residual connection in the block.

X′ = X^(l) + MHA(LayerNorm(X^(l)))

X′ now contains the original information in X^(l) plus contextual updates from attention. The feedforward layer will read from X′.

Feedforward Network (position-wise)

›

A second layer norm is applied to X′, then the feedforward MLP transforms each position's vector independently: expand d → 4d with W_up, apply GELU, compress back 4d → d with W_down. A second residual connection adds this back to X′.

X^(l+1) = X′ + FFN(LayerNorm(X′))

This completes one block. X^(l+1) is the residual stream entering the next block. Steps 2–5 repeat for each of the L blocks.

Final Layer Norm and Unembedding

›

After the last block, a final layer norm is applied. The resulting T×d matrix is multiplied by the unembedding matrix — in GPT-2, the transpose of the embedding matrix E^⊤ (weight tying; Pythia, Llama and Mistral learn a separate matrix instead) — producing logits over the vocabulary for each position.

X^(L) → LayerNorm(X^(L)) ∈ ℝ^T×d

logits = LayerNorm(X^(L)) · E^⊤ ∈ ℝ^T×|V|

P(next token | context) = softmax(logits[T−1])

When generating text, only the last position's logits are used to pick the next token; in training, every position's logits are scored against the token that follows it. The output is a probability distribution over all |V| tokens (50,257 for GPT-2).

Complete forward pass summaryX⁽⁰⁾ = Embed(tokens) + PosEmbed // T × d for l = 0 to L−1: // block l maps X^(l) → X^(l+1) X′ = X^(l) + MHA(LayerNorm(X^(l))) // attention sub-layer X^(l+1) = X′ + FFN(LayerNorm(X′)) // feedforward sub-layer logits = LayerNorm(X^(L)) · E^⊤ // T × |V| P = softmax(logits[T−1]) // next-token distribution (last position, at generation)

Everything is differentiable. Every operation in this forward pass — embedding lookup, matrix multiplications, layer norms, softmax, GELU — is differentiable with respect to the model parameters. This means the entire stack can be trained end-to-end by gradient descent, minimizing the cross-entropy loss between predicted and actual next tokens.