31. Minimal Pre-Norm Transformer Decoder Block

Problem

Every transformer-based sequence model, no matter how much its attention mechanism is customized, is built from the same two-sublayer block repeated N times: a self-attention sublayer and a position-wise MLP sublayer, each wrapped in a residual connection and a normalization layer. Getting this scaffolding right — where the norm goes, where the residual branches off, that the causal mask actually reaches the attention call — is table-stakes before any more exotic architecture work makes sense.

Implement the forward method of a minimal pre-norm decoder block: x = x + SelfAttention(LayerNorm1(x)), then x = x + MLP(LayerNorm2(x)). __init__ already builds the submodules (two nn.LayerNorms, an nn.MultiheadAttention with batch_first=True, and a two-layer MLP) — the job is only to wire them together correctly in forward, including building the causal attention mask.

Source: src/31_transformer_decoder_block.py

class TransformerDecoderBlock(nn.Module):
    def __init__(self, d_model: int, n_heads: int, d_ff: int): ...
    def forward(self, x: torch.Tensor) -> torch.Tensor: ...  # (B, T, d_model) -> (B, T, d_model)

>>> torch.manual_seed(0)
>>> block = TransformerDecoderBlock(d_model=16, n_heads=2, d_ff=32)
>>> x = torch.randn(2, 5, 16)
>>> block(x).shape
torch.Size([2, 5, 16])

Step-by-Step Approach

1. Normalize x with self.ln1 first — this is pre-norm, so the norm happens before the sublayer, not after (post-norm, the original Transformer paper's order, is less common in modern LLMs because it trains less stably at depth).
2. Build a causal attention mask of shape (T, T): 0 on and below the diagonal, -inf above it, via torch.triu(..., diagonal=1) filled with -inf.
3. Pass the normalized tensor as query, key, AND value into self.attn (this is self-attention) along with the causal mask, and add its output back onto the un-normalized x as a residual connection.
4. Repeat the same normalize-sublayer-residual pattern for the MLP branch: x = x + self.mlp(self.ln2(x)).
5. Double check both sublayers read from the residual stream (x) fresh each time and write back into it — a common bug is normalizing once and reusing the normalized tensor for both sublayers instead of re-normalizing after the first residual add.
6. Verify shape is preserved end to end: input and output are both (B, T, d_model), since a decoder block never changes the sequence length or model dimension.

The key insight is that pre-norm's residual stream (x) is the thing that carries information forward unmodified in magnitude across many stacked blocks; each sublayer only ever contributes an additive delta computed from a normalized copy, which is exactly what keeps very deep transformer stacks trainable.

Reference solution

import torch
import torch.nn as nn


class TransformerDecoderBlock(nn.Module):
    def __init__(self, d_model: int, n_heads: int, d_ff: int):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.ln2 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d_model),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, _ = x.shape

        # pre-norm attention sublayer: norm first, then residual add around the sublayer
        normed = self.ln1(x)  # (B, T, d_model)
        # causal mask: True/-inf marks positions the query may NOT attend to (future positions)
        causal_mask = torch.triu(
            torch.full((T, T), float("-inf"), device=x.device, dtype=x.dtype), diagonal=1
        )
        attn_out, _ = self.attn(normed, normed, normed, attn_mask=causal_mask, need_weights=False)
        x = x + attn_out  # (B, T, d_model)

        # pre-norm MLP sublayer, same residual pattern
        x = x + self.mlp(self.ln2(x))  # (B, T, d_model)
        return x


TEST_CASES = [
    {"name": "shape is preserved", "d_model": 16, "n_heads": 2, "d_ff": 32, "B": 2, "T": 5},
    {"name": "single-head, single-example batch", "d_model": 8, "n_heads": 1, "d_ff": 16, "B": 1, "T": 4},
]


def main():
    torch.manual_seed(0)

    for i, case in enumerate(TEST_CASES):
        d_model, n_heads, d_ff, B, T = case["d_model"], case["n_heads"], case["d_ff"], case["B"], case["T"]
        print(f"Test {i}: {case['name']} (d_model={d_model}, n_heads={n_heads}, d_ff={d_ff}, B={B}, T={T})")

        torch.manual_seed(0)
        block = TransformerDecoderBlock(d_model, n_heads, d_ff).eval()
        x = torch.randn(B, T, d_model)
        print(f"  input x: {tuple(x.shape)}")

        with torch.no_grad():
            out = block(x)
        print(f"  output: {tuple(out.shape)}")
        assert out.shape == (B, T, d_model)
        assert torch.isfinite(out).all()

        x2 = x.clone()
        x2[:, -1, :] = torch.randn(B, d_model)
        with torch.no_grad():
            out2 = block(x2)
        torch.testing.assert_close(out[:, :-1, :], out2[:, :-1, :], atol=1e-5, rtol=1e-5)
        print("  causality check: earlier positions unaffected by last-token perturbation")

        print("PASSED")


if __name__ == "__main__":
    main()

Key Functions & Tricks

• nn.LayerNorm(d_model) — normalizes each token's feature vector to zero mean/unit variance (plus learned scale/shift), applied independently at every sequence position.
• nn.MultiheadAttention(d_model, n_heads, batch_first=True) — PyTorch's built-in multi-head attention; batch_first=True keeps tensors in (B, T, D) layout instead of the default (T, B, D).
• torch.triu(mask, diagonal=1) — keeps only the strictly-upper-triangular part of a matrix, the standard way to build a causal ("can't see the future") additive mask.
• attn(query, key, value, attn_mask=...) — self-attention is just multi-head attention called with the same tensor for all three roles.
• nn.Sequential(Linear, GELU, Linear) — the standard position-wise MLP/FFN sublayer, typically expanding to d_ff > d_model in the hidden layer.
• residual add: x = x + sublayer(norm(x)) — the pre-norm pattern applied identically to both the attention and MLP sublayers.

How to Recognize This Pattern