Transformer Block
How are all sub-components assembled into a Pre-LN Transformer Block, and what is its shape ledger?
How this lesson unfolds.
We have analyzed Attention, FFNs, LayerNorm, and Residual shortcuts as isolated pieces. Follow this visual module to see how they wire together into the complete Pre-LN Transformer Block and follow the rigorous shape ledger of every operation.
Why did the original Transformer's Post-LN architecture suffer from unstable training, and how does Pre-LN placement resolve it?
Imagine a river flowing from high mountains down to the sea. In a Post-LN architecture, the engineers build a giant hydroelectric dam and filtration gate right across the main riverbed after every single town. As water passes through 24 consecutive dams, water pressure drops, sediment chokes the gates, and salmon swimming upstream (gradients) are completely blocked. In a Pre-LN architecture, the main river flows completely unobstructed and dam-free. Small side canals branch off to power municipal water plants and return clean water back to the main river. Upstream swimmers have a clean, dam-free highway all the way to the source.
- Learning goalContrast the mathematical and architectural differences between Pre-LN and Post-LN layouts and explain why Pre-LN preserves residual gradient scale.
- Helpful foundationsLayer normalization formulation · Residual connection gradient calculus
- Common trapBelieving that Post-LN and Pre-LN produce identical model capacities. While both have the same parameter count, Post-LN restricts effective depth during training, whereas Pre-LN allows gradients to train all layers equally.
- Next connectionNext, trace the exact forward signal flow and backward residual gradient invariance equations in Part B.
More context & mechanics
The spatial placement of LayerNorm dramatically alters the mathematical properties of backpropagation:
- The Post-LN Architecture (Original Vaswani 2017):
x_(l+1) = LayerNorm(x_l + SubLayer(x_l))
- The Flaw: Normalization is placed directly on the residual stream.
- Consequence: The derivative of LayerNorm scales inversely with the vector norm: ~ 1/\|x\|. As depth increases, \|x\| accumulates, causing backpropagated gradients to shrink exponentially in early layers.
- Symptom: The model cannot train without a delicate, high-risk learning rate warm-up schedule.
- The Pre-LN Architecture (Modern Standards - Radford 2019):
x_(l+1) = x_l + SubLayer(LayerNorm(x_l))
- The Solution: Normalization is moved into the side branch.
- Consequence: The main residual spine x_l is never modified by LayerNorm. The identity shortcut remains pure addition: x_(l+1) = x_l + f(x_l).
- Benefit: Gradients flow backward across 100+ layers effortlessly, eliminating the need for delicate warm-up phases.
Virtually every modern decoder-only language model (GPT-2, GPT-3, GPT-4, Llama-1/2/3, Mistral, Qwen) uses Pre-LN or Pre-RMSNorm. Training large models in FP16/BF16 without Pre-LN frequently results in catastrophic loss spikes and NaN divergence.
The original 2017 Transformer paper utilized Post-LN. In 2019, Radford et al. adopted Pre-LN for GPT-2. In 2020, Xiong et al. published "On Layer Normalization in the Transformer Architecture", providing formal mathematical proof that Pre-LN preserves gradient scale across arbitrary depth.
- Clean the spine; normalize the branch: Keep the residual stream pure addition; put LayerNorm inside the branch.
- Pre-LN for painless scaling: Warm-up becomes optional; deep networks train without divergence.
How do we derive the exact forward signal accumulation and prove gradient scale invariance across a full Pre-LN transformer block?
Think of a ledger in a joint family bank account. The initial balance is \1,000. In a Pre-LN system, each family member consults a standardized budget chart (LN) before spending, but when they deposit their earnings, they add raw cash directly into the bank account (x + F(x)). The bank's main ledger is never subjected to arbitrary re-indexing or percentage cuts; every dollar deposited in month 1 is still sitting in the account in month 12.
- Learning goalProve that Pre-LN architectures maintain O(1) backward gradient norms across depth, and trace forward signal accumulation across stacked blocks.
- Helpful foundationsPre-LN architecture concepts · Matrix norm bounds and asymptotic notation
- Common trapThinking that Pre-LN eliminates the need for any final normalization. Because residual additions accumulate signal across L layers, a final LayerNorm before the vocabulary unembedding head is strictly necessary.
- Next connectionNext, inspect the full macro-architecture and parameter accounting of GPT-2 124M in Chapter 24 Screen 1.
More context & mechanics
#### 1. Forward Signal Accumulation Calculus
Let x₀ be the token embedding. In a Pre-LN transformer with L blocks:
x_l = x₀ + ∑_(i=0)^(l-1) MHAᵢ(LN(xᵢ)) + ∑_(i=0)^(l-1) FFNᵢ(LN(xᵢ^((1))))
Because every sublayer output is added directly to x, the expected magnitude of x_l grows with depth:
E[\|x_l\|²] = O(l)
To prevent the final linear unembedding layer from receiving huge magnitudes, modern architectures apply one final LayerNorm immediately before the output projection:
y_(final) = LayerNorm(x_L) W_(vocab)
Let L be the training loss. The gradient with respect to layer l is:
(∂ L) / (∂ x_l) = (∂ L) / (∂ x_L) (∂ x_L) / (∂ x_l) = (∂ L) / (∂ x_L) ( I + ∑_(i=l)^(L-1) (∂ Sublayerᵢ) / (∂ x_l) )
Notice that in Pre-LN:
(∂ Sublayerᵢ) / (∂ x_l) = (∂ Sublayerᵢ) / (∂ LN(xᵢ)) · (∂ LN(xᵢ)) / (∂ x_l) = O( (1) / (√(i)) )
As depth i grows, the gradient contributed by the sublayer scales down, while the identity matrix I contributes an unbroken O(1) signal:
lim_(L → ∞) \| (∂ L) / (∂ x_l) \| = O(1)
In Post-LN, by contrast, the gradient scales as O(1 / L), vanishing exponentially. Pre-LN mathematically guarantees O(1) gradient scale invariance across all layers! \blacksquare
Let L = 2 blocks.
- x₀ = [1.0, 1.0].
- Sublayer 1 computes update Δ₁ = [0.1, 0.2] ⟹ x₁ = [1.1, 1.2].
- Sublayer 2 computes update Δ₂ = [0.3, -0.1] ⟹ x₂ = [1.4, 1.1].
Backward pass:
- Suppose incoming loss gradient ∇_(x₂) L = [2.0, 3.0].
- The identity pathway passes [2.0, 3.0] · I = [2.0, 3.0] directly to x₁ and x₀.
- Even if all weights in block 2 have zero gradients, x₀ receives the full [2.0, 3.0] signal.
In the official PyTorch `nn.TransformerEncoderLayer`, setting `norm_first=True` activates this exact Pre-LN formulation. Failure to set `norm_first=True` defaults to legacy Post-LN, which frequently fails to converge on deep architectures.
Xiong et al. (2020) proved this exact O(1) vs O(1/L) theorem, demonstrating why learning rate warm-up was required for Post-LN (to prevent early divergent gradients) but unnecessary for Pre-LN.
- O(1) gradient survival: In Pre-LN, the gradient norm is invariant to network depth.
- `norm_first=True`: The single flag that saves deep transformer training from collapse.
What is the complete tensor shape ledger as activations pass through every sub-layer of a Transformer Block?
Like an accountant keeping a double-entry ledger of every cent flowing in and out of a bank, an AI engineer must track every dimension of a tensor through GPU memory. At the entrance and exit of the block, the tensor balance is always exactly (B, T, 768). Deep inside the attention mechanism, it unfolds into a 4D tensor (B, 12, T, 64). Deep inside the FFN, it broadens into (B, T, 3072). But like a well-behaved guest, it always returns to its original shape before leaving the room.
- Learning goalRecord and verify every intermediate tensor shape across all sub-layers of a standard GPT-2 124M Transformer block.
- Helpful foundationsAttention Is All You Need (Vaswani et al., 2017) · Language Models are Unsupervised Multitask Learners / GPT-2 (Radford et al., 2019) · GPU Tensor Memory Layout & Shape Debugging
- Common trapAssuming tensor shapes change between consecutive Transformer blocks. The block has strictly identical input and output dimensions (B, T, d_model), enabling arbitrary stacking depth.
- Next connectionWith the Transformer Block fully assembled and verified, we examine how GPT-2 124M constructs its initial input tensor from discrete token and positional embeddings.
More context & mechanics
Here is the definitive Shape Ledger for a standard GPT-2 Small (124M) block where B = Batch Size, T = Sequence Length ≤ 1024, d = 768, h = 12, dₖ = 64, and d_(ff) = 4d = 3072:
| Step | Operation | Input Shape | Output Shape | Notes |
| :--- | :--- | :--- | :--- | :--- |
| 0 | Block Input x | — | (B, T, 768) | Native model dimension |
| 1 | LN₁(x) | (B, T, 768) | (B, T, 768) | Mean/var over dim=-1 |
| 2 | Q, K, V Projections | (B, T, 768) | (B, T, 768) | Linear transform |
| 3 | Reshape & Transpose | (B, T, 768) | (B, 12, T, 64) | Split into h=12 heads |
| 4 | Scaled Dot Products | Q × K^\top | (B, 12, T, T) | Score matrix |
| 5 | Causal Mask + Softmax | (B, 12, T, T) | (B, 12, T, T) | Attention weights A |
| 6 | Value Aggregation | A × V | (B, 12, T, 64) | Blended head values |
| 7 | Transpose & Concat | (B, 12, T, 64) | (B, T, 768) | Flatten heads |
| 8 | Output Proj W_O | (B, T, 768) | (B, T, 768) | Mix head features |
| 9 | First Residual Add | (B, T, 768) + (B, T, 768) | (B, T, 768) | x^((1)) = x + Δ₁ |
| 10 | LN₂(x^((1))) | (B, T, 768) | (B, T, 768) | Second Pre-LN |
| 11 | FFN Expansion W₁ | (B, T, 768) | (B, T, 3072) | 4x expansion |
| 12 | GELU Activation | (B, T, 3072) | (B, T, 3072) | Smooth gating |
| 13 | FFN Projection W₂ | (B, T, 3072) | (B, T, 768) | Contract back to d |
| 14 | Second Residual Add | (B, T, 768) + (B, T, 768) | (B, T, 768) | x_(out) = x^((1)) + Δ₂ |
This shape ledger is the primary diagnostic reference when debugging CUDA memory leaks, setting up distributed tensor parallelism (Megatron-LM), or writing custom Triton kernels.
The strict preservation of input-output shapes is the defining architectural virtue of the Transformer, borrowed from the modular block philosophy established in ResNets (He et al., 2015).
- Golden Axiom: "Enters as (B, T, 768), leaves as (B, T, 768)."
- The block flexes internally into 4D attention and 4x FFN, but its external interface is always strictly (B, T, d).
How does The Complete Transformer Block: Pre-LN Layer Topology operate mathematically at scale?
Vertical execution flow through one complete block:
- Learning goalMaster the algebraic and geometric derivation of The Complete Transformer Block: Pre-LN Layer Topology.
- Helpful foundationsFoundations of Build Llm From Scratch · Coordinate Invariance & Geometric Projections · Radford et al. (2019), Language Models are Unsupervised Multitask Learners, OpenAI
- Common trapAssuming The Complete Transformer Block: Pre-LN Layer Topology is an arbitrary computational trick rather than an invariant spatial law.
- Next connectionConnects upstream foundational lemmas to modern high-dimensional computational systems.
More context & mechanics
A complete layer anatomy blueprint in square 1:1 format on dark slate (#0f172a) with modular containers.
Top section: Title in white sans-serif: "The Complete Transformer Block Architecture". Subtitle in muted cyan: "Detailed data flow through Pre-LayerNorm, Causal Multi-Head Attention, and MLP."
Middle section: Vertical execution flow through one complete block:
- Input Tensor x ∈ R^(B × L × D).
- Fork to Residual Shortcut 1.
- LayerNorm / RMSNorm.
- Causal Multi-Head Self-Attention (with optional RoPE).
- Join Residual Shortcut 1: x₁ = x + MHA(LN(x)).
- Fork to Residual Shortcut 2.
- LayerNorm / RMSNorm.
- Feed-Forward Network: Linear(D → 4D) → GELU/SwiGLU → Linear(4D → D).
- Join Residual Shortcut 2: x_(out) = x₁ + FFN(LN(x₁)).
Bottom section: Dimension invariance: "Input shape (B, L, D) strictly equals output shape (B, L, D), allowing dozens of identical blocks to be chained sequentially."
Pristine block diagrams, clearly distinguished residual lines, elegant dark theme.
Every transformation preserves structural invariants while mapping coordinate spaces. When operating in high dimensions, numerical precision and computational efficiency require adhering strictly to canonical mathematical laws.
Formal derivation and peer-reviewed analysis:
Radford et al. (2019), Language Models are Unsupervised Multitask Learners, OpenAI
"The Complete Transformer Block: Pre-LN Layer Topology: Invariance, symmetry, and structure govern mathematical truth across continuous and discrete domains."