The Transformer Architecture End-to-End

Autocomplete has read function returns the and needs one more token. result feels plausible, but picking it requires each token to carry its own information while the consults its history without peeking at the answer. What changes between that text and the final choice?

Follow one small decoder-only Transformer forward pass. Three token IDs become positioned vectors, attention exchanges information along legal paths, a block updates every row without changing its width, and the final row becomes vocabulary scores. Small NumPy experiments expose individual operations; three linked PyTorch cells then assemble the complete model. The weights are illustrative, not trained, so a plausible continuation isn't evidence of language understanding.

The preceding autoencoder lesson squeezed an input through a narrow latent code. Here we keep one vector per token and let those vectors communicate. You'll reuse matrix multiplication, nonlinear layers, and the softmax calculation from Softmax & Cross-Entropy. No new differentiation rules are needed for this forward pass.

Before reading the pipeline, predict two shapes: with three tokens and width four, the hidden state should stay [3, 4]; with six vocabulary entries, the selected row should produce six logits. Keep those predictions in view as the data moves.

That path is what we mean by a Transformer here: a stack of attention and position-wise feed-forward updates that carries a same-width representation forward.

The architecture introduced in 2017 paired an encoder with a decoder for translation. A text-generation decoder-only model keeps the decoder's masked self-attention path and predicts each token from preceding context.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762[2]Reference 2Improving Language Understanding by Generative Pre-Training.https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf}

Tokens become positioned vectors

Start at the input boundary with a tiny autocomplete vocabulary. These aren't IDs from a production tokenizer. We're choosing them so the forward pass is readable:

A model doesn't operate on raw token strings. It looks up one learned vector per ID, then adds position information so returns the differs from the returns. Without positional cues, pure self-attention is permutation equivariant: permuting the rows of an input matrix produces the exact same permutation in the output. Shuffling function returns the into the returns function would yield the same token vectors in shuffled row order. Adding or injecting position breaks that symmetry so the network knows the follows returns.

Architectures handle position in three distinct ways:

