Residual Connections & Layer Normalization

A residual connection (or skip connection) wraps a sublayer $F$ (attention or feed-forward) so that the block computes

instead of $y = F(x)$ directly. The input $x$ is added back to the sublayer's output, forming a shortcut path that bypasses $F$ entirely. This means $F$ only needs to learn the residual — the difference between what the input already represents and what the output should represent — rather than reconstructing the entire representation from nothing.

Layer normalization (LayerNorm) is a normalization technique applied per token (per position in the sequence), independent of other tokens in the batch. For a single token's activation vector $x \in \mathbb{R}^d$, LayerNorm computes

where $\gamma, \beta \in \mathbb{R}^d$ are learned scale and shift parameters, and $\epsilon$ is a small constant preventing division by zero. Unlike batch normalization, LayerNorm's statistics are computed across the feature dimension of a single example, not across the batch dimension — this makes it insensitive to batch size and well-suited to sequences of varying length, since each token is normalized using only its own $d$ features. The interaction of these two components inside a Transformer sublayer is called a residual block, and its exact wiring — whether normalization happens before or after the residual add — is a design choice with real consequences, discussed next.

A Transformer layer contains two sublayers: multi-head attention and a position-wise feed-forward network. Each is wrapped in a residual-plus-norm block. There are two standard arrangements.

Post-LN (original Transformer, 2017). Normalize after adding the residual:

Pre-LN (modern default, used in GPT-2 and later). Normalize before feeding into the sublayer, and add the raw residual afterward:

The pre-LN arrangement keeps an unbroken identity path from the very first layer to the very last — the residual stream $x$ is never itself passed through a normalization layer before being added back, only a normalized copy is fed to the sublayer. This detail turns out to matter enormously for gradient flow at depth (see Complexity & Correctness). Each sublayer in practice also uses dropout on the sublayer output before the residual add, as a regularizer: ``` Sublayer output → Dropout → Add to residual stream → (Post-LN: Normalize here) ```

Pseudocode: a full pre-norm Transformer block

``` function TransformerBlockPreLN(x, W_attn, W_ffn, gamma1, beta1, gamma2, beta2):

# x: matrix of shape (seq_len, d_model), one row per token

# --- Sublayer 1: Multi-Head Attention --- normed1 = LayerNorm(x, gamma1, beta1) # normalize per token attn_out = MultiHeadAttention(normed1, W_attn) # see Multi-Head Attention attn_out = Dropout(attn_out, p=dropout_rate) x = x + attn_out # residual connection

# --- Sublayer 2: Position-wise Feed-Forward --- normed2 = LayerNorm(x, gamma2, beta2) ffn_out = FeedForward(normed2, W_ffn) # e.g. Linear -> ReLU -> Linear ffn_out = Dropout(ffn_out, p=dropout_rate) x = x + ffn_out # residual connection return x

function LayerNorm(x, gamma, beta, eps=1e-5):

# x: (seq_len, d_model); normalize each row independently for each token row t in x: mu = mean(t) # scalar var = mean((t - mu)^2) # scalar (biased estimator) t_hat = (t - mu) / sqrt(var + eps) # normalized row t_out = gamma t_hat + beta # elementwise affine, gamma,beta in R^d output_row[t] = t_out return output_row ``` The post-LN variant is identical except the normalization call moves after* the residual add: ``` function TransformerBlockPostLN(x, W_attn, W_ffn, gamma1, beta1, gamma2, beta2): attn_out = MultiHeadAttention(x, W_attn) attn_out = Dropout(attn_out, p=dropout_rate) x = LayerNorm(x + attn_out, gamma1, beta1) # normalize AFTER residual add ffn_out = FeedForward(x, W_ffn) ffn_out = Dropout(ffn_out, p=dropout_rate) x = LayerNorm(x + ffn_out, gamma2, beta2) # normalize AFTER residual add return x ``` A full encoder or decoder stacks $N$ of these blocks, feeding the output of one directly as the input to the next — a design elaborated fully in the Transformer Encoder Stack and Transformer Decoder Stack articles.

Time and space complexity. Both operations are cheap relative to attention. The residual add is elementwise: $O(n \cdot d)$ time and space for a sequence of length $n$ and hidden dimension $d$. LayerNorm computes a mean and variance per token, each an $O(d)$ reduction, over $n$ tokens: $O(n \cdot d)$ time, $O(n \cdot d)$ space for the normalized output (or $O(1)$ extra space beyond the input if done in place). Compare this to self-attention's $O(n^2 d)$ time — residual and norm operations are asymptotically negligible, adding no new bottleneck to the architecture's overall complexity, which remains dominated by attention.

Why residual connections fix gradient flow: an invariant argument. Consider a stack of $N$ residual blocks in pre-LN form, $x_{l+1} = x_l + F_l(\mathrm{LN}(x_l))$. By the chain rule, the gradient of the loss $\mathcal{L}$ with respect to an early layer's input $x_l$ is

