Residual connections and layer norm placement

Overview

Residual connections (skip connections) and layer normalization placement are critical architectural choices in Transformers that determine training stability, gradient flow, and final performance. The interaction between these two mechanisms is non-obvious: where you place LayerNorm relative to the residual path fundamentally changes how gradients propagate and how the model learns.

Figure — Residual connections and layer norm placement

Core Concepts

Definition Layer Normalization

Layer normalization normalizes activations across the feature dimension for each example independently:

$\text{LayerNorm}(x) = \gamma \odot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta$

where $\mu = \frac{1}{d}\sum_{i=1}^{d} x_i$ and $\sigma^2 = \frac{1}{d}\sum_{i=1}^{d}(x_i - \mu)^2$ are computed per-example across the $d$ features. $\gamma$ and $\beta$ are learnable scale and shift parameters.

Why Layer Norm, not Batch Norm? Transformers process variable-length sequences. Batch statistics would depend on padding and sequence length distribution—Layer Norm depends only on the individual example, making it stable across different sequence lengths.

The Two Architectures: Post-LN vs Pre-LN

Post-LN (Original Transformer, 2017)

Formula Post-LN Architecture Derivation

Why this structure? The original "Attention is All You Need" paper placed LayerNorm after the residual addition. Let's derive what this means for gradient flow.

For a single sublayer: $x_{i+1} = \text{LayerNorm}(x_i + \text{Sublayer}(x_i))$

Step 1: Compute the sublayer transformation $h_i = \text{Sublayer}(x_i) \quad \text{(attention or FFN)}$

Step 2: Add residual connection $z_i = x_i + h_i$

Step 3: Apply LayerNorm $x_{i+1} = \text{LayerNorm}(z_i) = \gamma \odot \frac{z_i - \mu(z_i)}{\sqrt{\sigma^2(z_i) + \epsilon}} + \beta$

Gradient flow analysis: During backprop, the gradient must pass through LayerNorm's normalization. The Jacobian of LayerNorm has bounded entries (due to the division by $\sigma$ ), but this creates a bottleneck: large updates to $z_i$ get dampened by the normalization.

Critical issue: At initialization, if $\text{Sublayer}(x_i)$ produces large values, the residual sum $z_i$ has high variance. LayerNorm will normalize this, but the gradients become sensitive to the scale of activations. This is one reason Post-LN training is delicate and typically relies on careful initialization and learning-rate warmup.

Post-LN structure for a full layer (attention + FFN):

\begin{align} y &= \text{LayerNorm}(x + \text{MultiHeadAttention}(x)) \\ z &= \text{LayerNorm}(y + \text{FFN}(y)) \end{align}

Worked example Post-LN Training Dynamics

Scenario: Training a deep Post-LN Transformer (e.g. 12+ layers) from scratch.

What is generally observed (from the analysis in Xiong et al., 2020, "On Layer Normalization in the Transformer Architecture"):

Early training is fragile: Without warmup, Post-LN training can diverge, because at initialization the expected gradient magnitude in the deeper layers is large and uneven across depth.
Warmup is important: Post-LN Transformers are typically trained with a learning-rate warmup schedule (ramping the LR up from a small value over some thousands of steps) precisely to avoid this early instability.
Gradient imbalance across depth: Theoretical and empirical analyses show that Post-LN gradients are not uniform across layers at initialization, which is what makes the warmup necessary.

Why this step? At initialization the sublayer outputs have essentially random scale. LayerNorm-after-residual makes the backward gradients scale-sensitive and uneven across depth. Warmup lets the model settle into a stable activation regime before taking large steps.

Pre-LN (Modern Standard)

Formula Pre-LN Architecture Derivation

Why this structure? Pre-LN applies LayerNorm before the sublayer, inside the residual branch. This design (analyzed formally by Xiong et al., 2020) became the standard for large models because it improves gradient behavior at initialization.

For a single sublayer: $x_{i+1} = x_i + \text{Sublayer}(\text{LayerNorm}(x_i))$