1. Fixed sinusoidal encodings: The original 2017 Transformer generated sine and cosine waves of varying frequencies, allowing relative shifts to be expressed as linear transformations without adding learned parameters.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762}
2. Learned absolute positional embeddings: GPT and BERT learned an explicit embedding table $W_{\text{pos}} \in \mathbb{R}^{T_{\text{max}} \times d_{\text{model}}}$, adding row $t$ directly to token $t$.^{[2]Reference 2Improving Language Understanding by Generative Pre-Training.https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf[3]Reference 3BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.https://arxiv.org/abs/1810.04805} While simple, this creates a hard cutoff: the model can't process sequences longer than $T_{\text{max}}$.
3. Rotary Position Embedding (RoPE): Modern LLMs such as Llama and Mistral rotate query and key vectors inside each head by an angle proportional to sequence index.^{[4]Reference 4RoFormer: Enhanced Transformer with Rotary Position Embedding.https://arxiv.org/abs/2104.09864[5]Reference 5LLaMA: Open and Efficient Foundation Language Models.https://arxiv.org/abs/2302.13971} Their dot product then depends directly on relative token distance $(m - n)$, enabling context extension.

Our arithmetic uses small additive vectors to keep every coordinate inspectable by hand:

A row belongs to one token position; a column is one feature of that position. Real models use far wider vectors and batches, but the axes mean the same thing:

The full vocabulary is function, returns, the, result, None, value, in that ID order. The output head will score those same six tokens. Input and output weights needn't be tied, but a generated ID must be meaningful when fed back into the embedding table.

Attention lets a token read earlier tokens

At the, the model needs a useful read of its history. An RNN transports information through one recurrent state. Self-attention takes a different route: every position computes a read from the visible positions, then blends their value vectors.

For a decoder, "visible" means the current and earlier tokens only. That rule lets the model process a whole training sequence without letting a later answer flow backward.

Each token produces three learned projections:

• A query asks what information this position needs.
• A key advertises what a position contains.
• A value supplies information if the query matches its key.

Why separate queries, keys, and values instead of dotting the hidden state $X$ with itself? Direct self-matching $X X^T$ creates a symmetric score matrix where token A attends to token B with the exact same weight that token B attends to token A. Human language isn't symmetric: in function returns the, returns needs to inspect its subject function, but function doesn't need to read returns with equal strength. Three separate projection matrices $W_q, W_k, W_v \in \mathbb{R}^{d_{\text{model}} \times d_k}$ decouple those roles, allowing queries to search for specific properties while keys advertise them.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762}

Before checking the numbers, predict the score shape: three query rows compared with three key rows should produce a [3, 3] matrix. The the row can score all three keys at this stage; the causal mask will remove any illegal entries next.

To see the multiplication, take one two-feature attention head. These handpicked arrays isolate the dot product; they aren't the projections used by the complete model later:

There is one score for every query-key pair. That square T x T matrix also explains the quadratic cost of standard attention. Doubling T creates four times as many pair scores before the value mix.

Scores only decide where to read. After scaling and masking, softmax turns each row into weights, and multiplying those weights by V produces one mixed value row for each position. The the row becomes a two-feature update that the heads will merge back into model width.

Scale before softmax

Scaled dot-product attention divides query-key scores by the square root of the head width before applying softmax:

$$\operatorname{Attention}(Q, K, V) = \operatorname{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

Why divide by $\sqrt{d_k}$? Consider the statistical distribution of raw scores. If the components of query $q \in \mathbb{R}^{d_k}$ and key $k \in \mathbb{R}^{d_k}$ are independent random variables with mean zero and variance one, their dot product is:

$$q \cdot k = \sum_{i=1}^{d_k} q_i k_i$$

Each individual product $q_i k_i$ has mean $\mathbb{E}[q_i k_i] = \mathbb{E}[q_i]\mathbb{E}[k_i] = 0$ and variance:

$$\operatorname{Var}(q_i k_i) = \mathbb{E}[(q_i k_i)^2] - (\mathbb{E}[q_i k_i])^2 = \operatorname{Var}(q_i)\operatorname{Var}(k_i) = 1 \times 1 = 1$$

Because the components are independent, their variances sum directly:

$$\operatorname{Var}(q \cdot k) = \sum_{i=1}^{d_k} \operatorname{Var}(q_i k_i) = d_k$$

The standard deviation of the raw dot product grows with $\sqrt{d_k}$. For a typical head width $d_k = 64$ or $128$, raw dot products swing into magnitudes of $\pm 20$ to $\pm 50$. Exponentiating numbers with spreads that large drives softmax into saturation: the largest entry gets a probability near $1.0$, while the rest drop to zero. In that flat regime, softmax gradients $\frac{\partial s_i}{\partial z_j} = s_i(\delta_{ij} - s_j)$ vanish, halting learning.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762} Dividing by $\sqrt{d_k}$ resets the score variance to $1.0$, keeping gradients healthy.

Their base model used $d_{\text{model}}=512$ with $h=8$ heads, so each head had $d_k=64$. The calculation below compares the unscaled and scaled versions of concrete scores [12.4, 9.1, 3.8]:

The unscaled row puts nearly all attention mass on one key, while the scaled row spreads it across visible choices. Scaling doesn't prevent a confident choice after training; it prevents head width alone from making early scores overly sharp.

Use the earlier the row to finish one attention read. Its raw scores [0.2, 1.28, 0.96] become [0.141, 0.905, 0.679] after division by $\sqrt{2}$. Softmax gives weights about [0.21, 0.44, 0.35]; applying them to the three value rows produces approximately [0.34, 0.65]. This is a mixture of values, not a probability distribution over next tokens. The vocabulary prediction comes much later.

The causal mask prevents answer leakage

Training can process the positions in function returns the result in parallel within each layer, but position the is still responsible for predicting result. If its query can read the result representation, it sees the answer. A causal mask sets future scores to negative infinity before softmax; those entries receive probability zero. Setting them to zero wouldn't work: $e^0 = 1$, so forbidden keys would still receive active attention weight. Using $-\infty$ guarantees $e^{-\infty} = 0$.

Predict the allowed cells before using the table: the the row should include function, returns, and the, but not result. The final row can see all four inputs because no later input is present in this example.

Here is a separate four-position score matrix so the forbidden the → result cell is present. It isn't the three-position dot-product example above.

The output shows future attention mass: 0.0. Read row the (row index 2) as [0.067, 0.111, 0.821, 0.000]: its weights sum to one across the current and earlier positions, and the result column is zero. If a causal language model reports nonzero future attention during training, it has a data-leak bug, not an impressive loss curve. Check this invariant before celebrating a low training loss.

Why can't the representation at the attend to result while training a decoder-only next-token model?

Answer

result is the target token the the position is supposed to predict. Letting the read it would expose the answer before prediction, so training loss would no longer measure generation from preceding context.

Heads run in parallel; blocks run in sequence

A single attention head only computes one weighted average across the sequence. But a token often needs to attend to multiple different relationships at once: one head might track syntactic dependencies (matching a verb to its subject function), while another tracks typing or semantic scoping, and a third tracks the immediate prior token.

Multi-head attention gives a block several such views: it projects model features into heads, applies attention independently in each head, concatenates their outputs, and projects back to d_model.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762} Each projection can use every input feature. Splitting the projected coordinates doesn't restrict head 1 to reading only the first half of the original embedding.

