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.
image/svg+xml Matplotlib v3.10.8, https://matplotlib.org/

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}))

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))

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.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

Reading 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: 0.681 ms
RMSNorm: 0.652 ms

No 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:

image/svg+xml Matplotlib v3.10.8, https://matplotlib.org/

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(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))

One block, every configuration

Norm, activation, arrangement as flags; attention and FFN as factories (MoE and grouped-query attention drop in later):

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)))

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.26/1.65/1.31
 swiglu: loss at step 200/400/600: 2.24/1.51/1.17

More 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.