Each factor is the identity matrix $I$ plus a Jacobian term. Expanding the product, the identity terms alone contribute $\partial \mathcal{L}/\partial x_N$ directly to $\partial \mathcal{L}/\partial x_l$ — an unattenuated path exists regardless of how small or ill-conditioned the sublayer Jacobians become. Without the residual connection ($x_{l+1} = F_l(x_l)$ alone), the same product would consist purely of sublayer Jacobians, and if each has, say, spectral norm $< 1$, the product shrinks geometrically with depth — the vanishing gradient problem. The identity term structurally guarantees a gradient floor, which is the essence of why residual connections make arbitrarily deep networks trainable: the network is never forced to lose the input signal, only permitted to add to it.

Why depth cannot hurt (in principle). With residual connections, if the ideal function for a block is the identity, $F$ only needs to learn $F(x) = 0$, which is trivially reachable (e.g., zero-initialized output weights). This means adding a residual block can only make the hypothesis class the network can represent a superset of what a shallower network could represent — the identity is always achievable, so a deeper residual network is never forced to represent a strictly worse function than a shallower one. Without residuals, learning an identity mapping through several nonlinear layers is nontrivial, since generic nonlinear layers do not naturally sit near identity at initialization.

Why LayerNorm stabilizes optimization. LayerNorm rescales each token vector to have unit variance across features at every layer, which keeps the input distribution to each sublayer within a numerically stable range independent of depth. This mitigates the internal covariate shift problem: without renormalization, activation magnitudes can compound multiplicatively across layers (either shrinking toward zero or exploding), making learning rates that work for one layer inappropriate for another. Because LayerNorm's statistics are computed per token rather than per batch, this stabilization is exact and deterministic at both training and inference time — there is no dependency on batch statistics that must be estimated and can behave differently at test time (a known pitfall of batch normalization in sequence models).

Pre-LN vs. post-LN: a real tradeoff. Post-LN, as used in the original Transformer, applies LayerNorm to the sum $x + F(x)$, meaning the residual stream itself is renormalized every layer — the identity shortcut is not perfectly clean, since normalization is a nonlinear (though gentle) operation inserted directly into the gradient path back through the residual stream. Empirically and theoretically, this causes gradients at initialization to grow with depth in post-LN Transformers, requiring learning-rate warmup to train stably beyond a modest number of layers. Pre-LN keeps the residual stream itself untouched by any normalization — LayerNorm only appears on the branch feeding into each sublayer — preserving the clean identity path derived above and enabling stable training without warmup, even for very deep stacks (dozens to over a hundred layers). The tradeoff: pre-LN stacks can suffer from representation collapse in extremely deep settings (very late layers contribute a shrinking relative fraction of the final normalized output) and sometimes require an additional final LayerNorm after the last block to control output scale — a detail present in GPT-2's architecture.

Trace one pre-LN block on a toy input with $d_{model} = 4$, one token, to make every number concrete. Let the token's residual-stream vector be

Step 1 — LayerNorm before attention. Assume $\gamma = [1,1,1,1]$, $\beta = [0,0,0,0]$ (identity affine, common at initialization), $\epsilon = 10^{-5}$.

Step 2 — Toy attention sublayer. Suppose (for illustration) the attention sublayer, after self-attention and output projection, produces $F(\hat{x}) = [0.10,\ -0.05,\ 0.20,\ 0.00]$ for this token (a small update, as is typical near initialization).

Step 3 — Residual add. The key operation: add back to the original, un-normalized $x$, not to $\hat x$.

Notice $x'$ is close to the original $x$ — the block made a small, controlled correction rather than replacing the representation. This is exactly the intended residual behavior.

Step 4 — Feed into sublayer 2 (feed-forward). Repeat: normalize $x'$ (mean $\approx 1.0125$, similar rescaling), pass through the feed-forward network to get some small update $G$, and add it to $x'$ (not to the normalized copy) to produce the block's final output.

This trace illustrates the core mechanic: LayerNorm operates on a branch copy used only to compute the sublayer's output, while the actual residual stream accumulates small additive updates across layers, exactly the structure the gradient-flow argument above relies on.

Residual connections originated in computer vision, and layer normalization was developed for recurrent and sequence models; both were combined into the Transformer architecture, whose motivating problem was training arbitrarily deep sequence models without recurrence.

Every major Transformer-based system — encoder-only models like BERT, decoder-only models like the GPT family, and encoder-decoder models like T5 — is built from residual blocks wrapped in layer normalization, stacked dozens to over a hundred times. Without this pairing, none of these models could be trained at their actual depths; the vanishing-gradient problem that plagued deep recurrent networks would reappear immediately in deep attention stacks. Production LLM training pipelines depend on the specific pre-LN vs. post-LN choice to determine whether learning-rate warmup schedules are necessary and how large a model can be scaled before training destabilizes. Variants of this idea — RMSNorm (a simplified LayerNorm without mean-centering), and various residual-scaling schemes — appear in more recent large models seeking marginal efficiency or stability gains, but the underlying principle, an identity shortcut plus per-token normalization, remains unchanged. This foundation is exactly what the Transformer Encoder Stack and Transformer Decoder Stack articles build upon when assembling full multi-layer architectures, and it is inseparable from why Multi-Head Attention can be stacked at all at production scale.

Ask the tutor: Why does pre-LN Transformer training remain stable without learning-rate warmup, while post-LN does not?