Normalization and Activations: RMSNorm, Pre-Norm, and the Gated FFN

1. Quick anchor

Attention gets the headlines, but the reason a modern transformer with 80+ layers trains at all is the boring scaffolding around it: normalization, residuals, and the FFN. A transformer block is x -> x + Attn(Norm(x)) -> x + FFN(Norm(x)). The residual gives gradients a clean highway from the loss all the way back to layer 0. Pre-norm (normalize inside the residual branch, never on the highway itself) is what keeps that highway un-distorted so you can stack depth. RMSNorm is LayerNorm with the mean-subtraction deleted — same stabilization, fewer ops. And the FFN does two-thirds of the parameter work in the model; the modern version is a gated MLP (SwiGLU) that multiplies two projections together for a learned, input-dependent gate. Get these four pieces right and the model converges; get them wrong and a 70-layer model NaNs in the first thousand steps.

2. Why interviewers probe this

• IC4 — do you understand the mechanics? Can you write LayerNorm and RMSNorm from scratch, name the difference (mean term), and explain why a residual connection's "+1" in the gradient is the whole point? Can you state what an FFN block actually computes and why it's d_model -> 4·d_model -> d_model?
• IC5 — do you understand the system tradeoffs? Why is pre-norm the default despite post-norm sometimes scoring higher? Why did the field move LayerNorm -> RMSNorm and ReLU -> SwiGLU — is it quality, speed, or both, and can you quantify it? You should be able to reason about activation memory, kernel fusion, and what breaks at 100 layers vs 12.
• Both levels — taste. These choices are nearly free to get right and catastrophic to get wrong. Interviewers use them to check whether you've actually trained or debugged a large model, or just read about attention. A candidate who shrugs at "where do you put the norm?" has never watched a loss curve diverge.

3. Concept build-up

3.1 LayerNorm: the baseline

