Pre-LN vs Post-LN Residual Pattern

Where does normalization happen?

A residual block adds a sublayer's output back to its input. Moving LayerNorm across that addition changes both the forward computation and the path gradients follow.

This exercise compares two blocks with the same feed-forward network:

FFN(z) = Dense(D)(relu(Dense(d_ff)(z)))

Pre-LN:   y = x + FFN(LayerNorm(x))
Post-LN:  y = LayerNorm(x + FFN(x))

In pre-LN, the network receives normalized features and adds its result to the original x. In post-LN, the network receives x directly, and LayerNorm acts on the residual sum. Neither formula is the same as x + LayerNorm(FFN(x)).

Normalize features within each token

The input has shape (T, D). Linen's default LayerNorm() computes a separate mean and variance over the final D features of each row. It does not mix statistics across tokens. With learned scale gamma and offset beta, each row is transformed as:

LN(z) = gamma * (z - mean(z)) / sqrt(variance(z) + 1e-6) + beta

Use Linen's defaults: epsilon 1e-6, scale initialized to one, and offset initialized to zero. The normalized values are not guaranteed to have exactly unit variance because of epsilon. Learned scale and offset can further change their mean and variance.

Both Dense layers act on the final axis and share their parameters across rows. The hidden shape is (T, d_ff); the second projection restores (T, D) so the residual addition is valid.

A small example of the ordering

To isolate the difference, use f(z) = 2*z in this example only. The exercise itself uses the two-layer FFN above.

import jax.numpy as jnp

def ln(z):
    mean = jnp.mean(z, axis=-1, keepdims=True)
    variance = jnp.mean((z - mean) ** 2, axis=-1, keepdims=True)
    return (z - mean) / jnp.sqrt(variance + 1e-6)

x = jnp.array([[1., 3.]], dtype=jnp.float32)
pre = x + 2 * ln(x)
post = ln(x + 2 * x)

assert jnp.allclose(pre, jnp.array([[-1., 5.]]), atol=2e-6)
assert jnp.allclose(post, jnp.array([[-1., 1.]]), atol=2e-6)

Pre-LN leaves the residual sum unnormalized inside this block. A larger model may still normalize it later, for example after a stack of blocks.

Why the choice matters for training

Pre-LN has a direct identity contribution to the input gradient through the residual addition. Post-LN sends that path through the final LayerNorm too. This changes how gradients behave, but an identity contribution does not guarantee constant gradient norms or stable training at every depth.

Xiong et al. (2020) demonstrated pre-LN training without warmup in particular settings. That does not make warmup unnecessary for every model, optimizer or learning rate. This exercise compares forward passes; it does not establish which training setup will work best.

Both branches have the same parameter shapes and count. Swapping their placement still changes the function, so matching shapes alone does not make it safe to switch a trained model from one pattern to the other.