$$\operatorname{MultiHead}(Q, K, V) = \operatorname{Concat}(\operatorname{head}_1, \dots, \operatorname{head}_h) W^O$$

Suppose our tiny model width is four and it has two heads. Each head owns two features. Predict the shape trace before running the cell: [1, 3, 4] should become [1, 2, 3, 2] after the split, then return to [1, 3, 4] after the merge. Splitting and merging must preserve every token position:

The head axis describes independent work inside one block. The block axis is sequential depth: output from block 1 enters block 2. Stacking repeats the structure, not the weights; each block has its own projections. Mixing heads with blocks loses that dependency.

A block alternates communication and local transformation

Attention lets token positions exchange information across sequence steps. The next sublayer changes each row on its own, using the same weights at every position. That feed-forward network (FFN) is where a token's mixed context gets a per-position nonlinear transformation. Geva et al. observed that FFN layers act like key-value memories: the first linear projection detects input patterns, while the second projection writes out factual associations.^{[6]Reference 6Transformer Feed-Forward Layers Are Key-Value Memories.https://aclanthology.org/2021.emnlp-main.446/}

The original Transformer used two linear layers with a ReLU nonlinearity between them, widening from $d_{\text{model}}=512$ to $d_{\text{ff}}=2048$ (a $4\times$ expansion) and back:^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762}

$$\operatorname{FFN}(x) = \operatorname{ReLU}(xW_1 + b_1)W_2 + b_2$$

The first projection widens the feature dimension; the second returns it to d_model, which keeps blocks stackable. This isolated FFN demo feeds in our positioned vectors to show [3, 4] -> [3, 6] -> [3, 4]. In a complete pre-LN block, the FFN instead receives the normalized stream after attention's residual addition. Bias terms are omitted here so the matrix sizes stay visible.

Unlike attention's Q @ K.T, neither FFN multiplication creates a token-by-token matrix. Every row gets the same transformation without reading another row.