Step 1: Apply LayerNorm to the input $\tilde{x}_i = \text{LayerNorm}(x_i)$

Why this step? Normalize the input to the sublayer so it receives controlled, zero-mean, unit-variance activations regardless of what happened in previous layers.

Step 2: Compute the sublayer transformation on normalized input $h_i = \text{Sublayer}(\tilde{x}_i)$

Step 3: Add residual connection (no LayerNorm here!) $x_{i+1} = x_i + h_i$

Gradient flow analysis: The gradient during backprop flows through the residual path with a clean "+1" term in the Jacobian: $\frac{\partial x_{i+1}}{\partial x_i} = I + \frac{\partial h_i}{\partial x_i}$

The LayerNorm is inside the branch, so it doesn't bottleneck the main gradient highway. This creates much more stable, more depth-uniform gradient flow—the gradient can flow largely unchanged through the residual path while the normalized branch gradients add to it.

Key advantage: The residual stream $x_i$ is never normalized, so it accumulates information additively across all layers. This creates a "residual stream" interpretation: each layer reads from and writes to a shared information stream.

Pre-LN structure for a full layer:

\begin{align} y &= x + \text{MultiHeadAttention}(\text{LayerNorm}(x)) \\ z &= y + \text{FFN}(\text{LayerNorm}(y)) \end{align}

Worked example Pre-LN Training Dynamics

Scenario: A deep Pre-LN Transformer of the same depth. What is generally observed:

More stable early training: Loss tends to decrease smoothly from early in training, and the model is far more robust to the choice of learning rate.
Warmup can be relaxed but is NOT eliminated: Pre-LN reduces the dependence on warmup, but large modern Pre-LN models (GPT-3, PaLM, etc.) still use learning-rate warmup schedules for stable large-scale training. The correct statement is "Pre-LN is more forgiving of warmup choice," not "Pre-LN needs no warmup."
More uniform gradient norms across depth: Xiong et al. show the expected gradient magnitude is much better behaved across layers, which is the mechanistic reason Pre-LN scales to very deep networks.

Why this works: LayerNorm normalizes the input to each sublayer, so every layer receives well-behaved activations. The residual path provides a clean gradient highway. The model is much less sensitive to initialization and warmup details.

Practical difference: Training very deep, very large models (GPT-3, PaLM) in practice uses Pre-LN. Post-LN becomes progressively harder to train stably as depth grows without careful tuning.

Mathematical Comparison: Gradient Flow

Formula Gradient Flow Derivation

Let's derive exactly how gradients flow in both architectures. Suppose we have loss $L$ and want $\frac{\partial L}{\partial x_i}$ given $\frac{\partial L}{\partial x_{i+1}}$ .

Post-LN Gradient:

$x_{i+1} = \text{LN}(x_i + F(x_i))$ By chain rule: $\frac{\partial L}{\partial x_i} = \frac{\partial L}{\partial x_{i+1}} \frac{\partial x_{i+1}}{\partial (x_i + F(x_i))} \left(I + \frac{\partial F}{\partial x_i}\right)$

The Jacobian of LayerNorm $\frac{\partial \text{LN}}{\partial z}$ depends on the statistics of $z = x_i + F(x_i)$ : $\frac{\partial \text{LN}(z)}{\partial z_j} = \frac{\gamma}{\sigma}\left(\delta_{ij} - \frac{1}{d} - \frac{(z_i - \mu)(z_j - \mu)}{\sigma^2}\right)$

Key observation: This Jacobian has $O(1/d)$ off-diagonal terms and depends on the variance $\sigma^2$ of $z$ . If $F(x_i)$ produces large outputs, $\sigma$ is large, and the gradient is scaled down by $1/\sigma$ .

Pre-LN Gradient:

$x_{i+1} = x_i + F(\text{LN}(x_i))$

