Lesson 3 · Forward Pass

Quick review from Lesson 2

Before tracing the full forward pass, here are the four building blocks from the previous lesson. Each block operates on the residual stream — a sequence of d-dimensional vectors, one per token position — updating it additively.

Residual Stream

A d-dimensional vector at each token position, updated additively by every sublayer. Every block writes a small delta onto the stream rather than replacing it.

x^(l+1) = x^(l) + Δ^(l)

Attention

Dynamic routing across positions. Each query attends to the keys at its own and earlier positions; the resulting weights determine a weighted sum of value vectors.

A = softmax(QK^⊤/√d_k), Out = AV

Feedforward Network

Position-local transformation. Expands d → 4d, applies a GELU activation, compresses 4d → d. Thought to store many factual associations in its weights.

FFN(x) = GELU(x W_up + b_up) W_down + b_down

Layer Norm

Normalizes each position's vector to zero mean, unit variance, then applies learned scale γ and shift β. Applied before each sublayer (pre-norm).

LN(x) = γ · (x − μ)/σ + β

What's new in this lesson

Three mechanisms that weren't covered in Lesson 2 are central to how a trained model actually runs:

Causal masking — The attention score matrix S ∈ ℝ^T×T has its upper triangle set to −∞ before softmax. This prevents any token from attending to future positions, enforcing the autoregressive constraint.

Weight tying — In GPT-2, the matrix used to convert final representations back into vocabulary logits is the transpose of the same embedding matrix used at the input: logit(j) = h · e_j. This is a geometric dot product between the representation and each token's embedding vector. (Pythia, Llama and Mistral instead learn a separate unembedding matrix.)

Autoregressive generation — At inference, the model produces one token at a time. Each new token is appended to the context and the forward pass runs again. The KV cache makes each step O(T) rather than O(T²).

The complete forward pass

Step through all eleven stages of a transformer forward pass. The architecture diagram on the left highlights the active component; the panel on the right shows its equation, tensor shape, and role.

Example input: "The cat sat on" — four tokens, model dimension d = 768, vocabulary |V| = 50,257.

The loop: Steps 3–8 (LN → Attention → ⊕ → LN → FFN → ⊕) repeat L times. GPT-2 small has L = 12; GPT-2 XL has L = 48. Steps 9–11 run once, after all blocks. Everything is differentiable — see Tab 6.

Causal masking

Inside each attention head, every token computes a score against every other token. Without a constraint, token 3 could attend to token 5 — using information it hasn't "seen" yet. Causal masking prevents this.

How the mask works

The raw score matrix S ∈ ℝ^T×T is computed first — a float for every (query, key) pair. Then the upper triangle (future positions: key position s > query position t) is set to −∞. Softmax maps −∞ → 0, so those positions receive exactly zero attention weight.

Masked attention S_t,s = (q_t · k_s) / √d_k
M_t,s = S_t,s + mask_t,s where mask_t,s = 0 if s ≤ t, else −∞
A = softmax(M) ← rows sum to 1 over attended positions only

Interactive score matrix

Six-token sequence: "The cat sat on the mat". Toggle between the three states the score matrix passes through inside an attention head.

Why this is necessary

During training, the model predicts the next token at every position simultaneously. The loss at position 3 uses the model's prediction for token 4; that prediction must be made from tokens 0–3 only, not from tokens 4+ (which haven't been "read" yet in sequential generation). Without the mask, the model could simply "look up" the answer from the future context, learning nothing useful.

A subtle point: The mask isn't about the model's "knowledge" — it's about enforcing consistent semantics between training and inference. At inference, token 5 literally doesn't exist yet when you're generating token 4. The mask makes training respect the same constraint.

Encoder vs decoder attention

BERT-style encoder models use bidirectional attention — no causal mask. Every token can attend to every other token. This is valid when the task is understanding an existing sequence (classification, extraction), not generating new tokens. GPT-style decoders always use causal masking because they generate left-to-right.

Weight tying

After the final layer norm, the model has a matrix X̂^(L) ∈ ℝ^T×d — one d-dimensional vector per token position. To produce next-token probabilities, these vectors must be projected into vocabulary space (|V| ≈ 50,000 entries). This projection is called unembedding.

The shared matrix

In GPT-2, the unembedding matrix is not a separate learned parameter. It is the transpose of the same embedding matrix E used at the input (many newer models untie the two; see the end of this tab):