LayerNorm normalizes across the feature dimension of a single token (not across the batch — that's BatchNorm, which is brittle for variable-length sequences). For a token vector $x \in \mathbb{R}^d$:

Symbols: $x_i$ is feature $i$ of the token; $\mu$ is the mean over the $d$ features; $\sigma^2$ is the variance over those features; $\epsilon$ (e.g. $10^{-5}$) prevents divide-by-zero; $\gamma, \beta \in \mathbb{R}^d$ are learned scale and shift; $\odot$ is element-wise multiply. The job: keep each token's activations at a controlled scale so the next sublayer never sees exploding or vanishing inputs, which stabilizes training and lets you use higher learning rates.

3.2 RMSNorm: delete the mean

RMSNorm (Zhang & Sennrich, 2019) makes one observation: most of LayerNorm's benefit comes from the rescaling (dividing by a measure of magnitude), not the re-centering (subtracting the mean). So drop $\mu$ entirely:

Here the denominator is the root-mean-square of the features. No $\mu$, no $\beta$ (typically). Why is this safe? A mechanistic finding is that models trained with RMSNorm spontaneously learn representations roughly orthogonal to the uniform (all-ones) vector — and LayerNorm's mean-subtraction is exactly the projection that removes the component along that uniform vector. If the model already avoids putting signal there, subtracting the mean is removing something close to zero: redundant work. Empirically RMSNorm is 7–64% faster than LayerNorm (the range depends on shape and hardware) with no meaningful quality loss — verified across LLaMA, OLMo, and others. That's why it's ubiquitous in 2024–2026 LLMs.

The savings: you skip computing $\mu$, skip the subtraction broadcast, and skip the $\beta$ add. On a memory-bandwidth-bound norm kernel, fewer passes over the activation tensor matters more than the FLOP count.

3.3 Residual connections: why depth is possible at all

The residual (He et al., 2015) is x_{l+1} = x_l + f(x_l). The magic is in the backward pass. Differentiate:

That +1 is everything. In a plain deep network, the gradient at layer $l$ is a long product of Jacobians, $\prod_l \frac{\partial f_l}{\partial x_l}$; if each factor has magnitude < 1, the product decays exponentially with depth and early layers get no signal (vanishing gradients). The residual replaces each factor with (1 + ∂f/∂x_l). Even if the sublayer's Jacobian $\frac{\partial f}{\partial x_l} \to 0$, the gradient still flows through the +1 term — an identity highway from the loss back to layer 0. This is the single reason you can stack 80 transformer blocks instead of 6.

A second view: a residual stack computes an additive refinement. Each block reads the residual stream, writes a small update, and adds it. The stream is a shared scratchpad; layers compose by accumulation, not by repeated multiplication. This framing (the "residual stream") is also how mechanistic interpretability reasons about transformers.

3.4 Pre-norm vs post-norm: where you put the norm decides whether you can go deep

The original transformer (Vaswani et al., 2017) used post-norm: x_{l+1} = LN(x_l + Sublayer(x_l)). The norm sits on the residual highway, after the add. This is a problem: the clean identity path is now wrapped in a normalization at every layer, so the gradient highway is repeatedly rescaled and the effective signal can blow up or shrink as depth grows. Post-norm transformers need careful learning-rate warmup and are finicky past a couple dozen layers.

Pre-norm (Xiong et al., 2020) moves the norm inside the residual branch:

Now the skip path x_l + ... is a pure identity — nothing touches it. The norm only conditions the input that flows into the sublayer. Gradient flow analysis shows pre-norm keeps gradient magnitudes well-behaved across depth without warmup, which is why every large modern LLM (LLaMA, Qwen, Gemma, GPT-style) is pre-norm. The tradeoff: pre-norm lets the residual stream grow in magnitude with depth (each block adds to it, and nothing renormalizes the stream itself), so the final layers' relative contributions shrink and you typically add one final norm before the output head. Post-norm, by contrast, has been reported to give slightly better results in some narrow regimes (e.g. a notable BLEU gain in certain machine-translation setups) precisely because it does renormalize the stream — but it pays for that with training instability you have to fight. At staff scale the verdict is settled: pre-norm, plus a final norm, plus sometimes tricks like scaling residual branches by $1/\sqrt{2N}$ to tame the growing stream.

3.5 Activations: ReLU -> GELU -> SiLU -> gated GLU

The FFN's nonlinearity has drifted over time:

• ReLU: $\max(0, x)$. Cheap, but a hard zero kills gradients for negative inputs and the kink is non-smooth.
• GELU (Hendrycks & Gimpel, 2016): $\text{GELU}(x) = x \cdot \Phi(x)$, where $\Phi$ is the standard-normal CDF. A smooth, probabilistic gate — it scales each input by the probability that a standard Gaussian is below it. Smooth gradients everywhere; used in BERT, GPT-2/3.
• SiLU / Swish: $\text{SiLU}(x) = x \cdot \sigma(x)$ with $\sigma$ the logistic sigmoid. Very close to GELU, slightly cheaper; the base unit inside SwiGLU.

The bigger leap is gating (GLU variants, Shazeer 2020). Instead of one projection through a nonlinearity, split into two projections and multiply:

$W_g$ produces the gate, $W_v$ produces the value; the gate (squashed by SiLU) decides, per dimension and per token, how much of the value passes. GeGLU is identical with GELU instead of SiLU. Because gating consumes the hidden dimension across two matrices, implementations shrink the hidden width to $\frac{2}{3} \cdot 4 d_{model} \approx \frac{8}{3} d_{model}$ to keep parameter count matched to a standard FFN. The payoff is consistently lower perplexity and better downstream scores at equal parameters — small but reliable, which is why SwiGLU/GeGLU is standard in LLaMA, OLMo, Gemma. The cost is one extra matmul (three weight matrices instead of two: $W_g, W_v, W_{out}$).

3.6 The FFN block itself

The FFN (a.k.a. MLP) is applied independently to every token position — no mixing across positions (attention already did that). A standard FFN:

It expands $d \to 4d$, applies the nonlinearity, contracts back. This is where roughly two-thirds of a transformer's parameters live (the $4\times$ expansion dominates attention's projections). Interpretively, the FFN acts as a per-token key-value memory: $W_1$ rows detect features, $W_2$ columns write associated outputs into the residual stream. The gated version replaces $\sigma(xW_1)$ with the SwiGLU expression above.

4. Minimal implementation

A complete, runnable pre-norm transformer block with RMSNorm and a SwiGLU FFN — production-shaped, the way LLaMA-family models are built.

import torch
import torch.nn as nn
import torch.nn.functional as F
 