By chain rule: $\frac{\partial L}{\partial x_i} = \frac{\partial L}{\partial x_{i+1}} \left(I + \frac{\partial F}{\partial \text{LN}(x_i)} \frac{\partial \text{LN}(x_i)}{\partial x_i}\right)$

Critical difference: The gradient $\frac{\partial L}{\partial x_{i+1}}$ flows directly through the identity term $I$ without any normalization. The LayerNorm Jacobian only affects the branch gradient, not the main pathway.

Why this matters: In a 24-layer Post-LN Transformer the main gradient must pass through a LayerNorm Jacobian at every layer, whose scaling depends on the activation variance $\sigma$ at each layer. In Pre-LN, the main gradient path is a simple additive accumulation—no per-layer multiplicative normalization on the highway.

Qualitative estimate: Xiong et al. (2020) show that in Pre-LN the expected gradient norm at initialization is essentially independent of depth, whereas in Post-LN it grows with depth in a way that makes the deeper layers dominate—this depth-dependence is exactly what warmup compensates for. (Avoid quoting a specific decay constant like $0.8^L$ ; the precise factor depends on architecture and initialization and is not a universal number.)

Practical Implications

Worked example Real-World Architecture Choices

Original Transformer / early BERT / GPT-2: Used Post-LN. These are relatively shallow (12–24 layers) and are trained with learning-rate warmup and gradient clipping.

GPT-3 (2020) and later large models: Use Pre-LN, which is far more stable when scaling to very deep networks and very large parameter counts. These models still use warmup schedules—Pre-LN reduces sensitivity to warmup, it does not remove the need for it.

T5 (2019): Uses a Pre-LN-style placement with standard LayerNorm (T5 uses a LayerNorm variant that omits the bias/mean-subtraction step in some implementations, but the widely stated fact is that it uses LayerNorm, not RMSNorm).

PaLM (2022): Pre-LN with SwiGLU activations in the FFN. The Pre-LN structure makes it easier to experiment with new activation functions without hitting training instability, and it is trained with a warmup schedule.

Common mistake Common Misconceptions About Layer Norm Placement

Mistake 1: "Post-LN normalizes the output, which should make training more stable."

Why it feels right: Normalization after the residual add seems like it would control activation scales better—you're normalizing the final output of each layer.

The fix: This intuition is backward! Normalizing the output means gradients must flow through the normalization during backprop, on the main pathway. This creates a bottleneck and depth-dependent gradient imbalance. Pre-LN normalizes the input to each sublayer (where it actually helps) while leaving the main gradient path clean.

Steel-man: Post-LN does control forward activation scales tightly. But in deep learning, gradient flow matters more than forward activation control. Pre-LN sacrifices some forward activation control (the residual stream is unnormalized) for much better gradient flow.

Mistake 2: "Pre-LN and Post-LN should converge to the same solution since they're just rearranging the same operations."

Why it feels right: Both architectures have the same LayerNorm and residual connections, just in different orders. Surely they express the same function class?

The fix: The order matters profoundly for optimization, not just expressiveness. Post-LN creates a difficult optimization landscape with depth-dependent gradients at initialization. Pre-LN creates a smoother landscape. They may express similar functions, but one is far easier to optimize—especially at large depth.

Mistake 3: "Pre-LN needs no learning-rate warmup at all."

Why it feels right: Pre-LN clearly makes training more stable and less warmup-sensitive, so it's tempting to conclude warmup is unnecessary.

The fix: Pre-LN reduces the dependence on warmup and can sometimes train without it at small scale, but large models (GPT-3, PaLM, etc.) still use warmup schedules in practice. The accurate claim is "Pre-LN is more robust to warmup choice," not "warmup is unnecessary."

Mistake 4: "The residual connection adds $x$ , so the output scale grows linearly with depth."

Why it feels right: In Pre-LN, we keep adding to $x$ : after $L$ layers, shouldn't the scale be $O(L)$ times larger?