Unembedding (weight tied) logits = X̂^(L) · E^⊤ ∈ ℝ^T×|V|
logit_t(token j) = h_t · e_j (dot product of final rep with token j's embedding)

This means the input embedding matrix E ∈ ℝ^|V|×d is used at both ends of the network — once to read tokens into the residual stream, and once to score the final representation against every vocabulary token.

Geometric interpretation

The logit for token j measures how similar the final representation h is to the embedding vector e_j (a dot product, so direction and length both count). Tokens whose embeddings point in similar directions to h get high logits; tokens pointing away get low (or negative) logits.

After the forward pass, the model should have moved the representation at the last position toward the embedding direction of likely next tokens. This is the sense in which the model "expects" certain tokens: their embeddings are close to where the representation landed.

Example: logits after "The cat sat on the"

Illustrative numbers, not GPT-2 output: the final-position representation h dotted with eight token embeddings. mat gets the highest logit: its embedding has the largest dot product with h.

Why tie the weights?

Parameter efficiency: E has |V| × d ≈ 38M parameters (GPT-2). Tying saves a separate 38M unembedding matrix.

Training signal: Gradients flow through both uses of E simultaneously. The embedding for "cat" is updated both as an input representation and as a target direction in logit space.

Empirical results: Weight tying improves perplexity (Press & Wolf, 2017; Inan et al., 2017). It encodes a strong inductive bias: the residual stream should finish in the same geometric space that tokens occupy.

Key insight: After training, the residual stream at the last position has been shaped by every attention head and FFN to point toward the embedding direction of the most likely next token. Unembedding just reads off the dot products.

Not universal

Weight tying is used in GPT-2 and GPT-Neo, but it is not universal. Pythia, Llama 2, Llama 3 and Mistral, among others, use a separate unembedding matrix. The separate matrix adds parameters but gives the model more flexibility to place representations and token embeddings independently.

Autoregressive generation

At inference, the model generates text one token at a time. Each step runs a full forward pass over all tokens seen so far; the output is a probability distribution over the vocabulary at the last position; one token is sampled and appended; repeat.

Temperature

Before softmax, logits can be divided by a temperature parameter τ (written τ here because T is the sequence length). Lower τ sharpens the distribution (the model becomes more deterministic); higher τ flattens it (more uniform, more surprising outputs).

Temperature scaling P(token_j | context) = softmax(logits / τ)_j
τ → 0 : argmax sampling τ = 1 : standard τ → ∞ : uniform

The logits are invented and the picks (green) are scripted to show one plausible run; the slider changes only the bars. A real sampler would rarely pick ⟨eos⟩ in the last step at low τ, and with τ → 0 it would always take the top bar.

Temperature τ 1.0

Top candidates for next token

Generation complete — ⟨eos⟩ token produced.

The KV cache

Without caching, generating token t requires recomputing Q, K, V for all T tokens — O(T²) total work. The KV cache stores the key and value matrices from all previous positions. At step t, the model only computes Q, K, V for the new token, then concatenates the new K and V to the cache. Each generation step is then O(T) in attention.

KV cache math: At step t, cached keys K_0..t−1 and values V_0..t−1 are reused directly. Only the new token's query Q_t is multiplied against all cached keys to get attention scores. Memory cost: O(L × T × d) per sequence — linear in sequence length.

Sampling strategies

In practice, models rarely use plain temperature sampling over the full vocabulary. Common strategies include top-k sampling (sample from the k highest-probability tokens only), nucleus (top-p) sampling (sample from the smallest set of tokens whose cumulative probability exceeds p), and greedy decoding (always pick the highest-probability token). Each trades off diversity against coherence differently.

Everything is differentiable

Training a language model requires computing the gradient of the loss with respect to every parameter. This is only possible because every operation in the forward pass — all eleven steps — has a well-defined gradient. Backpropagation can then propagate the error signal from the loss all the way back to the embedding matrix.

Gradients for each stage

Operation	Key gradient	Notes
Embedding lookup	∂Loss/∂E[id_t] = ∂Loss/∂x_t	Gradient flows directly into the embedding row of the token at position t (summed over every position where that token appears)
Layer Norm	∂LN/∂x = γ/σ · (I − 11^⊤/d − x̂x̂^⊤/d), x̂ = (x−μ)/σ	O(d) per position (cheap); stable gradient magnitude
Linear (W)	∂Loss/∂W = x^⊤ · δ	Outer product of input and output gradient (row-vector convention y = xW); same for W_Q, W_K, W_V, W_O, W_up, W_down
Softmax (attention)	∂a/∂s = diag(a) − aa^⊤, a = softmax(s)	Jacobian is T×T; causal mask zeroes the upper triangle before this
GELU	∂GELU(x)/∂x = Φ(x) + xφ(x) (GELU(x) = x·Φ(x))	Smooth approximation to ReLU; non-zero gradient almost everywhere
Residual add	∂(x + f(x))/∂x = I + ∂f/∂x	The identity term means gradient always passes through; no vanishing
Unembedding (E^⊤)	∂Loss/∂h = ∂Loss/∂logits · E	With weight tying (GPT-2), this gradient also updates E, the input embedding
Softmax → cross-entropy	∂l_t/∂logits_t = P_t − y_t	Elegantly simple: predicted probs minus one-hot true label (per-token loss l_t)

Why residual connections help: The gradient of the loss with respect to layer l's input includes an identity term from the residual add. This means that even in very deep models (L = 48 in GPT-2 XL), gradients flow directly from the loss to early layers without being multiplied through all 2L sublayer Jacobians — the highway that makes deep networks trainable.

From forward pass to training

The forward pass computes a probability distribution over vocabulary at each position. The cross-entropy loss measures how much probability the model assigned to the actual next token:

Training objective Loss = − (1/T) Σ_t=0^T−1 log P(token_t+1 | token_0..t) (a training sequence of T + 1 tokens gives T predictions)
∂Loss/∂logits_t = (P_t − y_t) / T where y_t is a one-hot for the true token

Notation: Loss is the training loss; L stays the number of blocks, as in Lesson 2. Positions are numbered from 0.

Backpropagation propagates this gradient through all eleven forward-pass stages, accumulating ∂Loss/∂W for every weight matrix. An optimizer (usually AdamW) then updates each parameter in the direction that reduces the loss. This cycle — forward pass → loss → backward pass → update — is what training is.

Bridge to Lesson 4: Everything in this lesson — the forward pass, the loss, the gradients — is the mechanism that training operates on. Lesson 4 covers how the optimizer uses gradients to update weights, why training is computationally expensive, and what shapes the model ultimately learns.

Continue to Lesson 4 · Training →