The Transformer Architecture End-to-End

An autoencoder turned a return-photo patch into a compact representation. A Transformer must build a useful representation at every token position, because a support reply such as Order shipped from depends on what came earlier and must still predict what comes next.

This chapter follows a small decoder-only Transformer forward pass. We'll keep the dimensions tiny enough to print, then assemble the same operations in PyTorch. The next chapter will add the learning objective that trains these predictions.

The Transformer introduced in 2017 used an encoder and a decoder for translation. Its decoder combines masked self-attention, encoder-decoder attention, and a feed-forward network. A text-generation decoder-only model keeps the causal self-attention path and predicts the next token from preceding context.^[1][2]

Tokens become positioned vectors

Use a simplified support-chat vocabulary. We won't claim these are IDs from a production tokenizer; we're choosing IDs deliberately so the forward pass is readable:

A model doesn't operate on token labels. It looks up one learned vector per ID, then adds or applies position information so shipped from differs from from shipped. In this first calculation, position is a small vector added to each token embedding.

Our tiny model uses three tokens and four features, so its hidden-state matrix has shape [3, 4]. Real models use far wider vectors and batches, but the axes mean the same thing:

Attention lets a token read earlier tokens

An RNN transports information through one recurrent state. Self-attention takes a different route: at each position it computes relevance scores against other visible positions, then blends their value vectors. For a decoder, "visible" means the current and earlier tokens only.

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.

The projections come from the same hidden-state matrix x, using different learned weights W_q, W_k, and W_v.^[1]

To see the multiplication, take one two-feature attention head. These arrays stand in for outputs of its learned projections:

Three queries compared with three keys produce a [3, 3] matrix. The square T x T attention-score matrix is why full attention becomes expensive as context length grows: doubling T gives four times as many pair scores.

Scale before softmax

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

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

Here, d_k is the number of features in one key or query vector. Vaswani et al. explain that without scaling, larger dot products push softmax toward very small gradients when the key width is large.^[1]

This calculation uses a 64-feature head and compares unscaled and scaled scores for one token:

Without scaling, nearly all attention mass lands on one key. Scaling doesn't prevent a confident choice after training; it prevents the head width alone from making early scores overly sharp.

The causal mask prevents answer leakage

During training, we can process Order shipped from hub in one matrix operation. Position from is supposed to predict hub, so it must not read the hub representation first. A causal mask sets future scores to negative infinity before softmax; those entries receive probability zero.

The important output is future attention mass: 0.0. If a causal language model reports nonzero future attention during training, it has a data-leak bug, not an impressive loss curve.

Heads run in parallel; blocks run in sequence

One attention head can learn one set of relevance patterns. Multi-head attention projects the model features into several heads, applies attention in parallel, concatenates their outputs, and projects back to d_model.^[1]

Suppose our tiny model width is four and it has two heads. Each head owns two features. Splitting and merging must preserve every token position:

The head axis is parallel work inside one block. The block axis is sequential depth: output from block 1 enters block 2. Mixing these concepts leads to configuration and shape mistakes.

A block alternates communication and local transformation

Attention lets token positions exchange information. A feed-forward network (FFN) then transforms each position independently, using the same weights at every position. The original Transformer used a two-layer FFN with a ReLU nonlinearity between the layers.^[1]

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

The first projection commonly widens the feature dimension; the second returns it to d_model, which keeps blocks stackable.

Residual paths and normalization

A residual connection adds a sublayer result back to its input. It creates an identity path through deep stacks, following the residual-network idea introduced by He et al.^[3] Layer normalization computes statistics within each token vector; its learned scale and offset can then adjust the normalized output.^[4]

There are multiple valid placements. The 2017 Transformer applied normalization after each residual addition. The diagram and PyTorch block below use pre-normalization: normalize first, compute the sublayer update, then add it to the residual stream. The two layouts share attention, FFN, and residual concepts; don't accidentally describe one while coding the other.^[1] Xiong et al. connect Pre-LN with better-behaved gradients at initialization and show that it can train without learning-rate warmup in their evaluated tasks.^[5]

Implement masked attention in PyTorch

Now implement a single causal multi-head self-attention layer. The projections are learned linear maps; the mask is created from token positions. masked_fill is the key line: it turns every future score into negative infinity before softmax.

The output shape matches the input shape, while the [batch, heads, tokens, tokens] weight tensor exposes each head's routing decisions.

Assemble one decoder block

With attention implemented, a small pre-normalized block is short: normalize, attend, add; normalize, run FFN, add. PyTorch's LayerNorm includes the learned scale and offset that the earlier NumPy calculation omitted. The block below uses GELU instead of the original Transformer's ReLU, so you can see that the activation choice can change while the FFN's structural role stays the same: widen, transform, and return to d_model.

This block doesn't generate a word yet. 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

After the final block, many pre-normalized decoder stacks apply one final normalization before the output head. For example, GPT-2 applies its final ln_f immediately before computing logits.^[6] A learned output matrix then converts each hidden vector into one raw score per vocabulary item. These scores are logits. During generation, use only the final position's scores because that position represents all text currently visible to the decoder. The small calculation below starts from an already normalized final-position vector.

The result hub completes Order shipped from hub. Conceptually, generate a longer support answer by appending that token, running the stack with the longer context, and choosing the next one. Optimized serving reuses earlier attention keys and values through a KV cache instead of recomputing them for every new token. The next chapter turns this forward pass into a trained language model.

Encoder, decoder, and encoder-decoder visibility

Decoder-only generation is one member of the Transformer family. Architecture choice is largely a question of which positions may read which information:

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 encoder sequence, so its core path is causal self-attention.

Diagnose a broken forward pass

The fastest way to debug a Transformer is to ask what invariant was violated:

Use a deliberately bad mask to make leakage visible:

If that second value isn't zero in a causal decoder, fix the mask before trusting any metric.

Mastery check

Key concepts

• Token IDs become positioned vectors with shape [B, T, d_model].
• Scaled dot-product attention mixes visible token values using query-key relevance.
• A causal mask gives future positions zero attention weight.
• Heads run in parallel inside a block; blocks run sequentially.
• Attention communicates across positions; the FFN transforms each position.
• Residual paths and normalization let same-shaped blocks stack.
• In the pre-LN pipeline shown here, a final normalization prepares the residual stream before the output projection turns the last hidden state into next-token logits.

Evaluation rubric

• Foundational: Traces Order shipped from from token IDs through hidden states to a next-token distribution.
• Foundational: Explains why the causal mask must hide hub from the from position during training.
• Intermediate: Derives the score and output shapes for one attention head and for a multi-head merge.
• Intermediate: Implements causal attention with a future-mask assertion.
• Intermediate: Explains attention, FFN, residual, and normalization roles without treating them as interchangeable.
• Advanced: Distinguishes encoder-only, decoder-only, and encoder-decoder visibility rules and selects a variant for a stated task.

Common pitfalls

• Treating query, key, or value as a fixed property of a word rather than a learned per-layer projection.
• Applying softmax before causal masking, which allows future answers to influence the output.
• Confusing parallel attention heads with sequential decoder blocks.
• Projecting to vocabulary size inside every block instead of once after the stack.
• Assuming every Transformer has the same normalization placement or cross-attention path.

Follow-up questions