Layer Normalization: Pre-LN vs Post-LN

In a 32-layer decoder, each token carries its own hidden vector through the stack. The collection of those states is the residual stream. Take one token's state: [3.0, 1.0, -1.0, 5.0]. Attention reads other visible tokens and writes an update; a feed-forward network writes another. Repeated additions can change the stream's scale. We want to control what each branch reads without accidentally blocking useful gradients.

LayerNorm subtracts the token's mean and divides by the square root of its variance plus a small epsilon, then applies a learned scale and shift. Which vector should it normalize: the residual sum after the addition, or the branch input before the addition? Trace those two choices and you get different activations and gradient paths.

Positional techniques like RoPE and ALiBi tell the model where tokens sit relative to one another. Normalization addresses a different question: what scale does each operation read? We'll calculate LayerNorm by hand, contrast it with BatchNorm and RMSNorm, then test what residual placement can and cannot tell us about training stability.

One residual stream, many writes

A plain residual sublayer follows this blueprint, before choosing normalization or residual scaling:

$x_{l+1} = x_l + \text{Sublayer}(x_l)$

The input $x_l$ takes a direct skip path, while the Sublayer computes an attention or feed-forward update. If each sublayer produces perturbations with roughly independent variance $\sigma_{\text{sub}}^2$, the variance of the unnormalized sum climbs steadily with layer depth $l$:

$\operatorname{Var}(x_l) \approx \operatorname{Var}(x_0) + l \cdot \sigma_{\text{sub}}^2$

This is a model of uncorrelated writes, not a law of trained transformers. In general, $\operatorname{Var}(x+\Delta)=\operatorname{Var}(x)+\operatorname{Var}(\Delta)+2\operatorname{Cov}(x,\Delta)$. Correlated, cancelling, or rescaled updates can behave differently. Even linear variance growth means standard deviation grows like $\sqrt{l}$, not like $l$.

In Post-LN, normalization wraps the sum after the write. In Pre-LN, normalization sits on the input to the sublayer branch, leaving the skip path untouched. During backpropagation, Pre-LN's block Jacobian therefore has an additive identity term; Post-LN's includes a normalization Jacobian outside the sum. Later we'll construct a branch that cancels the identity contribution in one direction.^{[1]Reference 1On Layer Normalization in the Transformer Architecture.https://arxiv.org/abs/2002.04745}

We'll start with the bare arithmetic of a single token vector.

LayerNorm by hand

Take our running token vector as the hidden state of one token:

x = [3.0, 1.0, -1.0, 5.0]

First, calculate its mean over the 4 features:

$\mu = \frac{3.0 + 1.0 + (-1.0) + 5.0}{4} = \frac{8.0}{4} = 2.0$

Subtracting this mean centers the features around zero, giving [1.0, -1.0, -3.0, 3.0]. Centering shifts the coordinate origin without altering pairwise distances between features.

Next, calculate the population variance across the centered values:

$\sigma^2 = \frac{1.0^2 + (-1.0)^2 + (-3.0)^2 + 3.0^2}{4} = \frac{1.0 + 1.0 + 9.0 + 9.0}{4} = \frac{20.0}{4} = 5.0$

LayerNorm uses the biased population variance (dividing by $d = 4$) rather than the sample variance ($d - 1$).

For this hand calculation, ignore epsilon. Take the square root: $\sigma = \sqrt{5.0} \approx 2.236$. Dividing each centered feature by $\sigma$ produces:

$\hat{x} = \left[ \frac{1.0}{2.236}, \frac{-1.0}{2.236}, \frac{-3.0}{2.236}, \frac{3.0}{2.236} \right] \approx [0.447, -0.447, -1.342, 1.342]$

Before any learned affine transform, its mean is 0 and its variance is 1 in this epsilon-free calculation. With positive epsilon the variance is slightly smaller here; a constant input has zero normalized variance.

Finally, learned scale ($\gamma$) and shift ($\beta$) vectors can reshape the standardized features. When initialized to $\gamma = [1.0, 1.0, 1.0, 1.0]$ and $\beta = [0.0, 0.0, 0.0, 0.0]$, the values pass through unchanged. During training, the optimizer tunes $\gamma$ and $\beta$ so the layer can restore whatever dynamic range the downstream computation needs.

Because these operations run strictly across the feature coordinates of one token, other tokens in the sequence or batch don't influence this token's statistics.

Here's the calculation packaged as a standalone Python function with strict input validation:

from math import isclose, isfinite, sqrt

def layer_norm(values, gamma=None, beta=None, eps=1e-5):
    if not values or not isfinite(eps) or eps <= 0:
        raise ValueError("Use a nonempty vector and positive finite epsilon")
    if gamma is None:
        gamma = [1.0] * len(values)
    if beta is None:
        beta = [0.0] * len(values)
    if len(gamma) != len(values) or len(beta) != len(values):
        raise ValueError("gamma and beta must match the feature count")
    if not all(isfinite(v) for v in [*values, *gamma, *beta]):
        raise ValueError("All values must be finite")
    mean = sum(values) / len(values)
    variance = sum((value - mean) ** 2 for value in values) / len(values)
    scale = sqrt(variance + eps)
    return [
        gamma_i * ((value - mean) / scale) + beta_i
        for value, gamma_i, beta_i in zip(values, gamma, beta)
    ]

normalized = layer_norm([3.0, 1.0, -1.0, 5.0])
rounded = [round(value, 2) for value in normalized]
assert rounded == [0.45, -0.45, -1.34, 1.34]
assert isclose(sum(normalized), 0.0, abs_tol=1e-12)
assert layer_norm([3.0, 1.0], gamma=[0.0, 0.0]) == [0.0, 0.0]
assert layer_norm([7.0, 7.0], beta=[2.0, -1.0]) == [2.0, -1.0]

The general formula