class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.gamma = nn.Parameter(torch.ones(dim))  # learned scale; no beta, no mean
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # rms over the feature dim; upcast to fp32 for numerical stability
        norm = x.float().pow(2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt()
        return (x.float() * norm).type_as(x) * self.gamma
 
class SwiGLU(nn.Module):
    """Gated FFN. Hidden width shrunk to ~8/3*d to match a standard 4*d FFN's params."""
    def __init__(self, dim: int, hidden: int | None = None):
        super().__init__()
        hidden = hidden or int(8 * dim / 3)
        hidden = 256 * ((hidden + 255) // 256)  # round to a hardware-friendly multiple
        self.w_gate = nn.Linear(dim, hidden, bias=False)
        self.w_value = nn.Linear(dim, hidden, bias=False)
        self.w_out = nn.Linear(hidden, dim, bias=False)
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w_out(F.silu(self.w_gate(x)) * self.w_value(x))
 
class Block(nn.Module):
    def __init__(self, dim: int, n_heads: int):
        super().__init__()
        self.attn_norm = RMSNorm(dim)
        self.ffn_norm = RMSNorm(dim)
        self.attn = nn.MultiheadAttention(dim, n_heads, batch_first=True)
        self.ffn = SwiGLU(dim)
 
    def forward(self, x: torch.Tensor, attn_mask=None) -> torch.Tensor:
        # PRE-NORM: norm is inside the residual branch; the skip path is a clean identity.
        h = self.attn_norm(x)
        x = x + self.attn(h, h, h, attn_mask=attn_mask, need_weights=False)[0]
        x = x + self.ffn(self.ffn_norm(x))
        return x
 
if __name__ == "__main__":
    blk = Block(dim=512, n_heads=8)
    tok = torch.randn(2, 16, 512)           # (batch, seq, dim)
    out = blk(tok)
    print(out.shape)                        # torch.Size([2, 16, 512])
    # sanity: a fresh block is near-identity, so output ≈ input (residual highway)
    print((out - tok).abs().mean().item())  # small number

Three things to notice. First, the upcast to fp32 inside RMSNorm — the squaring can overflow or lose precision in bf16, so production code computes the norm in fp32 then casts back; getting this wrong is a real source of training instability. Second, the clean residual (x = x + sublayer(norm(x))) — the norm never touches the x on the left of the +. Third, rounding the hidden dim to a multiple of 256 so the matmuls hit fast tensor-core paths; the "8/3·d" is theoretical, hardware wants alignment.

5. Production tradeoffs

Cost. The FFN is ~two-thirds of the model's parameters and FLOPs, so the SwiGLU "third matmul" is not free — it's the dominant compute in the block. You recover the parameter count by shrinking hidden width to ≈8/3·d, but you still pay one extra matmul of activation read/writes. At inference this shows up as FFN being the throughput bottleneck, which is why MoE (sparse FFNs) is the next lever: route each token to a few experts so you scale parameters without scaling per-token FFN FLOPs.

Latency. Norm kernels are memory-bandwidth bound, not compute bound — they read and write the whole activation tensor and do little arithmetic. That's why RMSNorm's win is about fewer passes over memory, and why production stacks fuse norm + the following matmul into one kernel to avoid a round-trip to HBM.

Quality and failure modes. The catastrophic failures are at training time: post-norm + no warmup at depth -> divergence; RMSNorm in bf16 without fp32 upcast -> slow instability; forgetting the final norm in a pre-norm model -> a poorly-calibrated output head. None of these show up in a 12-layer toy model — they emerge at scale, which is exactly why interviewers ask.

What changes at scale. Deep pre-norm models grow the residual stream's magnitude with depth; large training runs counter this with branch-scaling (e.g. dividing each sublayer output by $\sqrt{2N}$ for $N$ layers) or with careful initialization so the stream stays in range for the final norm.

6. How it's asked

7. Pitfalls & flashcards

• bf16 RMSNorm without fp32 upcast — the x^2 sum loses precision; compute the norm in fp32, cast back. Silent training instability otherwise.
• Forgetting the final norm in pre-norm models — the residual stream grows with depth; without a norm before the output head, logits are poorly scaled.
• Putting the norm on the residual (post-norm) at depth — works in toy models, diverges at scale without aggressive warmup.
• SwiGLU parameter blowup — three matrices not two; if you don't shrink hidden width to ≈8/3·d you've silently grown the model ~33%.
• Confusing LayerNorm with BatchNorm — LayerNorm normalizes across features of one token; BatchNorm across the batch, which is fragile for variable-length sequences and is not what transformers use.
• Assuming RMSNorm is a quality upgrade — it's a speed upgrade at quality parity, not a quality win. The win is cheaper, not better.

Flashcard. Pre-norm = x + Sublayer(Norm(x)) (clean skip, trains deep). Post-norm = Norm(x + Sublayer(x)) (renormalizes stream, diverges deep). RMSNorm = LayerNorm minus the mean and beta. Residual's gradient +1 is why depth works. FFN = ~2/3 of params; SwiGLU gates it for reliable perplexity gains at ~3 matmuls.

8. Further reading

• Vaswani et al., 2017, Attention Is All You Need — original post-norm transformer and the FFN block (https://arxiv.org/abs/1706.03762).
• Zhang & Sennrich, 2019, Root Mean Square Layer Normalization — RMSNorm derivation and ablations (https://arxiv.org/abs/1910.07467).
• Shazeer, 2020, GLU Variants Improve Transformer — SwiGLU/GeGLU and the 2/3 hidden-width rule (https://arxiv.org/abs/2002.05202).
• Xiong et al., 2020, On Layer Normalization in the Transformer Architecture — the pre-norm vs post-norm gradient analysis (https://arxiv.org/abs/2002.04745).
• Hendrycks & Gimpel, 2016, Gaussian Error Linear Units (GELUs) (https://arxiv.org/abs/1606.08415).
• He et al., 2015, Deep Residual Learning for Image Recognition — where the residual +1 came from (https://arxiv.org/abs/1512.03385).

Next: Positional encodings — RoPE, ALiBi, and length extrapolation.