Chapter 4

Inside a Transformer Block

Assemble attention and a position-wise MLP as two updates to a persistent residual stream. Trace normalization, ordering, width changes, activations, stacking, and component ablations through one complete block.

Attention is one update inside a Transformer block. A block keeps a model-width residual stream, adds an attention result, then adds a position-wise MLP result. Normalization controls the inputs or outputs of those branches according to a declared order.

This chapter uses a pre-norm block as its implementation spine:

u=x+Attention⁡(N1(x)),u=x+\operatorname{Attention}(N_1(x)), y=u+MLP⁡(N2(u)).y=u+\operatorname{MLP}(N_2(u)).

The two normalization operations have separate learned parameters. Every residual-stream tensor retains shape (B,T,dmodel)(B,T,d_{model}).

The original Transformer used a post-norm order:

u=N1(x+Attention⁡(x)),u=N_1(x+\operatorname{Attention}(x)), y=N2(u+MLP⁡(u)).y=N_2(u+\operatorname{MLP}(u)).

These equations define different functions. We will calculate both, but we will not switch conventions halfway through an implementation.

The chapter then opens the two operations hidden behind short names:

  • LayerNorm and RMSNorm calculate different statistics along the feature axis;
  • an MLP expands each token record, applies an activation or gate, and contracts it back to model width;
  • dropout changes branch outputs during training but not evaluation;
  • stacking repeats the same shape-preserving interface with different parameters.

After this chapter

  • Trace attention and MLP updates through a residual stream.
  • Compare LayerNorm, RMSNorm, pre-norm, and post-norm computations.
  • Inspect the effect of removing or changing a block component.

Lessons

  1. 01
    The Residual Stream Receives Two Updates

    Trace attention and MLP residual updates, distinguish stream values from branch values, and design controlled branch ablations.

    1 exercise
  2. 02
    LayerNorm Normalizes Each Token Record

    Calculate LayerNorm statistics, learned affine parameters, feature-axis behavior, and per-token independence.

    1 exercise
  3. 03
    RMSNorm Removes Re-Centering

    Calculate RMSNorm and compare its statistics, invariances, parameter count, and axis contract with LayerNorm.

    1 exercise
  4. 04
    Pre-Norm and Post-Norm Change Computation Order

    Derive and compare pre-norm and post-norm residual computations, complete block equations, and reproducibility requirements.

    1 exercise
  5. 05
    The MLP Mixes Features at Each Position

    Trace a position-wise Transformer MLP, derive its shapes and parameter count, and test feature mixing without position mixing.

    1 exercise
  6. 06
    Activations and Gates Change the MLP

    Compare ReLU, GELU, GLU, and SwiGLU-style MLPs through calculations, projection shapes, gates, and parameter budgets.

    1 exercise
  7. 07
    Dropout and Stacking Preserve the Interface

    Calculate inverted dropout, distinguish train and evaluation behavior, state placement, and trace shape-preserving block stacks.

    1 exercise
  8. 08
    Audit a Tiny Pre-Norm Block

    Implement and audit a tiny pre-norm block with intermediate traces, invariants, failure signatures, and branch ablations.

    1 exercise

Review and practice

  1. •
    Review

    Review the complete Transformer block, its numerical distinctions, interface, ablations, and debugging checks.

  2. •
    Exercises

    Solve Transformer block calculation, shape, parameter, operation-order, independence, dropout, and debugging exercises.

Chapter progress