For any token state $x \in \mathbb{R}^d$, LayerNorm computes:^{[2]Reference 2Layer Normalization.https://arxiv.org/abs/1607.06450}

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

where:

• $\mu = \frac{1}{d} \sum_{i=1}^d x_i$: feature mean of this token
• $\sigma^2 = \frac{1}{d} \sum_{i=1}^d (x_i - \mu)^2$: feature variance of this token
• $\gamma, \beta \in \mathbb{R}^d$: learnable per-channel scale and bias parameters
• $\epsilon$: small positive constant (typically $10^{-5}$ or $10^{-6}$) guarding against division by zero

Before the affine transform, variance equals $\sigma^2/(\sigma^2+\epsilon)$. It is near 1 when $\sigma^2\gg\epsilon$, exactly $1/2$ when $\sigma^2=\epsilon$, and zero for a flat vector. Learned per-channel $\gamma$ and $\beta$ can change both output mean and variance. Epsilon keeps the denominator nonzero, but the gain near a flat input can still approach $1/\sqrt{\epsilon}$.

In mixed precision, distinguish the input dtype from the dtype used for reductions and epsilon addition. FP16's smallest positive subnormal is about $5.96\times10^{-8}$; 1e-12 rounds to zero if represented directly in FP16. BF16 has a wider exponent range. Accumulating statistics in FP32 avoids that particular epsilon underflow, but cannot recover differences already lost when inputs were rounded. Match the checkpoint's epsilon and check the actual kernel; changing it changes the function.

Why BatchNorm is awkward in causal sequence models

BatchNorm usually reduces each channel across the batch and spatial or temporal positions. That creates different constraints from per-token LayerNorm:^{[3]Reference 3BatchNorm1d (PyTorch 2.14)https://docs.pytorch.org/docs/2.14/generated/torch.nn.BatchNorm1d.html}

1. Padding and causality: Padding can contaminate statistics unless explicitly excluded. BatchNorm1d on (batch, channels, time) reduces across batch and time, so a training-time token can depend on future tokens through those statistics. A causal attention mask does not repair that separate dependency.
2. Cross-sample coupling: With batch statistics, changing the other examples changes a token's normalization. Whether that helps or hurts depends on the task and sampling.
3. Train/eval statistics: With PyTorch's default running-statistics policy, evaluation uses stored estimates and can handle batch size one. If track_running_stats=False, evaluation also uses current batch statistics. One value per channel is insufficient for PyTorch's training-style BatchNorm; batch size one with several temporal positions is a different case.

Per-token LayerNorm avoids batch/sequence reduction. It does not erase padding's effects elsewhere in the model; attention masks and loss masks still matter.

Here's an executable demonstration comparing how LayerNorm and BatchNorm react when neighboring batch elements change:

from math import sqrt

def layer_norm(values, eps=1e-5):
    mean = sum(values) / len(values)
    variance = sum((value - mean) ** 2 for value in values) / len(values)
    return [(value - mean) / sqrt(variance + eps) for value in values]

def batch_norm_one_feature(values, eps=1e-5):
    mean = sum(values) / len(values)
    variance = sum((value - mean) ** 2 for value in values) / len(values)
    return [(value - mean) / sqrt(variance + eps) for value in values]

token = [1.0, 3.0, 5.0, 7.0]
ln_before = [layer_norm(row) for row in [token, [2.0, 4.0, 6.0, 8.0]]][0]
ln_after = [layer_norm(row) for row in [token, [100.0, 100.0, 100.0, 100.0]]][0]
batch_a = batch_norm_one_feature([token[0], 2.0, 3.0])
batch_b = batch_norm_one_feature([token[0], 2.0, 100.0])

print("LayerNorm token:", [round(value, 3) for value in ln_before])
print("Batch feature with peers=2,3:", round(batch_a[0], 3))
print("Batch feature with peers=2,100:", round(batch_b[0], 3))
assert ln_before == ln_after
assert round(batch_a[0], 3) != round(batch_b[0], 3)

LayerNorm token: [-1.342, -0.447, 0.447, 1.342]
Batch feature with peers=2,3: -1.225
Batch feature with peers=2,100: -0.718

The LayerNorm output for token is identical regardless of whether its neighbor has values near 2 or 100. The batch-normalized feature shifts from -1.225 to -0.718 purely because its batch neighbors shifted.

PyTorch implementation

For (batch, seq_len, d_model), nn.LayerNorm(d_model) reduces over hidden features. More generally, normalized_shape selects all matching trailing dimensions: nn.LayerNorm((seq_len, d_model)) also mixes time. That can leak future-token information in a causal decoder.^{[4]Reference 4LayerNormhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.LayerNorm.html}

Try changing just the last token. Which normalization keeps the first token unchanged?

import torch

x = torch.tensor([[[1., 2., 4.], [2., 5., 8.]]], dtype=torch.float64)
changed = x.clone()
changed[:, 1] += 100
per_token = torch.nn.LayerNorm(3, elementwise_affine=False).double()
over_time = torch.nn.LayerNorm((2, 3), elementwise_affine=False).double()
same_token = torch.allclose(per_token(x)[:, 0], per_token(changed)[:, 0])
same_time = torch.allclose(over_time(x)[:, 0], over_time(changed)[:, 0])
assert same_token and not same_time
print("First token unchanged with hidden-only reduction:", same_token)
print("First token unchanged with time-and-hidden reduction:", same_time)

First token unchanged with hidden-only reduction: True
First token unchanged with time-and-hidden reduction: False

Here is a single-hidden-dimension implementation. It upcasts half-precision statistics, keeps the output dtype, and broadcasts the affine parameters. It is an educational module, not a fused-kernel performance replacement:

from math import isfinite
import torch
import torch.nn as nn

class LayerNorm(nn.Module):
    """Layer Normalization from scratch."""
    def __init__(self, d_model: int, eps: float = 1e-5):
        super().__init__()
        if (not isinstance(d_model, int) or isinstance(d_model, bool)
                or d_model < 1 or not isfinite(eps) or eps <= 0):
            raise ValueError("Use a positive feature count and finite positive epsilon")
        self.gamma = nn.Parameter(torch.ones(d_model))
        self.beta = nn.Parameter(torch.zeros(d_model))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x shape: (batch, seq_len, d_model)
        if not x.is_floating_point() or x.shape[-1:] != self.gamma.shape:
            raise ValueError("Use a floating tensor whose final dimension matches d_model")
        work = x.float() if x.dtype in (torch.float16, torch.bfloat16) else x
        mean = work.mean(dim=-1, keepdim=True)
        var = work.var(dim=-1, keepdim=True, unbiased=False)
        x_norm = (work - mean) / torch.sqrt(var + self.eps)
        return (self.gamma * x_norm + self.beta).to(x.dtype)

# Quick check with the hand-worked example
x = torch.tensor([[[3.0, 1.0, -1.0, 5.0]]])  # shape (1, 1, 4)
ln = LayerNorm(d_model=4)
# Fix gamma=1, beta=0 to match the hand calculation
nn.init.constant_(ln.gamma, 1.0)
nn.init.constant_(ln.beta, 0.0)
out = ln(x)
print(out.round(decimals=2).detach())
expected = torch.tensor([[[0.45, -0.45, -1.34, 1.34]]])
assert torch.allclose(out.round(decimals=2), expected)

tensor([[[ 0.4500, -0.4500, -1.3400,  1.3400]]])

The gradients of scale and shift

Understanding LayerNorm during the backward pass requires looking at how gradients flow back through the reduction operations. Let $y = \text{LayerNorm}(x) = \gamma \odot \hat{x} + \beta$, where $\hat{x} = (x - \mu) / s$ and $s = \sqrt{\sigma^2 + \epsilon}$.

Let $g = \partial \mathcal{L} / \partial y$ represent upstream gradients. Because $\beta$ and $\gamma$ are shared across all tokens in the batch and sequence, their gradients sum over both axes:

$\frac{\partial \mathcal{L}}{\partial \beta_j} = \sum_{b, t} g_{btj}, \qquad \frac{\partial \mathcal{L}}{\partial \gamma_j} = \sum_{b, t} g_{btj} \hat{x}_{btj}$

For the input gradient $\partial \mathcal{L} / \partial x$ of one token, define $u = g \odot \gamma$. Differentiating through the mean and variance reductions yields:

$\frac{\partial \mathcal{L}}{\partial x} = \frac{1}{s} \left( u - \operatorname{mean}(u) - \hat{x} \odot \operatorname{mean}(u \odot \hat{x}) \right)$

Every mean in this expression runs across the feature coordinates of that token. The three terms have intuitive roles:

1. $u$: direct scaled gradient signal
2. $-\operatorname{mean}(u)$: enforces that the gradient of the centered vector has mean zero
3. $-\hat{x} \odot \operatorname{mean}(u \odot \hat{x})$: adjusts the component along the centered input. With positive epsilon this is not a complete orthogonal projection; changing the input scale still changes the normalized output slightly.

Here's an exact numerical test validating all three analytical gradients against PyTorch autograd:

import torch

torch.manual_seed(7)
x = torch.randn(2, 3, 4, dtype=torch.float64, requires_grad=True)
norm = torch.nn.LayerNorm(4, eps=1e-5).double()
with torch.no_grad():
    norm.weight.copy_(torch.tensor([0.5, 1.0, -2.0, 3.0]))
    norm.bias.copy_(torch.tensor([1.0, -1.0, 0.2, 0.0]))
g = torch.randn_like(x)
s = (x.var(-1, keepdim=True, unbiased=False) + norm.eps).sqrt()
xhat = (x - x.mean(-1, keepdim=True)) / s
manual = norm.weight * xhat + norm.bias
torch.testing.assert_close(norm(x), manual)
dx, dgamma, dbeta = torch.autograd.grad(norm(x), (x, norm.weight, norm.bias), g)
u = g * norm.weight
expected_dx = (u - u.mean(-1, keepdim=True)
               - xhat * (u * xhat).mean(-1, keepdim=True)) / s
torch.testing.assert_close(dx, expected_dx)
torch.testing.assert_close(dgamma, (g * xhat).sum((0, 1)))
torch.testing.assert_close(dbeta, g.sum((0, 1)))
print("Forward, input gradient, scale gradient, shift gradient: match")

Forward, input gradient, scale gradient, shift gradient: match

Where you put LayerNorm matters

Knowing how LayerNorm operates on a single vector sets up the bigger architectural choice: where does it go in a residual block?

One block, two placements

A transformer layer contains two residual sublayers: multi-head self-attention and a feed-forward network. Each sublayer takes an input $x_l$ and computes an update $F(x_l)$. The question is whether normalization wraps the sum or only the branch input:

• Post-LN: Normalizes after adding the residual update.
• Pre-LN: Normalizes before feeding the vector into the sublayer.

Post-LN: the original transformer

Vaswani et al. (2017) and early models like BERT placed LayerNorm on the output of the addition:^{[5]Reference 5Attention Is All You Need.https://arxiv.org/abs/1706.03762[6]Reference 6BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.https://arxiv.org/abs/1810.04805}

$x_{l+1} = \text{LayerNorm}(x_l + F(x_l))$

Compute the sublayer output, add it to the incoming skip connection, then normalize the sum. Normalization sits directly on the residual highway. Gradients backpropagating down the residual path must pass through a LayerNorm Jacobian at every single block.

Pre-LN: an identity contribution around the branch

GPT-2 reversed this order, normalizing the input before the sublayer:^{[7]Reference 7Language Models are Unsupervised Multitask Learners.https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf}

$x_{l+1} = x_l + F(\text{LayerNorm}(x_l))$

Normalize $x_l$, send the normalized vector into the sublayer, then add the sublayer's output to the raw, unnormalized $x_l$. The skip connection from $x_l$ to $x_{l+1}$ stays completely unobstructed.

GPT-2-style Pre-LN includes a final normalization before prediction. Its block input norms do not normalize the accumulated block output:^{[7]Reference 7Language Models are Unsupervised Multitask Learners.https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf[1]Reference 1On Layer Normalization in the Transformer Architecture.https://arxiv.org/abs/2002.04745}

$y = \text{Head}(\text{FinalNorm}(x_L))$

A two-layer walkthrough

Step through both layouts with a concrete numerical example. Let $F(x) = 0.2 \cdot x + 0.1$, and run two steps starting from $x_0 = [1.0, 2.0, 4.0, 8.0]$ with $\gamma = 1$ and $\beta = 0$.

With this example's fixed affine parameters, Post-LN outputs have zero mean and standard deviation close to one. Pre-LN accumulates these particular updates. A different $F$ could cancel or redirect them; neither layout forces a trained stream to follow this toy trajectory.

from math import sqrt

def layer_norm(values, eps=1e-5):
    mean = sum(values) / len(values)
    variance = sum((value - mean) ** 2 for value in values) / len(values)
    scale = sqrt(variance + eps)
    return [(value - mean) / scale for value in values]

def sublayer(values):
    # Deterministic stand-in for attention or an FFN update.
    return [0.2 * value + 0.1 for value in values]

def pre_ln_step(values):
    return [value + update for value, update in zip(values, sublayer(layer_norm(values)))]

def post_ln_step(values):
    raw = [value + update for value, update in zip(values, sublayer(values))]
    return layer_norm(raw)

start = [1.0, 2.0, 4.0, 8.0]
pre_after_two = pre_ln_step(pre_ln_step(start))
post_after_two = post_ln_step(post_ln_step(start))

assert sum(abs(value) for value in pre_after_two) > sum(abs(value) for value in start)
assert abs(sum(post_after_two) / len(post_after_two)) < 1e-12
assert abs(sum(pre_after_two) / len(pre_after_two) - 3.95) < 1e-12
print("Pre-LN after two steps:", [round(v, 3) for v in pre_after_two])
print("Post-LN after two steps:", [round(v, 3) for v in post_after_two])

Pre-LN after two steps: [0.79, 1.939, 4.237, 8.834]
Post-LN after two steps: [-1.026, -0.653, 0.093, 1.585]

Why Pre-LN can train more stably

The forward pass shows what gets normalized. The real breakthrough of Pre-LN lies in optimization dynamics: how gradients propagate from the loss down to layer 0.

The Post-LN gradient problem

Differentiating one Post-LN block with respect to its input yields:

$\frac{\partial x_{l+1}}{\partial x_l} = J_{\text{LN}, l} \cdot (I + J_{F, l})$

where $J_{\text{LN}, l}$ is evaluated at the sum $x_l + F(x_l)$.

Let $A_k=J_{\text{LN},k}(I+J_{F,k})$. The order matters because these matrices need not commute:

$\frac{\partial x_L}{\partial x_l}=A_{L-1}A_{L-2}\cdots A_l,\qquad g_l=A_l^TA_{l+1}^T\cdots A_{L-1}^Tg_L$

For one token, the exact normalization Jacobian is

$J_{\text{LN}}=\frac{\operatorname{diag}(\gamma)}{s}\left(I-\frac{\mathbf{1}\mathbf{1}^T}{d}-\frac{\hat{x}\hat{x}^T}{d}\right).$

It removes the common-offset direction and rescales other directions according to $s$, $\gamma$, and input geometry. A long product can attenuate or amplify a gradient; the placement equation alone does not give a universal depth law.

Xiong et al.'s formal last-FFN result uses a simplified initialization: single-head attention, zero Q/K projections (uniform attention), Gaussian remaining weights/inputs, and concentration assumptions. For width $d$ and depth $L$, their high-probability upper bounds on $\|\partial\tilde{\mathcal L}/\partial W_{2,L}\|_F$ are^{[1]Reference 1On Layer Normalization in the Transformer Architecture.https://arxiv.org/abs/2002.04745}

$\text{Post-LN}: O(d\sqrt{\ln d}),\qquad \text{Pre-LN}: O\!\left(d\sqrt{\frac{\ln d}{L}}\right).$

These concern the final FFN's second weight matrix, not every layer's weights or every trained checkpoint. The paper's broader analysis and experiments find large Post-LN gradients near the output and show that warmup helps its tested recipes. That is useful evidence for inspecting early update sizes, not a requirement that every Post-LN model use a fixed number of warmup steps.

The Pre-LN gradient advantage

Now differentiate a Pre-LN block:

$\frac{\partial x_{l+1}}{\partial x_l} = I + J_{F, l} \cdot J_{\text{LN}, l}$

Unrolling this across depth from layer $L$ to layer $l$ gives an additive decomposition:

$\frac{\partial x_L}{\partial x_l} = I + \sum_{k=l}^{L-1} J_{F, k} J_{\text{LN}, k} \frac{\partial x_k}{\partial x_l}$

The additive $I$ provides a path around the branch. But $I+B$ can have a very large norm, or $Bh=-h$ can cancel the identity on a direction $h$. The final norm and prediction head also participate in the loss gradient. Pre-LN does not guarantee nonvanishing gradients, bounded branch outputs, or safety at any learning rate.

Xiong et al. successfully removed warmup in their tested Pre-LN translation and BERT recipes.^{[1]Reference 1On Layer Normalization in the Transformer Architecture.https://arxiv.org/abs/2002.04745} Later Peri-LN experiments report Pre-LN gradient spikes and changing activation growth during training, and use warmup themselves.^{[8]Reference 8Peri-LN: Revisiting Normalization Layer in the Transformer Architecturehttps://arxiv.org/abs/2502.02732} Treat schedule, initialization, precision, and placement as one training recipe.

Here's an autograd check verifying the exact block Jacobians and their vector-Jacobian products:

import torch
from torch.autograd.functional import jacobian

torch.manual_seed(11)
x = torch.tensor([3.0, 1.0, -1.0, 5.0], dtype=torch.float64, requires_grad=True)
W = torch.randn(4, 4, dtype=x.dtype) * 0.2
identity = torch.eye(4, dtype=x.dtype)
def ln(v):
    return torch.nn.functional.layer_norm(v, (4,), eps=1e-5)
def pre(v):
    return v + W @ ln(v)
def post(v):
    return ln(v + W @ v)

J_pre = jacobian(pre, x)
J_post = jacobian(post, x)
torch.testing.assert_close(J_pre, identity + W @ jacobian(ln, x))
torch.testing.assert_close(J_post, jacobian(ln, x + W @ x) @ (identity + W))
g = torch.tensor([1.0, -2.0, 0.5, 3.0], dtype=x.dtype)
for block, J in [(pre, J_pre), (post, J_post)]:
    actual = torch.autograd.grad(block(x), x, grad_outputs=g)[0]
    torch.testing.assert_close(actual, J.T @ g)
print("Both block Jacobians and transposed gradient products: match")

# A fixed branch can cancel the identity on a particular direction.
point = torch.tensor([1., -1., 2., -2.], dtype=x.dtype)
direction = torch.tensor([2., -2., -1., 1.], dtype=x.dtype)
fixed_scale = (point.var(unbiased=False) + 1e-5).sqrt()
J_cancel = jacobian(lambda v: v - fixed_scale * ln(v), point)
cancelled = torch.allclose(J_cancel @ direction, torch.zeros_like(direction), atol=1e-12)
assert cancelled
print("Pre-LN identity term can cancel in one direction:", cancelled)

Both block Jacobians and transposed gradient products: match
Pre-LN identity term can cancel in one direction: True

A reported Pre-LN risk: representation collapse

Pre-LN can ease optimization while the residual stream dominates newer writes. This is a behavior to investigate, not a guaranteed outcome.

Because the residual stream bypasses normalization at each addition:

$x_L = x_0 + \sum_{l=0}^{L-1} F_l(\text{LayerNorm}(x_l))$

Suppose writes have bounded, comparable variance and are roughly uncorrelated with the accumulated stream. For fixed width, this motivates:

$\operatorname{Var}(x_l) \approx O(l) \implies \|x_l\| \approx O(\sqrt{l})$

Meanwhile, the update added at layer $l$ remains $\|\Delta x_l\| \approx O(1)$. The relative contribution of each new layer shrinks as depth increases:

$\frac{\|\Delta x_l\|}{\|x_l\|} \approx O\left(\frac{1}{\sqrt{l}}\right)$

A normalized input does not itself bound $F$'s output: learned weights can grow. Correlated writes can also invalidate the linear-variance argument. ResiDual studies increasingly similar normalized hidden states and diminishing later-block contributions in its analysis and translation experiments.^{[9]Reference 9ResiDual: Transformer with Dual Residual Connections.https://arxiv.org/abs/2304.14802} A small raw norm ratio alone does not prove that a layer is useless; measure normalized-state changes, direction, and the task effect of ablating it.

Residual magnitude growth across depth

This simulation deliberately adds the same aligned write eight times. Its Pre-LN norm grows linearly, so it illustrates write dilution, not the preceding uncorrelated-write $\sqrt{l}$ estimate. The DeepNorm-shaped column uses only residual scaling; it does not implement the full initialization or training recipe. The displayed write share is $\|\Delta x\|/(\|x_l\|+\|\Delta x\|)$, distinct from $\|\Delta x\|/\|x_l\|$:

from math import sqrt

def l2(values: list[float]) -> float:
    return sqrt(sum(v * v for v in values))

def layer_norm(values: list[float], eps: float = 1e-5) -> list[float]:
    mean = sum(values) / len(values)
    var = sum((v - mean) ** 2 for v in values) / len(values)
    scale = sqrt(var + eps)
    return [(v - mean) / scale for v in values]

depth = 8
update = [1.0, 0.0, 0.0, 0.0]
alpha = (2 * depth) ** 0.25  # encoder-only residual scale from DeepNet
x_pre = [1.0, 0.0, 0.0, 0.0]
x_post = [1.0, 0.0, 0.0, 0.0]
x_deep = [1.0, 0.0, 0.0, 0.0]

print(f"alpha={alpha:.3f}")
print("layer | Pre ||x|| | Pre write share | Post ||x|| | Deep ||αx|| | Deep ||x||")
for layer in range(1, depth + 1):
    pre_share = l2(update) / (l2(x_pre) + l2(update))
    x_pre = [a + b for a, b in zip(x_pre, update)]
    x_post = layer_norm([a + b for a, b in zip(x_post, update)])
    scaled = [alpha * a for a in x_deep]
    x_deep = layer_norm([a + b for a, b in zip(scaled, update)])
    print(
        f"{layer:5d} | {l2(x_pre):8.3f} | {pre_share:15.3f} | "
        f"{l2(x_post):10.3f} | {l2(scaled):9.3f} | {l2(x_deep):9.3f}"
    )

assert l2(x_pre) > l2(x_post)
assert abs(l2(x_post) - l2(x_deep)) < 1e-5
assert alpha == 2.0

alpha=2.000
layer | Pre ||x|| | Pre write share | Post ||x|| | Deep ||αx|| | Deep ||x||
    1 |    2.000 |           0.500 |      2.000 |     2.000 |     2.000
    2 |    3.000 |           0.333 |      2.000 |     4.000 |     2.000
    3 |    4.000 |           0.250 |      2.000 |     4.000 |     2.000
    4 |    5.000 |           0.200 |      2.000 |     4.000 |     2.000
    5 |    6.000 |           0.167 |      2.000 |     4.000 |     2.000
    6 |    7.000 |           0.143 |      2.000 |     4.000 |     2.000
    7 |    8.000 |           0.125 |      2.000 |     4.000 |     2.000
    8 |    9.000 |           0.111 |      2.000 |     4.000 |     2.000

In this aligned-write example, Pre-LN norm grows from 2.0 to 9.0 and write share drops from 0.50 to 0.11. With unit affine parameters and tiny epsilon, the normalized columns stay near $\sqrt{d}=2$. Learned affine weights would change that scale.

DeepNorm: scaled residual and matched init

Wang et al. (2022) introduced DeepNorm in DeepNet as a way to stabilize very deep Post-LN-shaped stacks.^{[10]Reference 10DeepNet: Scaling Transformers to 1,000 Layers.https://arxiv.org/abs/2203.00555}

DeepNorm preserves the Post-LN layout but modifies the residual sum with a depth-dependent scaling factor $\alpha$ and a matched initialization gain $\beta$:

$x_{l+1} = \text{LayerNorm}(\alpha \cdot x_l + G_l(x_l))$

The original recipe applies Xavier gain $\beta$ to FFN weights and attention's value/output projections, while Q/K projections use gain 1. Multiplying a weight by $\beta$ multiplies its initialization variance by $\beta^2$. This is not an extra $\beta G_l$ multiplication on every forward pass.

For an encoder-only stack with $N$ layers and two residual sublayers per layer:

$\alpha = (2N)^{1/4}, \qquad \beta = (8N)^{-1/4}$

A decoder-only stack uses the same formula with its layer count. An encoder-decoder is different because the decoder also has cross-attention:

The analyzed model update is a change in the same model's output after a parameter update, $\|F(x;\theta+\Delta\theta)-F(x;\theta)\|$. It is not the difference between the final and initial hidden states. Under its assumptions, DeepNet targets an output change on the order of the SGD step size, independent of increasing depth. The paper reports scaling to 1,000 layers; its documented translation recipes include warmup (for example, 4,000 updates on WMT-17), so that result does not establish a warmup-free recipe.^{[10]Reference 10DeepNet: Scaling Transformers to 1,000 Layers.https://arxiv.org/abs/2203.00555}

RMSNorm: scale-only normalization

Where normalization sits is one choice. What formula sits inside that box is another.

Zhang and Sennrich (2019) tested whether mean-centering was necessary in their networks.^{[11]Reference 11Root Mean Square Layer Normalization.https://arxiv.org/abs/1910.07467} RMSNorm drops it and uses the uncentered root mean square, usually with a learned gain and no learned shift. Their comparable-quality results are empirical evidence for the tested tasks, not proof that centering never matters:

$\text{RMSNorm}(x) = \gamma \odot \frac{x}{\operatorname{RMS}(x)} = \gamma \odot \frac{x}{\sqrt{\frac{1}{d} \sum_{i=1}^d x_i^2 + \epsilon}}$

Less arithmetic does not specify a GPU speedup

Llama, Qwen 2.5, and Gemma 2 provide examples of RMSNorm adoption; their placement recipes differ.^{[12]Reference 12Qwen2.5 Technical Reporthttps://arxiv.org/abs/2412.15115[13]Reference 13Gemma 2: Improving Open Language Models at a Practical Sizehttps://storage.googleapis.com/deepmind-media/gemma/gemma-2-report.pdf} The formula needs a mean-square statistic, without a mean subtraction or canonical bias. That saves work, but statistics are not the same as HBM passes or kernel launches.

LayerNorm can combine mean and variance estimation using Welford's algorithm, and fuse statistics, affine transformation, and output into a kernel. Triton's official tutorial demonstrates fused LayerNorm. It does not follow that RMSNorm always halves memory traffic, registers, barriers, or runtime.^{[14]Reference 14PyTorch CUDA LayerNorm kernelshttps://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/cuda/layer_norm_kernel.cu[15]Reference 15Layer Normalization: fused forward and backward tutorialhttps://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html}

The original RMSNorm paper reports 7%–64% running-time reductions across its different tested models and frameworks, not a fixed modern GPU kernel result.^{[11]Reference 11Root Mean Square Layer Normalization.https://arxiv.org/abs/1910.07467} Actual speed depends on width, row count, dtype, fusion, hardware, and backward work. Measure both the operator and the full step. If normalization takes 2% of a step, halving its time saves only 1% of that step.

Here's an implementation of RMSNorm:

from math import isfinite, sqrt

def rms_norm(values, weight=None, eps=1e-6):
    if not values or not isfinite(eps) or eps <= 0:
        raise ValueError("Use a nonempty vector and positive finite epsilon")
    if weight is None:
        weight = [1.0] * len(values)
    if len(weight) != len(values) or not all(isfinite(v) for v in [*values, *weight]):
        raise ValueError("Use one finite weight per finite feature")
    rms = sqrt(sum(value * value for value in values) / len(values) + eps)
    return [value * scale / rms for value, scale in zip(values, weight)]

values = [3.0, 1.0, -1.0, 5.0]
normalized = rms_norm(values)

assert round(sqrt(sum(value * value for value in normalized) / len(normalized)), 6) == 1.0
assert round(sum(normalized) / len(normalized), 6) != 0.0
assert rms_norm([3.0, 1.0], weight=[0.0, 0.0]) == [0.0, 0.0]
assert rms_norm([0.0, 0.0]) == [0.0, 0.0]

And the equivalent PyTorch module:^{[16]Reference 16RMSNormhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.RMSNorm.html}

from math import isfinite
import torch
import torch.nn as nn

class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization."""
    def __init__(self, d_model: int, eps: float = 1e-6):
        super().__init__()
        if (not isinstance(d_model, int) or isinstance(d_model, bool)
                or d_model < 1 or not isfinite(eps) or eps <= 0):
            raise ValueError("Use a positive feature count and finite positive epsilon")
        self.weight = nn.Parameter(torch.ones(d_model))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x shape: (batch, seq_len, d_model)
        if not x.is_floating_point() or x.shape[-1:] != self.weight.shape:
            raise ValueError("Use a floating tensor whose final dimension matches d_model")
        work = x.float() if x.dtype in (torch.float16, torch.bfloat16) else x
        rms = torch.sqrt(work.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return (self.weight * (work / rms)).to(x.dtype)

# Quick sanity check
x = torch.tensor([[[3.0, 1.0, -1.0, 5.0]]])
rms = RMSNorm(d_model=4)
out = rms(x)
print(out.round(decimals=4).detach())
feature_rms = torch.sqrt(out.pow(2).mean(dim=-1))
assert torch.allclose(feature_rms, torch.ones_like(feature_rms), atol=1e-5)
assert not torch.allclose(out.mean(dim=-1), torch.zeros_like(out.mean(dim=-1)))

tensor([[[ 1.0000,  0.3333, -0.3333,  1.6667]]])

To compare with PyTorch, specify nn.RMSNorm(4, eps=1e-6). This module's explicit epsilon differs from PyTorch 2.14's eps=None policy, which uses the computation type's machine epsilon (FP32 for half/BF16/FP32 inputs, FP64 for FP64 inputs).^{[16]Reference 16RMSNormhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.RMSNorm.html} Positive epsilon makes rescaling invariance approximate, particularly for tiny inputs. A zero vector remains zero; a constant nonzero vector generally remains nonzero. Both normalize before the learned gain, which can change the final RMS.

The cast before squaring matters. FP16 can represent 300, but its square, 90,000, exceeds FP16's largest finite value. torch.finfo(torch.float16).max reports that limit.^{[17]Reference 17Type Infohttps://docs.pytorch.org/docs/stable/type_info.html} Predict what the two RMS calculations will return, then compare them:

import torch
import torch.nn.functional as F

large = torch.tensor([300., -300.], dtype=torch.float16)
half_rms = (large.square().mean() + 1e-6).sqrt()
wide_rms = (large.float().square().mean() + 1e-6).sqrt()
assert torch.isinf(half_rms) and wide_rms.item() == 300.
print("RMS, FP16 accumulation:", half_rms.item())
print("RMS, FP32 accumulation:", wide_rms.item())

# A separate failure occurs before the normalization kernel runs.
original = torch.tensor([10000., 10001., 9999., 10002.])
rounded = original.half().float()
before = F.layer_norm(original, (4,), eps=1e-5)
after = F.layer_norm(rounded, (4,), eps=1e-5)
assert torch.equal(rounded, torch.full_like(rounded, 10000.))
assert torch.equal(after, torch.zeros_like(after))
assert before.abs().max() > 1.
print("FP16 input, cast back to FP32:", rounded.tolist())
print("LayerNorm of original:", [round(v, 4) for v in before.tolist()])
print("LayerNorm after input rounding:", after.tolist())

RMS, FP16 accumulation: inf
RMS, FP32 accumulation: 300.0
FP16 input, cast back to FP32: [10000.0, 10000.0, 10000.0, 10000.0]
LayerNorm of original: [-0.4472, 0.4472, -1.3416, 1.3416]
LayerNorm after input rounding: [0.0, 0.0, 0.0, 0.0]

Casting the squared tensor afterward leaves the infinity intact. Casting before the product repairs this accumulation failure. In the second case, FP16 rounds all four values to 10,000 before LayerNorm sees them. FP32 accumulation cannot recover those lost differences. When checking a low-precision kernel against a reference, feed both the same rounded input; compare with the original FP32 input separately to measure input-rounding error.

Reported normalization layouts

Different architectures combine normalization operators and placements in distinct configurations:

These rows name different design dimensions, not mutually exclusive architectures. Qwen3 combines Pre-RMSNorm with per-head Q/K RMSNorm before RoPE. Gemma 3 keeps its input/output residual-branch norms and replaces Gemma 2's attention soft-capping with QK-Norm. These concrete designs were checked October 10, 2026; the Qwen3 code reference is pinned to Transformers 5.17.0.^{[20]Reference 20Transformers 5.17.0 Qwen3 normalization and residual placementhttps://github.com/huggingface/transformers/blob/v5.17.0/src/transformers/models/qwen3/modeling_qwen3.py[19]Reference 19Gemma 3 Technical Reporthttps://arxiv.org/abs/2503.19786}

A naming trap: Qwen3's post_attention_layernorm is applied after the attention residual addition but before the MLP, followed by another residual addition. It is the MLP's input norm in a Pre-LN layout. Read the operations around the add instead of classifying placement from a variable name.^{[20]Reference 20Transformers 5.17.0 Qwen3 normalization and residual placementhttps://github.com/huggingface/transformers/blob/v5.17.0/src/transformers/models/qwen3/modeling_qwen3.py}

QK-Norm controls one source of attention scale

Residual normalization regulates activations between layers. But inside the attention head, another instability can trigger loss divergence: attention logit growth.

In standard attention, raw attention logits equal:

$S_{ij} = \frac{q_i k_j^T}{\sqrt{d_k}}$

Growing Q/K norms can enlarge score differences, making attention nearly one-hot. Absolute magnitude alone is insufficient: softmax([1000, 1000]) is [0.5, 0.5], and adding one constant to every score changes nothing. ViT-22B and Rybakov et al. report specific instability experiments; Rybakov also examines growth in other linear outputs and combines QK normalization with soft-capping.^{[18]Reference 18Scaling Vision Transformers to 22 Billion Parametershttps://arxiv.org/abs/2302.05442[21]Reference 21Methods of improving LLM training stabilityhttps://arxiv.org/abs/2410.16682}

For an increasingly one-hot distribution, the score-to-probability Jacobian entries approach zero:

$\frac{\partial \operatorname{Softmax}(S)_i}{\partial S_j} = \operatorname{Softmax}(S)_i (\delta_{ij} - \operatorname{Softmax}(S)_j) \approx 0$

This can weaken that route, not freeze every attention parameter. The value path still participates, and losses composed directly with softmax can have different gradient behavior (for example, cross-entropy has a logit gradient $p-y$).

One QK-Norm family applies LayerNorm or RMSNorm before the dot product:^{[21]Reference 21Methods of improving LLM training stabilityhttps://arxiv.org/abs/2410.16682}

$q' = \text{Norm}(q), \quad k' = \text{Norm}(k), \quad S_{ij} = \frac{q'_i {k'_j}^T}{\sqrt{d_k}}$

The original QKNorm paper instead uses per-head unit-L2 vectors and a learned scalar multiplying cosine similarity. That is different from unit-RMS vectors with the usual $1/\sqrt{d_k}$ scaling. Normalization axis, affine gain, temperature, and RoPE order are checkpoint choices.^{[22]Reference 22Query-Key Normalization for Transformershttps://arxiv.org/abs/2010.04245[20]Reference 20Transformers 5.17.0 Qwen3 normalization and residual placementhttps://github.com/huggingface/transformers/blob/v5.17.0/src/transformers/models/qwen3/modeling_qwen3.py}

Cauchy-Schwarz always gives:

$|q'_i \cdot k'_j| \le \|q'_i\|_2 \|k'_j\|_2$

Without affine parameters, unit-RMS vectors have norm $\sqrt{d_k}$, giving $|S_{ij}|\le\sqrt{d_k}$ under standard scaling. Learned gains enlarge this bound; biases and extra temperatures also matter. Even that unit-gain bound permits highly peaked attention. Normalization reduces dependence on raw Q/K magnitude, without guaranteeing entropy or stable optimization.

import torch

d = 256
q = torch.cat((torch.ones(d // 2), -torch.ones(d // 2))).double()
keys = torch.stack((q, -q))
# Both vectors already have zero mean, RMS 1, and L2 norm sqrt(d).
scores = keys @ q / d**0.5
probabilities = scores.softmax(0)
large_equal = torch.tensor([1000., 1000.]).softmax(0)
gain = 10
gained_scores = (gain * keys) @ (gain * q) / d**0.5
assert scores.tolist() == [16., -16.]
assert probabilities[1] < 2e-14
assert torch.equal(large_equal, torch.tensor([0.5, 0.5]))
assert gained_scores.tolist() == [1600., -1600.]
print("Uniform large logits:", large_equal.tolist())
print("Normalized logits:", scores.tolist())
print(f"Opposite-key probability: {probabilities[1].item():.2e}")
print("Gain 10 on both Q and K:", gained_scores.tolist())

Uniform large logits: [0.5, 0.5]
Normalized logits: [16.0, -16.0]
Opposite-key probability: 1.27e-14
Gain 10 on both Q and K: [1600.0, -1600.0]

from math import isclose, sqrt

def rms_rescale(values):
    rms = sqrt(sum(value * value for value in values) / len(values))
    if rms == 0:
        raise ValueError("This epsilon-free illustration requires a nonzero vector")
    return [value / rms for value in values]

def attention_logit(query, key):
    return sum(q * k for q, k in zip(query, key)) / sqrt(len(query))

query = [6.0, -3.0, 2.0, 1.0]
key = [4.0, -2.0, 1.0, 3.0]
large_query = [20.0 * value for value in query]

raw = attention_logit(query, key)
raw_large = attention_logit(large_query, key)
controlled = attention_logit(rms_rescale(query), rms_rescale(key))
controlled_large = attention_logit(rms_rescale(large_query), rms_rescale(key))

print("raw logits:", round(raw, 3), round(raw_large, 3))
print("RMS-rescaled logits:", round(controlled, 3), round(controlled_large, 3))
assert raw_large == 20.0 * raw
assert isclose(controlled, controlled_large, rel_tol=1e-12)

raw logits: 17.5 350.0
RMS-rescaled logits: 1.807 1.807

Peri-LN and sandwich normalization

Sandwich-style layouts surround the sublayer with normalizers on both sides:

$x_{l+1} = x_l + \text{Norm}_{\text{out}}(F(\text{Norm}_{\text{in}}(x_l)))$

Here $\text{Norm}_{\text{out}}$ sits inside the branch before the residual addition. It regulates each write while retaining an additive identity term. The stream can still grow through repeated writes, and learned gains affect its scale. Gemma 2 uses branch input/output RMSNorm; Kim et al.'s full Peri-LN recipe also normalizes embeddings and the final output. “Post-norm” on a branch is not the plain Post-LN equation that normalizes the entire sum.^{[8]Reference 8Peri-LN: Revisiting Normalization Layer in the Transformer Architecturehttps://arxiv.org/abs/2502.02732[13]Reference 13Gemma 2: Improving Open Language Models at a Practical Sizehttps://storage.googleapis.com/deepmind-media/gemma/gemma-2-report.pdf}

In Kim et al.'s experiments, 400M, 1.5B, and 3.2B models trained on 30B tokens show fewer gradient spikes and better measured losses with Peri-LN. The paper uses RMSNorm for its main runs and 10% warmup. Its observation of changing Pre-LN variance growth during training also cautions against treating an initialization estimate as a universal trained-model law.^{[8]Reference 8Peri-LN: Revisiting Normalization Layer in the Transformer Architecturehttps://arxiv.org/abs/2502.02732}

Debug normalization failures

When a transformer diverges during training, systematic logging isolates the failure mode:

Practice: compare residual MLP blocks

Test the difference in gradient propagation directly by building a configurable residual block that toggles between Pre-LN and Post-LN:

import torch
import torch.nn as nn

class ToggleBlock(nn.Module):
    """Residual MLP block with style='pre' or style='post'."""
    def __init__(self, d_model: int, style: str = "pre"):
        super().__init__()
        if style not in ("pre", "post"):
            raise ValueError("style must be 'pre' or 'post'")
        if not isinstance(d_model, int) or isinstance(d_model, bool) or d_model < 1:
            raise ValueError("d_model must be a positive integer")
        self.style = style
        self.norm = nn.LayerNorm(d_model)
        # Simplified sublayer: single linear + ReLU
        self.sublayer = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.ReLU(),
            nn.Linear(d_model * 4, d_model),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.style == "pre":
            # x_l + Sublayer(LayerNorm(x_l))
            return x + self.sublayer(self.norm(x))
        else:
            # LayerNorm(x_l + Sublayer(x_l))
            return self.norm(x + self.sublayer(x))

class TinyStack(nn.Module):
    def __init__(self, d_model: int = 64, depth: int = 6, style: str = "pre"):
        super().__init__()
        if not isinstance(depth, int) or isinstance(depth, bool) or depth < 1:
            raise ValueError("depth must be a positive integer")
        self.blocks = nn.ModuleList([ToggleBlock(d_model, style) for _ in range(depth)])
        self.final_norm = nn.LayerNorm(d_model) if style == "pre" else nn.Identity()
        self.head = nn.Linear(d_model, 16)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for block in self.blocks:
            x = block(x)
        return self.head(self.final_norm(x))

def inspect_grad(style: str, d_model: int = 64, depth: int = 6, seed: int = 0):
    torch.manual_seed(seed)
    model = TinyStack(d_model=d_model, depth=depth, style=style)
    x = torch.randn(2, 8, d_model)
    target = torch.randn(2, 8, 16)
    out = model(x)
    loss = nn.functional.mse_loss(out, target)
    loss.backward()
    grad_norms = [block.sublayer[0].weight.grad.norm().item() for block in model.blocks]
    rounded_norms = [round(value, 4) for value in grad_norms]
    print(f"{style:4} | loss={loss.item():.4f} | block grad norms={rounded_norms}")
    return grad_norms

pre_grad_norms = inspect_grad("pre")
post_grad_norms = inspect_grad("post")

assert len(pre_grad_norms) == len(post_grad_norms) == 6
assert all(value > 0 for value in pre_grad_norms + post_grad_norms)

pre  | loss=1.3247 | block grad norms=[0.2664, 0.2565, 0.2494, 0.2392, 0.2472, 0.2415]
post | loss=1.3183 | block grad norms=[0.2663, 0.2614, 0.2619, 0.2555, 0.2684, 0.2698]

Both layouts have healthy nonzero gradients in this small MLP example. It contains no attention or training loop, and is not a reproduction of the papers' transformers. The seed gives matching linear weights, inputs, targets, and heads here; the final Pre-LN norm initializes to unit scale/zero bias without drawing random values. For a training comparison, preserve those controls and compare schedules and output norms explicitly.

Mini exercise

1. Try depths 6, 24, and 48 across several seeds. Record first/last-block gradient ratios. Does a claimed placement effect survive seed changes, or do both layouts remain similar?
2. Replace the final Pre-LN norm with nn.Identity() while preserving the other weights. Measure head input RMS and logits. This is an architectural ablation, not a safe checkpoint-porting change.
3. Compare nn.LayerNorm(d_model, eps=1e-5) with nn.RMSNorm(d_model, eps=1e-5). Match epsilon and learned gains; then ask whether any difference comes from centering, placement, or the training recipe. CPU gradient norms do not measure GPU kernel speed.

Mastery check

• Calculate LayerNorm by hand: token mean, population variance, epsilon adjustment, and learned scale/shift.
• Trace the exact reduction axes. Explain BatchNorm's batch/time coupling and eval policy, and show why a hidden-only norm avoids future-token mixing.
• Formulate the block Jacobians for Pre-LN ($I + J_F J_{\text{LN}}$) and Post-LN ($J_{\text{LN}}(I + J_F)$).
• Scope Xiong's last-FFN initialization bound and explain why measured gradients/updates motivate a warmup experiment without establishing a universal schedule.
• State the assumptions behind relative-write dilution, distinguish aligned from uncorrelated writes, and use task ablations to investigate redundancy.
• Implement RMSNorm with explicit epsilon and gain. Separate formula simplification from measured operator or full-step performance.
• Preserve DeepNorm's $\operatorname{LN}(\alpha x+G(x))$ placement and selective initialization gain. Distinguish decoder-only from encoder-decoder scales, and demonstrate why normalized Q/K can still produce peaked attention.

Practice checkpoints