The fix: The variance grows roughly linearly, not the scale. If each sublayer adds a roughly independent contribution with variance $\sigma^2$ , the total variance after $L$ layers is $O(L\sigma^2)$ , so the standard deviation grows as $O(\sqrt{L})$ . This is manageable—and can be controlled by initializing sublayers with small output variance (scaling down the final projection).

This is actually a feature: The residual stream's variance growth is gradual ( $\sqrt{L}$ ). In Post-LN, LayerNorm keeps variance at roughly 1 at each layer, which sounds good but prevents the model from using different activation scales for different information across depth.

Advanced: The Residual Stream Interpretation

Intuition Residual Stream as Information Highway

In Pre-LN Transformers, we can think of the residual path $x$ as a stream that carries information across all layers:

Layer 1 reads from the stream, computes attention, writes back to the stream.
Layer 2 reads the updated stream, computes attention on the updated info, writes back.
...
Layer 24 reads the fully processed stream, outputs predictions.

Each layer performs a "read-modify-write" operation. The stream never gets reset or normalized—it's a persistent memory. This interpretation explains why Pre-LN works: the gradient can flow backward through this stream with no bottlenecks.

Analogy: Imagine a relay race where runners pass a baton (the residual stream). In Pre-LN, each runner adds their contribution and passes it on—the baton accumulates information. In Post-LN, each runner passes the baton through a "reset station" (LayerNorm) that standardizes it—this prevents information accumulation and creates gradient issues.

Implementation Details

Worked example Pre-LN Transformer Block (PyTorch-style pseudocode)

class PreLNTransformerBlock:
   def forward(self, x):
       # Self-attention with Pre-LN
       residual = x
       x = layer_norm_1(x)              # Normalize before attention
       x = multi_head_attention(x)      # Attention on normalized input
       x = residual + x                 # Add to unnormalized residual
       # Feedforward with Pre-LN
       residual = x
       x = layer_norm_2(x)              # Normalize before FFN
       x = feedforward(x)               # FFN on normalized input
       x = residual + x                 # Add to unnormalized residual
       return x

Why this step? The pattern is: x = x + sublayer(layernorm(x)). The residual is always added to the unnormalized $x$ , and LayerNorm is applied to a copy that goes into the sublayer.

Initialization detail: A common trick is to initialize the final linear projection in each sublayer (attention output projection and FFN second layer) with a small scale (e.g. proportional to $1/\sqrt{2L}$ for $L$ layers) so the residual stream variance stays roughly controlled across depth.

Connections

4.1.01-Self-attention-mechanism: Residual connections wrap around self-attention blocks
4.1.08-Multi-head-attention: Each head's output is projected and added to the residual stream
4.2.02-Training-stability-and-convergence: Pre-LN dramatically improves training stability in deep models
4.3.01-Positional-encoding: The residual stream carries positional information across layers
5.1.03-Gradient-flow-in-deep-networks: Core gradient flow concepts that explain why Pre-LN works
6.2.05-Layer-scaling-and-initialization: Proper initialization for residual connections
4.1.12-Encoder-decoder-attention: Encoder-decoder uses same Pre-LN pattern for cross-attention

Active Recall

Recall Explain to a 12-year-old: Why do we need skip connections?

Imagine you're playing a game of telephone with 24 people. Normally, the message gets garbled by the time it reaches the end—each person might mishear and change it slightly.

Now imagine if the first person could also whisper directly to the last person, and everyone in between also gets to whisper to the last person. The last person hears the original message clearly plus all the changes everyone made along the way. That's a skip connection!

In neural networks, skip connections let the original information travel directly to later layers without getting "garbled" by all the math in between. This helps the network learn better because it can compare the original input to the processed version and learn what changes actually helped.

#flashcards/ai-ml

What is the mathematical form of a residual connection in Transformers?

\text{output} = \text{input} + \text{Sublayer}(\text{input})

, where the input is added directly to the transformed output, creating an identity gradient path.

Why does Pre-LN provide better gradient flow than Post-LN?

