%matplotlib inline
from d2l import torch as d2l
import time
import torch
from torch import nn
from torch.nn import functional as F11.1 The Transformer Block
The previous chapter developed attention as a layer in which each position can read from the others. Deployed models place this layer in a transformer block. Attention communicates between positions, a position-wise feed-forward network transforms each position independently, and both add their outputs to a shared residual stream. This section studies two design choices that distinguish the original block from current ones: the placement of normalization and the form of the feed-forward network. Experiments examine signal propagation at initialization and compare a standard MLP with its gated successor at matched parameter count. We then combine these choices in a configurable TransformerBlock class used throughout the remainder of the chapter.
%matplotlib inline
from d2l import jax as d2l
from flax import nnx
import jax
from jax import numpy as jnp
import optax
import time11.1.1 Components of a Transformer Block
A transformer block contains two sublayers connected by a residual stream. Section 10.6 introduced the residual stream view: each token carries a running vector \(\mathbf{x} \in \mathbb{R}^d\) from layer to layer, and sublayers do not replace it — they read from it, compute, and add their result back. The block’s two sublayers divide the labor:
- Attention is the only place where positions interact. Each token queries the others as in Section 10.3 and adds the retrieved mixture to its stream. Contextual information enters through this sublayer.
- The feed-forward network (FFN) acts on each position separately: the same two- or three-matrix MLP is applied token by token, with no interaction between positions. It transforms each token’s accumulated state and contains most of the block’s parameters.
Writing \(\mathrm{Attn}\) for multi-head self-attention and \(\mathrm{Norm}\) for a normalization layer, the modern (pre-norm) block computes
\[ \mathbf{h} = \mathbf{x} + \mathrm{Attn}(\mathrm{Norm}_1(\mathbf{x})), \qquad \mathbf{y} = \mathbf{h} + \mathrm{FFN}(\mathrm{Norm}_2(\mathbf{h})). \tag{11.1.1}\]
Both sublayers preserve the shape \((n, d)\): the block maps a sequence of \(d\)-dimensional vectors to a sequence of the same length and width, which is what makes stacking dozens of copies possible. Note also what the block does not contain: positions. Attention is permutation equivariant (Section 10.4), and the block does nothing to change that — positional information is the model’s job, not the block’s, a division of responsibilities the next section exploits.
Figure 11.1.1 shows the wiring, in the two normalization arrangements compared below.
11.1.2 Where the Normalization Goes
Residual connections and normalization entered the block together in 2017 (Vaswani et al. 2017): without some control on scale, the sum of dozens of sublayer outputs drifts, and training deep stacks becomes a tuning nightmare. The normalizer of choice is layer normalization (Ba et al. 2016), which standardizes each token’s vector across its \(d\) features — unlike the batch normalization of Section 7.3, it involves no batch statistics, so it behaves identically in training and inference and is indifferent to sequence length. The 2017 paper also fixed where it goes, and a few years of painful experience revised that answer.
11.1.2.1 Two Arrangements
The original block normalizes the stream after each addition (“post-LN”):
\[ \mathbf{h} = \mathrm{Norm}_1(\mathbf{x} + \mathrm{Attn}(\mathbf{x})), \qquad \mathbf{y} = \mathrm{Norm}_2(\mathbf{h} + \mathrm{FFN}(\mathbf{h})). \tag{11.1.2}\]
The modern block Equation 11.1.1 normalizes each branch’s input (“pre-LN”) and leaves the stream alone. In code the two differ by two lines. The minimal block below takes the arrangement as a flag; its attention is the d2l.MultiHeadAttention of Section 10.3 and its FFN is the classic two-matrix MLP (we upgrade both ingredients later in this section).
class MiniBlock(nn.Module):
"""Attention + MLP on a residual stream; norm placement as a flag."""
def __init__(self, num_hiddens, pre_norm):
super().__init__()
self.pre_norm = pre_norm
self.attention = d2l.MultiHeadAttention(num_hiddens, 4, dropout=0)
self.W_1 = nn.Linear(num_hiddens, 4 * num_hiddens)
self.W_2 = nn.Linear(4 * num_hiddens, num_hiddens)
self.norm1 = nn.LayerNorm(num_hiddens)
self.norm2 = nn.LayerNorm(num_hiddens)
def ffn(self, X):
return self.W_2(F.gelu(self.W_1(X)))
def forward(self, X):
if self.pre_norm:
Y = self.norm1(X)
X = X + self.attention(Y, Y, Y, None)
return X + self.ffn(self.norm2(X))
X = self.norm1(X + self.attention(X, X, X, None))
return self.norm2(X + self.ffn(X))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))11.1.2.2 Signal Propagation at Initialization
Which arrangement is better? We can get a long way without training anything. A freshly initialized network is a fixed random function, so we can measure how it treats signals. We stack \(N\) blocks, pass a random sequence through them, and measure the output without training. Fixing the seed makes the experiment deterministic. We initialize every weight matrix as the original transformer did (Xavier initialization) and track three quantities: the scale of the residual stream, how distinct the tokens remain from one another, and how much gradient each block’s attention receives from a generic loss at the top.
For the distinctness measure we use the cross-token spread \(\|\mathbf{H} - \bar{\mathbf{H}}\| / \|\mathbf{H}\|\), where \(\bar{\mathbf{H}}\) copies the average token to every position: it is \(1\) when tokens are unrelated and \(0\) when they have collapsed onto a single point. One caution learned the hard way: after a final normalization layer the naive probe loss \(\|\mathbf{H}\|^2\) is constant by construction, so we probe with a generic linear readout \(\langle \mathbf{R}, \mathbf{H} \rangle\) for fixed random \(\mathbf{R}\) instead.
def signal_stats(pre_norm, num_blks=32, num_hiddens=256):
torch.manual_seed(0)
blks = nn.ModuleList([MiniBlock(num_hiddens, pre_norm)
for _ in range(num_blks)])
torch.manual_seed(1)
X = torch.randn(2, 64, num_hiddens)
H = X
for blk in blks:
H = blk(H) # materialize the lazy layers
for p in blks.parameters():
if p.ndim == 2: # the 2017 initialization
nn.init.xavier_uniform_(p)
H, spreads, rms = X, [], []
for blk in blks:
H = blk(H)
spreads.append(
((H - H.mean(1, keepdim=True)).norm() / H.norm()).item())
rms.append(H.pow(2).mean().sqrt().item())
(H * torch.randn_like(H)).mean().backward() # generic readout loss
grads = [blk.attention.W_q.weight.grad.norm().item() for blk in blks]
proj = {n: getattr(blks[-1].attention, n).weight.grad.norm().item()
for n in ('W_q', 'W_k', 'W_v', 'W_o')}
return spreads, rms, grads, proj
stats = {pre: signal_stats(pre) for pre in (False, True)}
for pre, name in ((False, 'post-LN'), (True, 'pre-LN')):
spreads, rms, grads, proj = stats[pre]
print(f'{name:>8}: stream RMS at k=32: {rms[-1]:5.2f}, '
'token spread at k=8,16,32: '
+ ' '.join(f'{spreads[k-1]:8.1e}' for k in (8, 16, 32)))
print(' grad norms at k=32: '
+ ' '.join(f'{n} {g:7.1e}' for n, g in proj.items()))
d2l.plot(torch.arange(1, 33), [stats[False][2], stats[True][2]],
'block index k', 'query projection gradient norm',
legend=['post-LN', 'pre-LN'], yscale='log') post-LN: stream RMS at k=32: 1.00, token spread at k=8,16,32: 3.7e-01 3.6e-02 1.2e-04
grad norms at k=32: W_q 2.9e-09 W_k 2.6e-09 W_v 5.8e-02 W_o 5.9e-02
pre-LN: stream RMS at k=32: 5.25, token spread at k=8,16,32: 7.4e-01 5.2e-01 3.7e-01
grad norms at k=32: W_q 1.9e-03 W_k 1.9e-03 W_v 8.1e-02 W_o 8.4e-02
class Stack(nnx.Module):
def __init__(self, num_blks, pre_norm, num_hiddens=256, rngs=None):
rngs = nnx.Rngs(0) if rngs is None else rngs
self.blks = nnx.List([MiniBlock(num_hiddens, pre_norm, rngs)
for _ in range(num_blks)])
def __call__(self, X):
for blk in self.blks:
X = blk(X)
return X
def signal_stats(pre_norm, num_blks=32, num_hiddens=256):
stack = Stack(num_blks, pre_norm, num_hiddens, rngs=nnx.Rngs(0))
init, key = nnx.initializers.xavier_uniform(), jax.random.key(0)
flat = []
for path, p in nnx.to_flat_state(nnx.state(stack, nnx.Param)):
if p[...].ndim == 2: # the 2017 initialization
key, sub = jax.random.split(key)
p = jax.tree.map(lambda v: init(sub, v.shape), p)
flat.append((path, p))
nnx.update(stack, nnx.from_flat_state(flat))
X = jax.random.normal(jax.random.key(1), (2, 64, num_hiddens))
R = jax.random.normal(jax.random.key(2), (2, 64, num_hiddens))
H, spreads, rms = X, [], []
for blk in stack.blks:
H = blk(H)
spreads.append(float(jnp.linalg.norm(H - H.mean(1, keepdims=True))
/ jnp.linalg.norm(H)))
rms.append(float(jnp.sqrt((H ** 2).mean())))
grads = nnx.grad(lambda m: (m(X) * R).mean())(stack) # generic readout
top = grads['blks'][num_blks - 1]['attention']
proj = {n: float(jnp.linalg.norm(top[n]['kernel'][...]))
for n in ('W_q', 'W_k', 'W_v', 'W_o')}
grads = [float(jnp.linalg.norm(
grads['blks'][k]['attention']['W_q']['kernel'][...]))
for k in range(num_blks)]
return spreads, rms, grads, proj
stats = {pre: signal_stats(pre) for pre in (False, True)}
for pre, name in ((False, 'post-LN'), (True, 'pre-LN')):
spreads, rms, grads, proj = stats[pre]
print(f'{name:>8}: stream RMS at k=32: {rms[-1]:5.2f}, '
'token spread at k=8,16,32: '
+ ' '.join(f'{spreads[k-1]:8.1e}' for k in (8, 16, 32)))
print(' grad norms at k=32: '
+ ' '.join(f'{n} {g:7.1e}' for n, g in proj.items()))
d2l.plot(jnp.arange(1, 33), [stats[False][2], stats[True][2]],
'block index k', 'query projection gradient norm',
legend=['post-LN', 'pre-LN'], yscale='log') post-LN: stream RMS at k=32: 1.00, token spread at k=8,16,32: 3.5e-01 2.5e-02 2.1e-04
grad norms at k=32: W_q 5.2e-09 W_k 4.8e-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-02
The two arrangements produce different networks at the same parameter values. Post-LN renormalizes the stream after every addition, so its RMS is pinned at \(1\) at every depth; pre-LN only ever adds, so its stream grows like the square root of the depth (the block outputs are nearly independent at initialization, and variances add). That growth is not a defect — it is a protection. Each pre-LN branch reads the stream through its own normalization, so as the stream grows, each successive block’s relative contribution shrinks, and the function stays tame.
The post-LN stack has no such damping, and the token-spread column shows what happens instead. At initialization attention is nearly uniform, so every attention sublayer pulls each token toward the average token; post-LN then renormalizes, locking in the contraction, and the spread decays geometrically — by block 32 the tokens are identical to about one part in ten thousand. Pre-LN dilutes the same contraction by the growing stream, and the spread falls only polynomially, still at \(0.4\) after 32 blocks. The decreasing token spread is an initialization-time diagnostic related to the rank-collapse phenomenon analyzed by Dong et al. (2021); it does not demonstrate collapse during training. In this 32-block stack, the post-LN query and key projections receive gradients of order \(10^{-9}\), about six orders of magnitude below their pre-LN counterparts, while the value and output projection norms are near \(0.06\). Once token representations become nearly identical, the attention mixture becomes insensitive to the query and key weights, which explains the small gradients in this constructed setting. Different initializations, residual scales, optimizers, and training dynamics can change the effect (Xiong et al. 2020).
The experiments of Xiong et al. (2020) likewise find that post-LN is more sensitive to learning-rate warmup than pre-LN under their training protocols. We therefore use pre-norm as the default in the teaching model and test both placements during training in Section 11.2.
11.1.2.3 RMSNorm
Modern models also simplified the normalizer itself. LayerNorm standardizes each token vector by subtracting its mean and dividing by its standard deviation, then rescales (and re-shifts) feature-wise. RMSNorm (Zhang and Sennrich 2019) drops the centering and the shift, keeping only the part that controls scale:
\[ \mathrm{RMSNorm}(\mathbf{x}) = \frac{\mathbf{x}}{\sqrt{\tfrac{1}{d} \sum_{i=1}^{d} x_i^2 + \epsilon}} \odot \boldsymbol{\gamma}. \tag{11.1.3}\]
It is a one-line function, and both frameworks ship it; let’s verify our implementation against the built-in and time the two normalizers.
def rms_norm(X, weight, eps=1e-6):
return X / torch.sqrt((X * X).mean(-1, keepdim=True) + eps) * weight
device = d2l.try_gpu()
X = torch.randn(64, 512, 1024, device=device)
layernorm = nn.LayerNorm(1024, device=device)
rmsnorm = nn.RMSNorm(1024, eps=1e-6, device=device)
print(f'ours vs. built-in: max deviation '
f'{(rms_norm(X, rmsnorm.weight) - rmsnorm(X)).abs().max():.1e}')
for name, f in (('LayerNorm', layernorm), ('RMSNorm', rmsnorm)):
for _ in range(10):
f(X) # warmup
if device.type == 'cuda':
torch.cuda.synchronize()
t0 = time.time()
for _ in range(200):
f(X)
if device.type == 'cuda':
torch.cuda.synchronize()
print(f'{name}: {(time.time() - t0) / 200 * 1e3:.3f} ms')ours vs. built-in: max deviation 9.5e-07
LayerNorm: 0.389 ms
RMSNorm: 0.412 msdef rms_norm(X, weight, eps=1e-6):
return X / jnp.sqrt((X * X).mean(-1, keepdims=True) + eps) * weight
X = jax.random.normal(jax.random.key(0), (64, 512, 1024))
layernorm = nnx.LayerNorm(1024, rngs=nnx.Rngs(0))
rmsnorm = nnx.RMSNorm(1024, epsilon=1e-6, rngs=nnx.Rngs(0))
print(f'ours vs. built-in: max deviation '
f'{jnp.abs(rms_norm(X, rmsnorm.scale[...]) - rmsnorm(X)).max():.1e}')
for name, f in (('LayerNorm', nnx.jit(layernorm)),
('RMSNorm', nnx.jit(rmsnorm))):
f(X).block_until_ready() # warmup + compile
t0 = time.time()
for _ in range(200):
Y = f(X)
Y.block_until_ready()
print(f'{name}: {(time.time() - t0) / 200 * 1e3:.3f} ms')ours vs. built-in: max deviation 9.5e-07
LayerNorm: 1.072 ms
RMSNorm: 0.719 msIn this benchmark, the PyTorch timings are indistinguishable at the displayed precision, while the JAX RMSNorm kernel is roughly one third faster. These are framework- and device-specific measurements rather than an architectural speed guarantee. RMSNorm omits centering and a bias parameter and computes one statistic instead of two; its quality must be assessed in the surrounding model. Llama is one prominent model family that uses it (Touvron et al. 2023a).
11.1.2.4 Normalizing Queries and Keys
One more normalization migrated inside attention. Nothing in the block above controls the size of the attention logits \(\mathbf{q}^\top \mathbf{k} / \sqrt{d}\) themselves: if training inflates the queries and keys, the softmax saturates, its Jacobian vanishes (Section 10.2), and training can spike or stall — an instability first met at vision-transformer scale, where attention entropy was observed collapsing in billion-parameter runs (Dehghani et al. 2023). QK-norm (Henry et al. 2020) applies RMSNorm to the queries and keys per head, right before the dot product, thereby controlling vector norms before scoring. It costs two extra norms per layer. Reports for Gemma 3, Qwen3, and OLMo 2 document its use in those model families. OLMo 2 combines QK-norm with normalization after each sublayer but outside the residual stream’s identity path (Team OLMo et al. 2025); Gemma 3 normalizes each branch both before and after the sublayer (Gemma Team 2025). Figure 11.1.2 compares these placements. Only original post-LN normalizes the residual stream after addition; the signal experiment above directly compares that arrangement with pre-LN. We leave QK-norm as an exercise rather than a flag; it drops into the block through the attn_factory hook introduced below.
11.1.3 The Feed-Forward Network
Between the attention sublayers sits the block’s other half: a small MLP applied to each position independently,
\[ \mathrm{FFN}(\mathbf{x}) = \mathbf{W}_2\, \phi(\mathbf{W}_1 \mathbf{x}), \tag{11.1.4}\]
with \(\mathbf{W}_1 \in \mathbb{R}^{4d \times d}\) and \(\mathbf{W}_2 \in \mathbb{R}^{d \times 4d}\) in the classic configuration. The factor-4 expansion is a convention the field has never found strong reason to revisit, and it implies something worth noticing: the FFN holds \(8d^2\) parameters against attention’s \(4d^2\), so about two thirds of a block’s parameters sit in these two matrices. When people say a transformer’s knowledge lives mostly in its FFNs, this ratio is the accounting behind the claim.
Two upgrades separate the 2017 FFN from today’s. The activation \(\phi\) moved from ReLU to GELU (Hendrycks and Gimpel 2016), a smoothed relative that ends the debate about the kink at zero (GPT onward (Radford et al. 2018)). Then Shazeer (2020) revived an older idea — gating — and found that a gated linear unit with the SiLU activation beat both:
\[ \mathrm{SwiGLU}(\mathbf{x}) = \mathbf{W}_2 \big(\mathrm{SiLU}(\mathbf{W}_g \mathbf{x}) \odot \mathbf{W}_1 \mathbf{x}\big), \qquad \mathrm{SiLU}(z) = z \cdot \mathrm{sigmoid}(z). \tag{11.1.5}\]
Instead of gating a linear transform through a fixed nonlinearity, SwiGLU computes two linear views of the input and multiplies them, one squashed into a soft gate, the other passed through untouched — the value is transmitted at full strength wherever the gate is open, rather than bent through the activation. The multiplicative interaction is the same trick gates played in LSTMs, here compressed into a single layer. A third matrix means more parameters at the same width, so fair comparisons shrink the hidden width to \(\tfrac{8}{3} d\), making the three matrices of the gated FFN cost the same \(8d^2\) as the classic one, up to rounding. Both variants, in one class:
class FeedForward(nn.Module):
"""Position-wise FFN: GELU MLP or SwiGLU at matched parameter count."""
def __init__(self, num_hiddens, act='swiglu', bias=False):
super().__init__()
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 = nn.Linear(num_hiddens, width, bias=bias)
self.W_1 = nn.Linear(num_hiddens, width, bias=bias)
self.W_2 = nn.Linear(width, num_hiddens, bias=bias)
def forward(self, X):
if self.act == 'gelu':
return self.W_2(F.gelu(self.W_1(X), approximate='tanh'))
return self.W_2(F.silu(self.W_g(X)) * self.W_1(X))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))(The GELU here is the tanh approximation, the exact variant GPT-2 shipped with — a detail that will matter when we load its weights in Section 11.2.) The parameter match is worth checking rather than trusting:
for act in ('gelu', 'swiglu'):
ffn = FeedForward(256, act)
print(f'{act:>7}: {sum(p.numel() for p in ffn.parameters())} parameters') gelu: 524288 parameters
swiglu: 524544 parametersfor act in ('gelu', 'swiglu'):
ffn = FeedForward(256, act)
n = sum(p.size for p in jax.tree.leaves(nnx.state(ffn, nnx.Param)))
print(f'{act:>7}: {n} parameters') gelu: 524288 parameters
swiglu: 524544 parametersWhether the gate improves performance at equal parameter count is an empirical question. The final experiment compares the alternatives in a small language model.
11.1.4 A Configurable Block
We assemble the block used throughout this chapter. Every design decision above becomes a constructor argument. norm selects LayerNorm or RMSNorm, act selects the FFN, pre_norm selects the arrangement, and two factory hooks let later sections swap whole sublayers — a mixture-of-experts FFN, or a cache-friendly grouped-query attention — without touching this class again. The constructor validates its flag strings: a typo must fail loudly, not silently select an architecture. The default configuration (pre-norm, RMSNorm, SwiGLU) is the one you would find inside an open-weights model released this year; TransformerBlock(num_hiddens, num_heads, norm='layer', act='gelu', pre_norm=False, bias=True) is a compact post-norm block in the 2017 encoder arrangement (with one anachronism kept for simplicity: GELU, where the 2017 original used ReLU).
class TransformerBlock(nn.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):
super().__init__()
assert norm in ('rms', 'layer'), f'unknown norm: {norm!r}'
assert num_hiddens % num_heads == 0
self.pre_norm = pre_norm
make_norm = nn.RMSNorm if norm == 'rms' else nn.LayerNorm
self.norm1, self.norm2 = make_norm(num_hiddens), make_norm(num_hiddens)
self.attention = (d2l.MultiHeadAttention(num_hiddens, num_heads,
dropout, bias=bias)
if attn_factory is None else attn_factory())
self.ffn = (FeedForward(num_hiddens, act, bias=bias)
if ffn_factory is None else ffn_factory())
self.dropout = nn.Dropout(dropout)
def forward(self, X, valid_lens=None):
if self.pre_norm:
Y = self.norm1(X)
X = X + self.dropout(self.attention(Y, Y, Y, valid_lens))
return X + self.dropout(self.ffn(self.norm2(X)))
X = self.norm1(X + self.dropout(self.attention(X, X, X, valid_lens)))
return self.norm2(X + self.dropout(self.ffn(X)))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)))11.1.4.1 Shapes and Parameters
The block preserves its input shape whatever the flags, and its parameter census confirms the accounting from above: \(4d^2\) in attention, \(8d^2\) in the FFN (however the activation slices it), and a rounding error of normalization weights — about \(12d^2\) per block, two thirds of it FFN.
X = torch.ones(2, 10, 256)
blk = TransformerBlock(256, num_heads=8)
blk2017 = TransformerBlock(256, num_heads=8, norm='layer', act='gelu',
pre_norm=False, bias=True)
d2l.check_shape(blk(X), X.shape)
d2l.check_shape(blk2017(X), X.shape)
count = lambda m: sum(p.numel() for p in m.parameters())
print(f'attention {count(blk.attention)}, ffn {count(blk.ffn)}, '
f'total {count(blk)} = {count(blk) / 256 ** 2:.2f} d^2')attention 262144, ffn 524544, total 787200 = 12.01 d^2X = jnp.ones((2, 10, 256))
blk = TransformerBlock(256, num_heads=8)
blk2017 = TransformerBlock(256, num_heads=8, norm='layer', act='gelu',
pre_norm=False, bias=True)
d2l.check_shape(blk(X), X.shape)
d2l.check_shape(blk2017(X), X.shape)
count = lambda m: sum(p.size for p in jax.tree.leaves(
nnx.state(m, nnx.Param)))
print(f'attention {count(blk.attention)}, ffn {count(blk.ffn)}, '
f'total {count(blk)} = {count(blk) / 256 ** 2:.2f} d^2')attention 262144, ffn 524544, total 787200 = 12.01 d^2One block is not yet a language model, but it is close enough to train. The same valid_lens mechanism that masked padding in Section 10.2 makes the block causal: pass each query a valid length equal to its own position, and every token attends only backward. A dozen lines wrap an embedding, a stack of blocks, and a tied output head (as in Section 10.4) around that trick:
class CharLM(nn.Module):
"""Minimal char LM: causal masking via per-query valid lengths."""
def __init__(self, vocab_size, num_hiddens=128, num_heads=4, num_blks=2,
act='swiglu', max_len=128):
super().__init__()
self.token_emb = nn.Embedding(vocab_size, num_hiddens)
self.pos_emb = nn.Embedding(max_len, num_hiddens)
for emb in (self.token_emb, self.pos_emb):
nn.init.normal_(emb.weight, std=0.02)
self.blks = nn.ModuleList([
TransformerBlock(num_hiddens, num_heads, act=act)
for _ in range(num_blks)])
self.norm = nn.RMSNorm(num_hiddens)
def forward(self, X):
B, T = X.shape
H = self.token_emb(X) + self.pos_emb(
torch.arange(T, device=X.device))
causal = torch.arange(1, T + 1, device=X.device).repeat(B, 1)
for blk in self.blks:
H = blk(H, causal)
return F.linear(self.norm(H), self.token_emb.weight)class CharLM(nnx.Module):
"""Minimal char LM: causal masking via per-query valid lengths."""
def __init__(self, vocab_size, num_hiddens=128, num_heads=4, num_blks=2,
act='swiglu', max_len=128, rngs=None):
rngs = nnx.Rngs(0) if rngs is None else rngs
init = nnx.initializers.normal(0.02)
self.token_emb = nnx.Embed(vocab_size, num_hiddens,
embedding_init=init, rngs=rngs)
self.pos_emb = nnx.Embed(max_len, num_hiddens,
embedding_init=init, rngs=rngs)
self.blks = nnx.List([
TransformerBlock(num_hiddens, num_heads, act=act, rngs=rngs)
for _ in range(num_blks)])
self.norm = nnx.RMSNorm(num_hiddens, rngs=rngs)
def __call__(self, X):
B, T = X.shape
H = self.token_emb(X) + self.pos_emb(jnp.arange(T))
causal = jnp.tile(jnp.arange(1, T + 1), (B, 1))
for blk in self.blks:
H = blk(H, causal)
return self.token_emb.attend(self.norm(H))11.1.4.2 Comparing GELU and SwiGLU
We compare the two alternatives on the same data (the character-level Time Machine corpus of Section 8.2), holding the model, seed, optimizer, learning rate, and 600-step budget fixed. Only the act flag differs, and the parameter counts match to a tenth of a percent.
data = d2l.TimeMachine(batch_size=64, num_steps=128, tokenization='char',
num_train=100000, num_val=3000)
for act in ('gelu', 'swiglu'):
torch.manual_seed(0)
model = CharLM(len(data.vocab), act=act)
losses = d2l.train_lm(model, data,
torch.optim.AdamW(model.parameters(), lr=1e-3,
weight_decay=0.0), 600)
print(f'{act:>7}: loss at step 200/400/600: ' + '/'.join(
f'{sum(losses[k-100:k]) / 100:.2f}' for k in (200, 400, 600))) gelu: loss at step 200/400/600: 2.26/1.65/1.31
swiglu: loss at step 200/400/600: 2.24/1.51/1.17data = d2l.TimeMachine(batch_size=64, num_steps=128, tokenization='char',
num_train=100000, num_val=3000)
for act in ('gelu', 'swiglu'):
model = CharLM(len(data.vocab), act=act, rngs=nnx.Rngs(0))
optimizer = nnx.Optimizer(model, optax.adamw(1e-3, weight_decay=0.0),
wrt=nnx.Param)
losses = d2l.train_lm(model, data, optimizer, 600)
print(f'{act:>7}: loss at step 200/400/600: ' + '/'.join(
f'{sum(losses[k-100:k]) / 100:.2f}' for k in (200, 400, 600))) gelu: loss at step 200/400/600: 2.27/1.72/1.40
swiglu: loss at step 200/400/600: 2.22/1.57/1.24The gated FFN ends more than a tenth of a nat ahead. This margin holds up across seeds, since rerunning with seeds 1 and 2 moves each number by a couple of hundredths, not the gap. More importantly, it agrees in direction with Shazeer (2020)’s systematic sweep and with the consistent choice of the major model families since Llama. It is a modest, real improvement at equal parameter count. Such incremental gains can accumulate when several architectural choices are combined.
The diagnostics in this section use one initialization, a 32-block synthetic stack, and framework-specific normalization kernels. They establish algebraic identities and report local signal, gradient, timing, and loss measurements; they do not by themselves establish causal explanations or universal model rankings. The cited large-scale studies provide separate evidence for the training practices discussed above.
11.1.5 Summary
A transformer block contains attention and a position-wise FFN connected through a residual stream. Attention communicates between positions; the FFN transforms each position and contains about two thirds of the parameters. In the original post-LN arrangement, normalization fixes the scale of the residual stream, but the query and key projections near the top of a 32-block stack receive gradients six orders of magnitude smaller than those near the bottom. Pre-norm instead normalizes each branch input and preserves usable gradients throughout the stack. RMSNorm retains only the scale statistic, while QK-norm controls the scale of attention logits. Finally, SwiGLU replaces a fixed activation with a learned gate and gives a small, seed-stable improvement at matched parameter count in our experiment. The resulting TransformerBlock exposes normalization, activation, and placement as options, with replaceable attention and FFN modules.
11.1.6 Exercises
- At initialization the block outputs are nearly independent, so the pre-norm stream RMS should grow like \(\sqrt{1 + 2N}\) after \(N\) blocks with two sublayers each — up to the sublayers’ output scale. Extend the signal-propagation experiment to fit the growth exponent from a log-log plot. How close is it to \(1/2\), and why does the attention sublayer contribute less than the FFN sublayer early in the stack?
- OLMo 2 places the normalization after each sublayer but off the stream: \(\mathbf{h} = \mathbf{x} +
\mathrm{Norm}(\mathrm{Attn}(\mathbf{x}))\). Add this third arrangement to
MiniBlockand rerun the experiment. Which of post-LN’s pathologies does it avoid, which of pre-LN’s properties does it give up, and what happens to the stream RMS? - The signal-propagation experiment uses the 2017 Xavier initialization. Rerun it with each framework’s default initialization and with all weights scaled by an extra factor of \(0.5\). How does the collapse depth change? Relate your observations to why post-LN models of GPT-1’s era (12 blocks) were trainable at all.
- Implement two more of Shazeer (2020)’s gated variants by changing one line of
FeedForward: ReGLU (ReLU gate) and GEGLU (GELU gate). Race all four at matched parameters onCharLM. Do your rankings agree with the paper’s? What does the spread between the gated variants, compared with the gap to the ungated MLP, tell you about which design choice carries the improvement? - Work out the exact parameter count of
TransformerBlock(d, h)for bothactsettings, including normalization weights, and check it against the census cell. Real models round the SwiGLU width up to a multiple of 256 for hardware efficiency; how far from \(8d^2\) does the Llama-7B configuration (\(d = 4096\), FFN width \(11008\)) land? - Implement QK-norm as an
attn_factory: wrapd2l.MultiHeadAttentionso that queries and keys pass through an RMSNorm (one per head dimension) before the dot product. At initialization, measure the standard deviation of the attention logits with and without it as you scalenum_hiddensfrom 64 to 1024. Which curve does Equation 10.2.1’s \(1/\sqrt{d}\) scaling predict, and which do you observe?