NLP Foundations
The Forward Pass End to End
Trace one token's numbers from embedding table to output logits and every "mysterious" transformer component turns out to be a shape-preserving or shape-mixing step in a fixed pipeline.
intermediate · 9 min read
Ask what a transformer "does" and most explanations jump straight to attention. But attention is one stage in a longer, entirely deterministic pipeline that takes a sequence of integer token IDs and produces a probability distribution over the next token. Walking that pipeline once, stage by stage, with shapes attached, demystifies more of a transformer than any diagram of an attention head does.
Stage 1: embedding
Token IDs (integers, shape (batch, seq_len)) are looked up in an embedding table of shape (vocab_size, d_model), producing (batch, seq_len, d_model). This is a gather operation, not a matmul, though it can be expressed as one (a one-hot vector times the embedding matrix). Positional information is then added or otherwise injected, since the embedding lookup alone carries no notion of order (see positional-encodings).
Stage 2: N identical blocks
The embedded sequence then passes through N structurally identical transformer blocks, each performing the same sequence of operations on a (batch, seq_len, d_model) tensor:
x = x + Attention(LN(x)) # mix information across positions
x = x + FeedForward(LN(x)) # transform each position independently
The residual additions (see residual-connections-skip) mean the tensor's shape never changes across a block; only its content is refined. This is a deliberate design choice: because shape is invariant, blocks can be stacked to arbitrary depth and even skipped or reordered in some inference-optimisation schemes, since nothing downstream depends on which block produced a given shape.
Inside attention, the mixing step, Q, K, and V are each a linear projection of x, attention weights are computed as softmax(QK^T / sqrt(d_k)) (see softmax-logits and attention-mechanism), and the weighted sum of V is projected back to d_model. Inside the feed-forward network, the position-independent transform, each position's vector is expanded to a larger intermediate width, passed through a non-linearity (see why-nonlinearity-matters), and projected back down; SwiGLU variants add a gating multiplication here (see feedforward-swiglu).
Stage 3: the output head
After the final block, one more layer norm is applied, and the resulting (batch, seq_len, d_model) tensor is projected by an output matrix of shape (d_model, vocab_size) (often the same weights as the input embedding table, transposed, in a scheme called weight tying) to produce logits of shape (batch, seq_len, vocab_size). Softmax turns the logits at the final position into a probability distribution used for sampling the next token (see sampling-decoding). During training, every position's logits are scored simultaneously against the shifted-by-one labels (see next-token-prediction-cross-entropy).
The whole pipeline, shapes attached
token_ids: (batch, seq_len) int
embeddings: (batch, seq_len, d_model) float
... N blocks, shape-preserving ...
final_hidden: (batch, seq_len, d_model)
logits: (batch, seq_len, vocab_size)
Everything between the first and last line preserves shape; only the embedding lookup and the output projection change dimensionality. That is the entire structural story of a decoder-only transformer.
When it falls down
- Training and inference use the same pipeline differently. Training scores every position in one parallel pass (teacher forcing); autoregressive generation runs the same pipeline once per new token, reusing cached keys and values from earlier positions (see kv-cache) rather than recomputing the whole sequence.
- "Identical blocks" hides real heterogeneity. Blocks share structure but not weights, and empirically different layers specialise, early layers often carry more syntactic and positional signal, later layers more semantic and task-specific signal, even though the forward-pass code treats them uniformly.
- The pipeline view undersells attention's asymmetry. Every other stage costs compute proportional to
seq_len; attention costs compute proportional toseq_len^2, so the "just another stage" framing breaks down precisely at the context lengths where it matters most. - Numerical drift accumulates depth-wise. Because each block's output feeds the next, small numerical differences (precision, kernel implementation, even summation order) compound across many blocks, which is why bit-for-bit reproducibility across hardware is hard even with identical weights.
Further reading
- Vaswani et al., 2017, Attention Is All You Need, arXiv:1706.03762 - the original encoder-decoder pipeline; most current models keep only the decoder stack.
- Elhage et al., 2021, A Mathematical Framework for Transformer Circuits - traces the same forward pass as a composition of linear and attention operations, useful for building intuition beyond the block diagram.