In Pre-LN, gradients flow through the residual path with a clean identity term (

I + ...

) on the main pathway, while in Post-LN gradients must pass through LayerNorm's Jacobian at every layer, creating depth-dependent, scale-sensitive bottlenecks.

Write the Pre-LN architecture formula for a Transformer layer

y = x + \text{Attention}(\text{LN}(x))

then

z = y + \text{FFN}(\text{LN}(y))

, where LayerNorm is applied before each sublayer inside the residual branch.

Write the Post-LN architecture formula for a Transformer layer

y = \text{LN}(x + \text{Attention}(x))

then

z = \text{LN}(y + \text{FFN}(y))

, where LayerNorm is applied after the residual addition.

What does the "+1" term in the residual gradient

\frac{dy}{dx} = 1 + \frac{dF}{dx}

guarantee?

It guarantees at least an identity gradient flows backward, preventing vanishing gradients even if the sublayer

F

learns nothing or has near-zero gradients initially.

Does Pre-LN eliminate the need for learning-rate warmup?

No. Pre-LN reduces sensitivity to warmup, but large models like GPT-3 and

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Dekho, yahan pe do cheezein samajhna zaroori hai jo Transformer ko itna powerful banati hain. Pehli cheez hai residual connection, jo basically ek shortcut ya "highway" hai. Jab aap bahut saari layers stack karte ho (jaise 24 layers), toh training ke time gradient ko peeche tak flow karna padta hai har layer ke through, aur har multiplication mein woh gradient ya toh bahut chota ho jaata hai ya bahut bada — dono cases mein network theek se seekh nahi paata. Residual connection ka formula simple hai: y = x + F(x). Iska matlab hai ki input x ko direct output mein add kar dete hain. Iska magic yeh hai ki jab derivative nikalte ho toh woh 1 + dF/dx ban jaata hai — yeh "+1" guarantee karta hai ki gradient kam se kam identity ke roop mein toh peeche flow karega hi, chahe layer shuru mein kuch bhi na seekhe. Isliye deep networks train ho paate hain bina gradient vanish hue.

Doosri cheez hai Layer Normalization, jo activations ko stable rakhta hai. Yeh har example ke features ko normalize karta hai — matlab mean nikaalta hai, variance nikaalta hai, aur values ko ek consistent scale pe le aata hai. Ek important baat: Transformers mein hum Batch Norm ki jagah Layer Norm use karte hain, kyunki humaari sequences alag-alag length ki hoti hain (kabhi chhota sentence, kabhi lamba). Agar batch ke statistics use karte toh padding aur length ke variation se sab kuch unstable ho jaata. Layer Norm sirf individual example pe depend karta hai, isliye woh reliable rehta hai.

Ab asli intuition yeh hai ki LayerNorm ko kahan rakhna hai — residual connection ke pehle (Pre-LN) ya baad mein (Post-LN) — yeh ek badi design decision hai. Original 2017 paper mein Post-LN use hua tha, jahan normalization residual add karne ke baad hota hai. Problem yeh hai ki isse gradient flow thoda delicate ho jaata hai, aur model ko train karne ke liye careful initialization aur learning-rate warmup ki zaroorat padti hai, warna training unstable ho jaati hai. Yeh baat matter karti hai kyunki jab aap real mein bade models banaoge, toh galat placement ki wajah se training crash ho sakti hai ya model converge hi nahi karega. Isliye yeh chhoti si placement wali choice actually decide karti hai ki aapka Transformer smoothly train hoga ya nahi.

Residual connections and layer norm placement

Overview

Core Concepts

The Two Architectures: Post-LN vs Pre-LN

Post-LN (Original Transformer, 2017)

Pre-LN (Modern Standard)

Mathematical Comparison: Gradient Flow

Practical Implications

Advanced: The Residual Stream Interpretation

Implementation Details

Connections

Active Recall

Concept Map

Hinglish (regional understanding)

Hinglish (regional understanding)

Go deeper — visual, from zero

Test yourself — Transformer Architecture

Connections