Pre-norm vs. post-norm: where LayerNorm sits, and why it changes trainability
The problem, from first principles
[01-transformer-block/06-assemble-full-block] placed [01-layer-normalization-forward]'s LayerNorm BEFORE each sublayer (attention, then FFN), the "Pre-Norm" arrangement, with [03-residual-connection]'s residual always adding a sublayer's output back onto the RAW, un-normalized input. The ORIGINAL 2017 Transformer paper did the opposite: it ran each sublayer FIRST, added the residual, and normalized AFTERWARD, "Post-Norm." Both are structurally valid: the same four ingredients, just reordered. But the reordering has a real, measurable, non-cosmetic consequence for how easy the resulting network is to actually TRAIN, which is why nearly every modern LLM has converged on Pre-Norm despite Post-Norm being the historical original.
The concrete difference: in Pre-Norm, the "residual stream," the running sum of every sublayer's output added back through the network, NEVER gets renormalized. It grows, unboundedly, with every added block (nothing in the architecture ever caps it). In Post-Norm, LayerNorm is the LAST operation at every single block, so the output's scale gets re-pinned at every layer, never growing with depth. Post-Norm's bounded activations sound like the safer choice, but Xiong et al. (2020) showed the opposite is true for TRAINING dynamics: Post-Norm's gradients, at the very deepest stacks, without careful learning-rate warmup, can grow unstably large near the INPUT layers specifically (the opposite failure mode from vanishing), while Pre-Norm's gradients stay comparatively well-behaved regardless of depth, which is the actual, practical reason modern architectures default to it.
From theory to code
Implement post_norm_transformer_block_forward (mirroring [06-assemble-full-block]'s Pre-Norm block, with the LayerNorm/sublayer order reversed), and stack_pre_norm_blocks/stack_post_norm_blocks (running the SAME block, with the SAME weights, repeatedly, to isolate the effect of depth alone).
Constraints
- Post-Norm order:
x = LayerNorm(x + Sublayer(x)), sublayer THEN residual THEN norm, for both the attention and FFN halves. stack_pre_norm_blocks/stack_post_norm_blocksreuse the exact SAMEblock_paramsdict at every iteration (deliberately: this isolates depth's effect from the effect of having different, independently-learned weights per block).- Both stacking functions chain blocks sequentially: block
i's output feeds blocki+1's input.
Hints
Hint 1: Post-Norm's attention half
attn_out, _ = multi_head_attention(x, x, x, num_heads, weight_o, bias_o, mask=mask), x = residual_connection(x, attn_out), x = layer_norm_forward(x, gamma1, beta1, eps). Compare directly against Pre-Norm's order (normalize, THEN attend, THEN residual): this is the same three operations, reordered.
Hint 2: Post-Norm's FFN half, and the stacking loops
Identical pattern: ffn_out = feedforward_sublayer(x, ...), x = residual_connection(x, ffn_out), x = layer_norm_forward(x, gamma2, beta2, eps). Both stacking functions are a simple for _ in range(num_blocks): x = <block_fn>(x, num_heads, **block_params) loop, exactly [01-transformer-block/07-stack-blocks]'s pattern with every block sharing one params dict instead of each having its own.
Click "Run Tests" to test your implementation