Modern production models such as Llama, Mistral, and Gemma replace standard MLPs with SwiGLU (Swish Gated Linear Unit):^{[7]Reference 7GLU Variants Improve Transformerhttps://arxiv.org/abs/2002.05202[5]Reference 5LLaMA: Open and Efficient Foundation Language Models.https://arxiv.org/abs/2302.13971}

$$\operatorname{SwiGLU}(x) = (\operatorname{Swish}(x W_{\text{gate}}) \odot (x W_{\text{up}})) W_{\text{down}}$$

Here, $\operatorname{Swish}(z) = z \cdot \sigma(z)$ (also called SiLU). Instead of a single upward projection, two linear projections run in parallel: one acts as a smooth gating filter that scales the linear path elementwise before projecting back down through $W_{\text{down}}$. Because SwiGLU introduces a third weight matrix, implementations typically size the intermediate dimension to $\frac{8}{3} d_{\text{model}}$ (rounded to multiples of 64 or 256) so total parameter counts match a standard $4\times d_{\text{model}}$ MLP.

Residual paths and normalization

A residual connection adds a sublayer result back to its input. Because attention and the FFN return to d_model, addition keeps the running stream at the same shape. It creates an identity path through deep stacks, following the residual-network concept from He et al.^{[8]Reference 8Deep Residual Learning for Image Recognition.https://arxiv.org/abs/1512.03385}

Placement of normalization determines training stability:

• Post-LayerNorm (2017 original): Normalization sits directly in the residual stream after each addition:^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762} $x_{l+1} = \operatorname{LayerNorm}(x_l + \operatorname{SubLayer}(x_l))$ As backpropagation flows backward from layer $L$ to layer $1$, gradients pass repeatedly through the LayerNorm Jacobian $\frac{\partial \operatorname{LN}}{\partial z}$ at every block. In deep networks ($L \ge 24$), gradients vanish or explode near early layers unless trained with delicate learning-rate warmup schedules.
• Pre-LayerNorm (modern standard): Normalization precedes each sublayer, leaving the main residual addition untouched:^{[9]Reference 9GPT-2 Source Implementation.https://github.com/openai/gpt-2/blob/master/src/model.py[10]Reference 10On Layer Normalization in the Transformer Architecture.https://arxiv.org/abs/2002.04745} $x_{l+1} = x_l + \operatorname{SubLayer}(\operatorname{LayerNorm}(x_l))$ The full recurrence becomes an additive telescope: $x_L = x_0 + \sum_{l=0}^{L-1} \operatorname{SubLayer}(\operatorname{LayerNorm}(x_l))$. Differentiating with respect to an early layer $x_l$ produces: $\frac{\partial \mathcal{L}}{\partial x_l} = \frac{\partial \mathcal{L}}{\partial x_L} \left( I + \sum_{j=l}^{L-1} \frac{\partial \operatorname{SubLayer}(\operatorname{LayerNorm}(x_j))}{\partial x_l} \right)$ The identity matrix $I$ acts as an unobstructed gradient highway, letting error signals travel directly across all layers without decaying. That's why Pre-LN models train reliably without warmups.

Architectures also differ in the normalization statistic itself:

• LayerNorm: Centers by subtracting mean $\mu$, then scales by standard deviation $\sqrt{\sigma^2 + \epsilon}$ with learned scale $\gamma$ and shift $\beta$.^{[11]Reference 11Layer Normalization.https://arxiv.org/abs/1607.06450}
• RMSNorm: Zhang & Sennrich showed that mean-centering is computationally unnecessary; stability comes almost entirely from scaling by the root-mean-square: $\operatorname{RMSNorm}(x) = \frac{x}{\operatorname{RMS}(x)} \odot \gamma, \quad \text{where } \operatorname{RMS}(x) = \sqrt{\frac{1}{d} \sum_{i=1}^d x_i^2 + \epsilon}$ Skipping mean subtraction saves memory bandwidth in GPU kernels, providing a 10% to 50% speedup in normalization runtime while matching LayerNorm's training stability. Llama and Mistral standardise on RMSNorm.^{[12]Reference 12Root Mean Square Layer Normalization.https://arxiv.org/abs/1910.07467[5]Reference 5LLaMA: Open and Efficient Foundation Language Models.https://arxiv.org/abs/2302.13971}

Before running the snippet, predict two checks: normalization should put each row near mean 0 and variance 1, while adding the update should leave the residual output at shape [2, 4]. The code simulates that pre-normalization order with a small update.

The small eps keeps division defined even for a constant row. It also explains why the variances are slightly below one. Learned scale and offset can change those statistics afterward; LayerNorm doesn't promise every output always has mean zero and variance one. In the block view, follow the attention update into x₁, then the FFN update into x₂. Both additions preserve [3, 4].

Implement masked attention in PyTorch

The NumPy example exposed the mask entry by entry. Now keep the same contract in one causal multi-head self-attention layer. The projections are learned linear maps; the mask is created from token positions. Calling masked_fill enforces that visibility constraint: it turns every future score into negative infinity before softmax.

Trace the shapes before running it: [1, 4, 8] becomes [1, 2, 4, 4] for each of q, k, and v; scores and weights stay [1, 2, 4, 4]; merging heads returns [1, 4, 8]. The largest future weight should be exactly zero.

The output shape matches the input shape, and the [batch, heads, tokens, tokens] weight tensor exposes each head's routing decisions. The zero future weight is the causal invariant in PyTorch form.

The PyTorch cell uses random inputs to test shapes and visibility separately from the hand-worked values. It deliberately omits padding, dropout, and a KV cache. All positions are real tokens, so the causal diagonal always supplies at least one visible key. Modern production stacks use fused kernels such as torch.nn.functional.scaled_dot_product_attention, which compute this exact causal calculation without materializing the full $T \times T$ attention matrix in GPU memory.

Assemble one decoder block

Attention now returns a same-width update, so it can sit inside a full pre-normalized block. Read the sequence before the code: normalize, attend, add; normalize, widen through the FFN, return to width eight, and add again. For input [2, 4, 8] and d_ff=24, every residual merge should still be [2, 4, 8].

PyTorch's LayerNorm includes the learned scale and offset that the earlier NumPy calculation omitted. This block uses GELU instead of the original Transformer's ReLU, the same activation family as GPT-2's MLP. GPT-2 uses a tanh approximation; nn.GELU() below uses the default exact formulation. Either way, the FFN's structural role stays the same: widen, transform, and return to d_model.^{[13]Reference 13Gaussian Error Linear Units (GELUs).https://arxiv.org/abs/1606.08415[9]Reference 9GPT-2 Source Implementation.https://github.com/openai/gpt-2/blob/master/src/model.py}

Continue in the same Python session so CausalSelfAttention remains defined. The attention module returns both its update and its weights; the block needs only the update for the residual addition.

The finite-gradient check catches a broken backward pass for this input; it doesn't prove that a deep stack will train well. The block still hasn't generated a word. It produces contextual vectors with the same [B, T, d_model] shape as its input, so another block can consume them.

Project the final position into a vocabulary

The block stack has finished its same-width work. The vocabulary head changes the stream's output width from model features to token scores. GPT-2 applies a final ln_f normalization immediately before computing logits.^{[9]Reference 9GPT-2 Source Implementation.https://github.com/openai/gpt-2/blob/master/src/model.py}

A learned output matrix converts each hidden vector into one raw score per vocabulary item. These scores are logits. Vaswani et al. share that matrix with the token embedding table; GPT-2's wte is reused the same way. This strategy, called weight tying, halves the parameter footprint of the vocabulary layers and enforces geometric consistency between input representations and output predictions.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762[2]Reference 2Improving Language Understanding by Generative Pre-Training.https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf} The toy below keeps a separate output_weight so the multiply is visible.

Now connect the actual prompt to the block. Continue in the same session with the DecoderBlock class above. This time use width four, two heads, and an FFN hidden width of six. The embedding table is exactly the one from the opening NumPy example.

The matrix multiply applies the vocabulary head to every position: [1, 3, 4] @ [4, 6] → [1, 3, 6]. Generation selects the last six scores. Training can use all three rows, with a shifted next-token target for each. The output weights below are deliberately chosen for a readable ranking; the random block hasn't learned English.

The vector sent to the output head has passed through attention, both residual additions, the FFN, and final normalization. Its values differ from the original embedding. The figure drops the batch axis of size one and shows the last-position slice; slicing before or after the position-wise vocabulary projection gives the same six logits.

Choosing the largest logit is greedy decoding. It appends result to make function returns the result.

Alternatively, sampling introduces temperature scaling $T > 0$ to soften or sharpen the logits before softmax:

$$P(w_i) = \frac{\exp(\text{logit}_i / T)}{\sum_j \exp(\text{logit}_j / T)}$$

When $T \to 0$, the probabilities collapse into an argmax choice. At $T = 1.0$, the model samples from its natural probability distribution. When $T > 1.0$, probabilities flatten toward a uniform distribution, producing diverse continuations.

To run another step, add a fourth row to the positional table as well as appending the new token ID. Our three-position demo isn't an unlimited text generator. Optimized serving reuses earlier attention keys and values through a KV cache instead of recomputing them for every new token.^{[14]Reference 14KV Cache Strategieshttps://huggingface.co/docs/transformers/main/en/kv_cache} We haven't implemented that optimization here. The next lesson turns this single step into a generation loop and a training objective.

Encoder, decoder, and encoder-decoder visibility

Decoder-only generation is one member of the Transformer family. Once the job changes, the visibility rule must change with it: classification can use the whole input, while translation needs a full source plus a causal target. Compare which positions each variant may read:

Why did decoder-only models come to dominate modern LLMs? Four architectural properties drove that outcome:

1. Unified pre-training objective: Autoregressive next-token prediction turns all text into an unsupervised training stream. Unlike masked language models that corrupt 15% of tokens or sequence-to-sequence models requiring paired data, every single token in the sequence provides a training loss signal.
2. Zero-cost prompt conditioning: In-context prompting and few-shot examples require no custom architecture. Prompts simply precede generation in the causal sequence, letting the model condition on instructions naturally through attention history.
3. Streamlined KV caching: During autoregressive decoding, a decoder-only model maintains a single KV cache for its causal self-attention layers. Encoder-decoder models must maintain two distinct caches per layer: a static cross-attention cache over encoder states and a growing causal self-attention cache over generated tokens.
4. Predictable scaling behavior: Empirical scaling laws demonstrated that autoregressive decoders scale smoothly across compute, parameters, and tokens, unlocking emergent reasoning without architectural rework.

In self-attention, queries, keys, and values come from one sequence. In encoder-decoder cross-attention, decoder queries read keys and values produced by the encoder. A plain decoder-only text generator has no separate source sequence, so its core path is causal self-attention.

Diagnose a broken forward pass

When a forward pass fails, don't inspect every weight first. Ask which contract broke: visibility, shape, device placement, or normalization. The symptom usually narrows the search:

Use a deliberately bad mask to make leakage visible. Predict the two totals before running the cell: an unmasked matrix should assign positive mass to future cells, while the causal version should assign 0.000.

The unmasked total exceeds one because it adds probability mass from several separately normalized rows. Each row still sums to one. A zero future total is necessary, but inspecting the mask alone can miss leakage elsewhere in the block. A stronger test changes a future input and checks whether an earlier output changes.

Test the whole block, not just the mask

Continue the PyTorch session. Replace the final token's features while keeping the first two positions fixed. Which outputs may change? Causality says the first two must be unchanged, even after the FFN and residual additions. The last output should change, so the test also catches a block that ignores its input.

Try removing masked_fill in CausalSelfAttention. The earlier outputs should now change, causing this test to fail. Keep dropout disabled for this comparison; otherwise random dropout masks can change outputs even when causality is correct.

You double the prompt from 3 tokens to 6, keeping width 4 and two heads. Which tensors grow fourfold, and which only double?

Answer

Each head's score and attention-weight matrices grow from [3,3] to [6,6], four times as many entries. Hidden states, Q/K/V, and FFN activations double their token dimension. Model parameter shapes don't change. This describes the explicit full matrices in our implementation, not every optimized attention kernel's memory use.

You replace the FFN's GELU with no activation at all. The code still preserves shape. What capability disappears from that sublayer?

Answer

Two consecutive linear transformations without an intervening nonlinearity collapse mathematically into one linear map ($W_1 W_2 = W_{\text{combined}}$). The intermediate expansion alone doesn't provide nonlinear per-token transformation. The overall Transformer remains nonlinear because attention and normalization remain, but the FFN loses its capacity to model complex feature interactions.

The attention output is [1,3,4], but you change the FFN's final projection to width 6. Why can't the residual addition be repaired just by reshaping?

Answer

The residual has 12 entries and the FFN update has 18. More importantly, addition must match the same token and feature coordinates. Project the six-wide update back to four features with a learned map; don't reinterpret token or feature axes to hide the mismatch.