Unit 7.5 · about 5 h
218 / 324Transformer blocks: residual stream, normalization and feed-forward
The block is the unit that stacks. Normalization order decides trainability at depth, and the encoder, decoder and cross choice decides which context each token may use.
- explain
- implement
You can now
- Contrast pre- and post-normalization block orders and predict where activations and gradients differ.
- Implement a pre-normalized residual attention-plus-MLP block with explicit positions and shape invariance.
- Select encoder, decoder or encoder-decoder structure from the task's context needs.
- Explain what a gated feed-forward adds over the plain two-matrix MLP and what it costs.
Covers
- Residual stream and pre versus post normalization
- Feed-forward and gated feed-forward blocks
- Encoder, decoder and encoder-decoder patterns
Study, in this order
3 steps
- Transformer layers, BERT, GPT and encoder-decoder modelsUnderstanding Deep Learning - Transformers
Section 12.4 (pp. 215 to 216) for the post-norm layer in equation 12.13; Section 12.6 (pp. 219 to 221) for BERT widths, heads and parameter counts; Section 12.7 (pp. 222 to 225) for autoregressive factorization, masked training and GPT-3 sizes; Section 12.8 (pp. 226 to 227) for cross-attention with Figure 12.13
Depth · Skip Section 12.5.1 tokenization and the GPT-3 few-shot anecdotes; stop Section 12.7 at generation mechanics.
- Feed-forward, `AddNorm`, encoder and decoder blocksD2L - attention and transformer implementation route
PositionWiseFFN, theAddNormresidual-plus-LayerNormclass,TransformerEncoderBlock,TransformerDecoderBlockwith its three norms and same-width comment, and the encoder shape-invariance checksDepth · Read Training for loss wiring only; do not run full machine-translation training here.
- Pre-norm vision transformer blockD2L - attention and transformer implementation route
ViTBlockwith its pre-versus-post-normalization paragraph,PatchEmbeddingwith the(96, 16, 512)shape check, and Putting It All Together for encoder-only assemblyDepth · Skip Training; the Fashion-MNIST run is not part of this unit.
Focus on
- Post-norm (UDL equation 12.13, D2L encoder) normalizes after each residual add; pre-norm (
ViTBlock) normalizes the branch inputs and adds untouched residuals. D2L states pre-norm trains more effectively. - Residuals force equal widths: D2L requires decoder
num_hiddensto equal the encoder width for the adds and the cross \(QK^\top\) product. - Neither D2L MLP is gated:
PositionWiseFFNisReLUandViTMLPisGELUwith dropout. A gated block adds a third projection and an elementwise product; hold parameters fixed by shrinking width, not by ignoring the extra matrix. - Pattern follows context: BERT-style encoders read both sides for representations, GPT-style decoders mask the future for generation, encoder-decoders add cross-attention for sequence-to-sequence tasks.
Reproduce closed-book
ViTBlock-pattern pre-norm block (norm, attention, add, norm, MLP, add) and show a (2, 100, 24) input returns (2, 100, 24). Then convert one D2L post-norm block to pre-norm and list exactly which lines moved.Traps
- Calling SDPA with nonzero
dropout_pin eval. Dropout applies in every mode, so a finished block still jitters; set0.0outside training. - Normalizing the wrong axis.
LayerNormstatistics run over the feature width (D2Lnorm_shapeequalsnum_hiddens); including batch or sequence axes computes nonsense moments. - Counting a gated MLP as two matrices. The gate is a third full projection; parameter and FLOP claims must include it.
Practice
TrenTorch: Pre-norm vs. post-norm
Implement post_norm_transformer_block_forward and the two stacking functions; the tests compare how the residual stream grows with depth
Pre-norm adds each sublayer to the raw stream, so its norm grows with depth; post-norm renormalizes after every residual. The section mini-project then builds the pre-normalized residual attention-plus-gated-MLP block with shape-invariance checks and serialization.
Self-test
-
Write the operation order of a post-norm and a pre-norm attention sublayer.
Post-norm: \(Y\) equals LayerNorm(\(X\) plus Attention(\(X\))). Pre-norm: \(Y\) equals \(X\) plus Attention(LayerNorm(\(X\))). Same parts, different order; only pre-norm keeps a clean residual path. -
Contrast BERT and GPT-3 on attention pattern, objective and size.
BERT is a bidirectional encoder trained to predict masked words: 24 layers, width 1024, 16 heads of width 64, MLP width 4096. GPT-3 is a causal decoder trained on next tokens: 96 layers, width 12288, 96 heads of width 128. -
Why must the D2L decoder match the encoder
num_hiddens?Residual adds need identical shapes, and encoder-decoder attention multiplies decoder queries against encoder keys. Different widths break both the adds and the score product. -
What is missing from D2L MLPs if the target is a gated feed-forward?
A third projection and the elementwise gate product. Both D2L MLPs are plain two-matrix stacks, so the gated block must be built and its extra matrix counted.
Completion criterion
Stop when
Your block holds shapes across depths, gives the same output on two eval-mode runs, and you can convert post-norm to pre-norm on demand.
Mark passed, then open 7.6 Only mark this when you can do it without notes.Leads to
- 7.6 Parameter, FLOP and memory accounting for attention Count the block you just built.
- 7.7 Architecture trade-offs and attention correctness tests Stress-test the block you just built.
Comes after
- 3.3 Batch, layer, group and RMS normalizationPlace normalization inside residual blocks.
- 5.6 Paired modalities: projector heads, alignment and zero-shot retrievalOpen the image and text encoder blocks that feed the shared space.
- 7.2 Causal masks, padding masks and fully masked rowsPlace the mask inside a full residual block.
- 7.3 Self-attention, cross-attention and multi-head shapesPlace multi-head attention inside a full residual block.
- 7.4 Positional information and position offsetsAdd positions before the first residual block.
Resources for this unit
- Understanding Deep Learning - Transformers Use Chapter 12.1-12.11 and the matching notebooks for dot-product self-attention, extension to transformers, transformer layers and vision transformers.
- D2L - attention and transformer implementation route Use queries-keys-values, attention-scoring-functions, multihead-attention, self-attention-and-positional-encoding, transformer and vision-transformer. Implement masked single-head attention, multi-head attention and one pre-normalized block before comparing with PyTorch.
Also in this section
- PyTorch - scaled dot product attention Use the official function documentation as a differential oracle for shapes, masks, dropout and backend behavior after the manual version passes intervention tests.
- Deep-ML candidate - attention problems Use 490 Scaled Dot-Product Attention as the fifth platform admission sample. If the judge catches a seeded masking error, use at most one of 490, 491, 53, 904 or 905 as a changed-prompt drill.
- TrenTorch - attention and block questions Solve the multi-head split, sinusoidal positional encoding, causal mask and pre-norm versus post-norm questions the units name, then carry the same checks into your own block.