class MiniBlock(nnx.Module):
"""Attention + MLP on a residual stream; norm placement as a flag."""
def __init__(self, num_hiddens, pre_norm, rngs=None):
rngs = nnx.Rngs(0) if rngs is None else rngs
self.pre_norm = pre_norm
self.attention = d2l.MultiHeadAttention(num_hiddens, 4, dropout=0,
rngs=rngs)
self.W_1 = nnx.Linear(num_hiddens, 4 * num_hiddens, rngs=rngs)
self.W_2 = nnx.Linear(4 * num_hiddens, num_hiddens, rngs=rngs)
self.norm1 = nnx.LayerNorm(num_hiddens, rngs=rngs)
self.norm2 = nnx.LayerNorm(num_hiddens, rngs=rngs)
def ffn(self, X):
return self.W_2(nnx.gelu(self.W_1(X)))
def __call__(self, X):
if self.pre_norm:
Y = self.norm1(X)
X = X + self.attention(Y, Y, Y, None)[0]
return X + self.ffn(self.norm2(X))
X = self.norm1(X + self.attention(X, X, X, None)[0])
return self.norm2(X + self.ffn(X))The Transformer Block
Dive into Deep Learning · §11.1
The transformer block
a residual stream and two sublayers · where the norm goes · GELU to SwiGLU · one configurable class
Anatomy: two sublayers, one highway
- Attention: the only place positions interact — each token queries the others and adds the retrieved mixture to its stream.
- FFN: per-position computation — and about two thirds of the parameters.
- Both read from the residual stream and add back; shape (n, d) in, (n, d) out — that is what makes stacking possible.
Two arrangements, two lines of code
Post-LN (2017): norm on the stream, after each addition \mathbf{h} = \mathrm{Norm}_1(\mathbf{x} + \mathrm{Attn}(\mathbf{x})), \qquad \mathbf{y} = \mathrm{Norm}_2(\mathbf{h} + \mathrm{FFN}(\mathbf{h}))
Pre-LN (modern): norm on the branch, stream untouched \mathbf{h} = \mathbf{x} + \mathrm{Attn}(\mathrm{Norm}_1(\mathbf{x})), \qquad \mathbf{y} = \mathbf{h} + \mathrm{FFN}(\mathrm{Norm}_2(\mathbf{h}))
Measure it before training it
Stack 32 freshly initialized blocks (Xavier, as in 2017), push a random sequence through — deterministic, seconds:
post-LN: stream RMS at k=32: 1.00, token spread at k=8,16,32: 3.5e-01 2.5e-02 2.0e-04
grad norms at k=32: W_q 5.6e-09 W_k 4.9e-09 W_v 6.5e-02 W_o 6.6e-02
pre-LN: stream RMS at k=32: 5.22, token spread at k=8,16,32: 7.2e-01 5.3e-01 3.7e-01
grad norms at k=32: W_q 2.3e-03 W_k 2.3e-03 W_v 8.2e-02 W_o 8.4e-02Reading the experiment
- Stream RMS: post-LN pinned at 1; pre-LN grows like \sqrt{\textrm{depth}} — each successive block’s relative contribution shrinks.
- Token spread: uniform attention pulls tokens toward their mean; post-LN locks it in → geometric collapse (to 10^{-4} by block 32). Pre-LN dilutes it → polynomial.
- Gradients: identical tokens make the mixture independent of the query/key weights — post-LN’s upper \mathbf{W}_q, \mathbf{W}_k get ~10^{6}\times less gradient; value, output, and FFN weights stay healthy in both.
Rank collapse (Dong et al., 2021) caught in the act; why post-LN needed warmup (Xiong et al., 2020) and GPT-2 went pre-norm.
RMSNorm: drop what you don’t need
\mathrm{RMSNorm}(\mathbf{x}) = \frac{\mathbf{x}}{\sqrt{\tfrac{1}{d}\sum_i x_i^2 + \epsilon}} \odot \boldsymbol{\gamma}
ours vs. built-in: max deviation 9.5e-07
LayerNorm: 1.076 ms
RMSNorm: 0.732 msNo centering, no bias, one statistic — comparable quality, less machinery (Zhang & Sennrich, 2019; Llama onward).
Four places for the norm
Post-LN (2017), pre-LN (the default), post-sublayer off-stream (OLMo 2), pre-and-post (Gemma 3) — only post-LN interrupts the stream’s identity path:
QK-norm: the same discipline, inside attention
- Nothing above controls the attention logits \mathbf{q}^\top\mathbf{k}/\sqrt{d}: if training inflates q and k, the softmax saturates and its gradient vanishes.
- QK-norm: RMSNorm on queries and keys, per head, right before the dot product — logit scale pinned by construction.
- Met at ViT-22B scale; standard in 2025 models (Gemma 3, Qwen3, OLMo 2).
OLMo 2 also re-litigated placement: norms after each sublayer — but still off the stream’s identity path.
The FFN: from fixed nonlinearity to learned gate
Classic: \mathrm{FFN}(\mathbf{x}) = \mathbf{W}_2\,\phi(\mathbf{W}_1\mathbf{x}), width 4d — ReLU, then GELU.
SwiGLU (Shazeer, 2020): two linear views, one squashed into a soft gate, multiplied \mathrm{SwiGLU}(\mathbf{x}) = \mathbf{W}_2\big(\mathrm{SiLU}(\mathbf{W}_g\mathbf{x}) \odot \mathbf{W}_1\mathbf{x}\big)
Three matrices → shrink the width to \tfrac{8}{3}d and the budget matches 8d^2.
class FeedForward(nnx.Module):
"""Position-wise FFN: GELU MLP or SwiGLU at matched parameter count."""
def __init__(self, num_hiddens, act='swiglu', bias=False, rngs=None):
rngs = nnx.Rngs(0) if rngs is None else rngs
assert act in ('swiglu', 'gelu'), f'unknown act: {act!r}'
self.act = act
if act == 'gelu':
width = 4 * num_hiddens
else: # 'swiglu': three matrices; width 8d/3 matches the MLP budget
width = round(8 * num_hiddens / 3)
self.W_g = nnx.Linear(num_hiddens, width, use_bias=bias,
rngs=rngs)
self.W_1 = nnx.Linear(num_hiddens, width, use_bias=bias, rngs=rngs)
self.W_2 = nnx.Linear(width, num_hiddens, use_bias=bias, rngs=rngs)
def __call__(self, X):
if self.act == 'gelu':
return self.W_2(nnx.gelu(self.W_1(X)))
return self.W_2(nnx.silu(self.W_g(X)) * self.W_1(X))One block, every configuration
Norm, activation, arrangement as flags; attention and FFN as factories (MoE and grouped-query attention drop in later):
class TransformerBlock(nnx.Module):
"""Configurable transformer block: attention + FFN on a residual
stream."""
def __init__(self, num_hiddens, num_heads, dropout=0, norm='rms',
act='swiglu', pre_norm=True, bias=False, attn_factory=None,
ffn_factory=None, rngs=None):
rngs = nnx.Rngs(params=0, dropout=1) if rngs is None else rngs
assert norm in ('rms', 'layer'), f'unknown norm: {norm!r}'
assert num_hiddens % num_heads == 0
self.pre_norm = pre_norm
make_norm = nnx.RMSNorm if norm == 'rms' else nnx.LayerNorm
self.norm1 = make_norm(num_hiddens, rngs=rngs)
self.norm2 = make_norm(num_hiddens, rngs=rngs)
self.attention = (d2l.MultiHeadAttention(num_hiddens, num_heads,
dropout, bias=bias,
rngs=rngs)
if attn_factory is None else attn_factory(rngs))
self.ffn = (FeedForward(num_hiddens, act, bias=bias, rngs=rngs)
if ffn_factory is None else ffn_factory(rngs))
self.dropout = nnx.Dropout(dropout, rngs=rngs)
def __call__(self, X, valid_lens=None):
if self.pre_norm:
Y = self.norm1(X)
X = X + self.dropout(self.attention(Y, Y, Y, valid_lens)[0])
return X + self.dropout(self.ffn(self.norm2(X)))
X = self.norm1(X + self.dropout(
self.attention(X, X, X, valid_lens)[0]))
return self.norm2(X + self.dropout(self.ffn(X)))The flags at work: GELU vs. SwiGLU
Same data, seed, optimizer, 600 steps; parameters matched to 0.1%:
gelu: loss at step 200/400/600: 2.27/1.73/1.41
swiglu: loss at step 200/400/600: 2.22/1.57/1.24More than a tenth of a nat, stable across seeds — the kind of small, real, equal-cost win architecture progress is made of.
Recap
- Attention communicates, the FFN computes; both add into the residual stream.
- Norm placement decides depth behavior at init: post-LN collapses tokens geometrically and starves upper attention layers of gradient; pre-LN grows the stream like \sqrt{\textrm{depth}} and keeps every block trainable.
- RMSNorm = the scale statistic only; QK-norm pins the attention logits.
- SwiGLU beats the matched-budget MLP by a small, seed-stable margin.
TransformerBlock: a decade of designs as one constructor signature — the chapter never builds another block.