Complete Transformer Decoder

Introduction

In the previous section we built one TransformerDecoderLayer — a single block that combines masked self-attention, cross-attention over encoder memory, and a position-wise FFN with LayerNorm and residuals between them. One layer, however, is not a decoder. A decoder that can translate, summarize, or generate language needs three additional pieces:

1. A token embedding that turns integer ids into d_model-dim vectors.
2. A positional encoding that injects order information.
3. A stack of N decoder layers that iteratively refine the hidden state, followed by an output projection that turns the final hidden state into vocabulary logits.

This section assembles those pieces into TransformerDecoder — the second-to-last building block before the full encoder-decoder Transformer in §6. We'll cover the math of the stacked forward pass, the clever parameter-sharing trick called weight tying, a step-by-step numerical walkthrough, the PyTorch implementation, and a careful parameter count.

Why Stack Decoder Layers

Shallow models can't capture hierarchical structure

A single attention block can only mix information once. Language has structure at many scales — morphology, syntax, semantics, discourse — and those scales compose. No finite amount of width can make up for depth: you eventually need to build abstractions on top of abstractions, which is exactly what stacking does.

Empirically, translation quality monotonically improves with depth up to N=6 in Vaswani et al. 2017 and keeps scaling in modern large language models. Kaplan et al. 2020 ("Scaling Laws for Neural Language Models") fit loss ≈ N^(−0.076), and Hoffmann et al. 2022 (Chinchilla) showed that compute-optimal training requires balancing parameters and data — adding layers alone isn't enough if you don't train on proportionally more tokens.

What each layer learns — the BERTology view, adapted

Tenney et al. 2019 ("BERT Rediscovers the Classical NLP Pipeline") and Rogers et al. 2020 ("A Primer in BERTology") probed encoder layers and found a rough hierarchy: shallow layers handle local and morphological patterns, middle layers capture syntax, and deep layers represent semantics and discourse. Decoders exhibit an analogous gradient. Shallow decoder layers tend to specialize in copying source content via cross-attention; middle layers handle local target-side fluency; deep layers assemble longer-range coherence across the target sentence.

Why This Matters: Stacking is not just repetition. Each layer starts where the last one stopped — so Layer 3 sees a syntactically organized representation, not raw embeddings. The shared weights between parallel heads in one layer give you breadth; stacking layers gives you depth. Both are needed.

Decoder Forward Pass

With the target input tokens $\mathrm{tgt}$ and the encoder output $\mathrm{memory}$ in hand, the decoder's forward pass has three stages:

Stage 1 — build h₀

Look up embeddings, scale them, add positional encoding, apply dropout: $h_0 = \mathrm{Dropout}\bigl(\mathrm{Embed}(\mathrm{tgt}) \cdot \sqrt{d_{model}} + \mathrm{PE}(\mathrm{tgt})\bigr)$.

The $\sqrt{d_{model}}$ scaling is explained in Vaswani 2017 §3.4. PyTorch's nn.Embedding initializes entries around $\mathcal{N}(0, 1)$, which means each embedding vector has variance per feature ≈ 1. Sinusoidal PE has amplitude ≈ 1 as well. Without scaling, the PE component would dominate the embedding. Multiplying the embedding by $\sqrt{d_{model}}$ (≈ 22.6 when d_model = 512) restores a balance where the identity of the token and its position both matter.

Stage 2 — iterate the stack

For layers $\ell = 1, \dots, N$: $h_\ell = \mathrm{DecoderLayer}_\ell(h_{\ell-1}, \mathrm{memory}, \mathrm{tgt\_mask}, \mathrm{memory\_mask})$.

Every layer receives the same memory and masks; only $h$ changes. The layer internals — masked self-attn, cross-attn, FFN, pre-norm, residuals, dropout — were specified in §4 and are not reopened here.

Stage 3 — project to vocabulary

A final LayerNorm $h_N \leftarrow \mathrm{LayerNorm}(h_N)$ (required when using pre-norm — without it the last residual sum has unbounded scale), then: $\mathrm{logits} = h_N \, W_{out}^\top$.

The logits have shape $[B, T_{tgt}, V]$ where V is the target vocabulary size. Softmax over the last axis gives token probabilities; cross-entropy against the gold ids gives the training loss.

Weight Tying

Notice that the embedding table $\mathrm{Embed} \in \mathbb{R}^{V \times d_{model}}$ and the output projection $W_{out} \in \mathbb{R}^{V \times d_{model}}$ have the same shape. Press & Wolf ("Using the Output Embedding to Improve Language Models", 2017) proposed using the same matrix for both: $W_{out} = \mathrm{Embed}$.

This has two benefits. First, it halves the parameters at the two largest matrices — for $V = 50{,}000$ and $d_{model} = 512$ that's 25.6M saved parameters. Second, it often improves quality: the inner product $h \cdot \mathrm{Embed}_w$ directly measures how close $h$ is to the embedding of candidate word $w$, which is the natural dual of the lookup operation done at the input.

When to tie: share a single vocabulary (language modeling, code models, shared-vocab translation). When not: separate source and target vocabularies (classical NMT). You could still tie the target embedding with the output projection — that's actually how Vaswani 2017 did it for their WMT models.

Plain-Python Stack Walkthrough

Before we reach for PyTorch, let's stack two tiny decoder layers by hand on the shared config $(B{=}1, T_{src}{=}4, T_{tgt}{=}3, d_{model}{=}8, H{=}2, d_{ff}{=}16, N{=}2)$. We'll watch the hidden state $h_0 \to h_1 \to h_2$ evolve across the stack using deterministic seeds so every number is reproducible.

        d0     d1     d2     d3     d4     d5     d6     d7
tok0  0.882  0.200  0.489  1.120  0.934 -0.489  0.475 -0.076
tok1 -0.052  0.205  0.072  0.727  0.381  0.061  0.222  0.167
tok2  0.747 -0.103  0.157 -0.427 -1.276  0.327  0.432 -0.371

        d0     d1     d2     d3     d4     d5     d6     d7
tok0  0.882  0.200  0.489  1.120  0.934 -0.489  0.475 -0.076
tok1 -0.052  0.205  0.072  0.727  0.381  0.061  0.222  0.167
tok2  0.747 -0.103  0.157 -0.427 -1.276  0.327  0.432 -0.371

        d0     d1     d2     d3     d4     d5     d6     d7
tok0  0.530  0.080  0.516  0.581  0.476 -0.617 -0.504 -0.161
tok1 -0.481  0.278  0.286  0.388 -0.082 -0.387 -1.063 -0.353
tok2  1.450  0.643  0.626 -0.598 -2.060 -0.356 -0.788 -0.535

        d0     d1     d2     d3     d4     d5     d6     d7
tok0  0.822 -0.090  0.344  0.919  0.321 -0.648 -1.121  1.099
tok1 -0.085  0.031  0.075  0.987 -0.478 -0.556 -1.299 -0.026
tok2  1.363  0.431  0.835 -0.339 -2.172  0.040 -1.211 -0.804

1import numpy as np
2
3np.random.seed(0)
4
5# Shared tiny config used in §4, §5, §6
6B, T_src, T_tgt, d_model, H, d_ff, N = 1, 4, 3, 8, 2, 16, 2
7d_k = d_model // H  # 4
8
9# Already-prepared h_0: Embed(tgt) * sqrt(d_model) + PE(tgt), then dropout.
10# We build it by hand so the numbers are traceable.
11tgt_in = (np.random.randn(T_tgt, d_model) * 0.5).astype(np.float32)
12memory = (np.random.randn(T_src, d_model) * 0.5).astype(np.float32)
13
14def layernorm(x, eps=1e-5):
15    mu = x.mean(-1, keepdims=True)
16    var = x.var(-1, keepdims=True)
17    return (x - mu) / np.sqrt(var + eps)
18
19def softmax(x, axis=-1):
20    x = x - x.max(axis=axis, keepdims=True)
21    e = np.exp(x)
22    return e / e.sum(axis=axis, keepdims=True)
23
24def mha(Q, K, V, W, causal=False):
25    q = (Q @ W["q"]).reshape(-1, H, d_k).transpose(1, 0, 2)
26    k = (K @ W["k"]).reshape(-1, H, d_k).transpose(1, 0, 2)
27    v = (V @ W["v"]).reshape(-1, H, d_k).transpose(1, 0, 2)
28    scores = q @ k.transpose(0, 2, 1) / np.sqrt(d_k)
29    if causal:
30        Tq = scores.shape[-2]
31        mask = np.triu(np.ones((Tq, Tq)), k=1).astype(bool)
32        scores = np.where(mask, -1e9, scores)
33    out = (softmax(scores) @ v).transpose(1, 0, 2).reshape(-1, d_model)
34    return out @ W["o"]
35
36def ffn(x, W):
37    return np.maximum(0, x @ W["w1"]) @ W["w2"]
38
39def decoder_layer(x, memory, W):
40    x = x + mha(layernorm(x), layernorm(x), layernorm(x), W["sa"], causal=True)
41    x = x + mha(layernorm(x), layernorm(memory), layernorm(memory), W["ca"])
42    x = x + ffn(layernorm(x), W["ff"])
43    return x
44
45def make_layer(seed):
46    r = np.random.default_rng(seed).standard_normal
47    def P(shape):
48        return (r(shape) * 0.2).astype(np.float32)
49    return {
50        "sa": {"q": P((d_model, d_model)), "k": P((d_model, d_model)),
51               "v": P((d_model, d_model)), "o": P((d_model, d_model))},
52        "ca": {"q": P((d_model, d_model)), "k": P((d_model, d_model)),
53               "v": P((d_model, d_model)), "o": P((d_model, d_model))},
54        "ff": {"w1": P((d_model, d_ff)), "w2": P((d_ff, d_model))},
55    }
56
57layers = [make_layer(1), make_layer(2)]
58
59h0 = tgt_in
60h1 = decoder_layer(h0, memory, layers[0])
61h2 = decoder_layer(h1, memory, layers[1])

The takeaway from the printed values of $h_0, h_1, h_2$ is that hidden states really do move as they pass through the stack. Token 2 ("tok2") starts at $d_4 = -1.276$, ends at $d_4 = -2.172$; its $d_0$ rises from 0.747 to 1.363. These shifts are not noise — they are the cumulative effect of attention and FFN remixing features using information the layer couldn't access before (the target's own history via self-attn, the source via cross-attn).

PyTorch Implementation

Now the production version. The class TransformerDecoder composes TransformerDecoderLayer (from §4), a token embedding, a positional encoding, a final LayerNorm, and an output projection — with the weight-tying toggle baked in.

1import math
2import torch
3import torch.nn as nn
4from typing import Optional
5
6# Assume TransformerDecoderLayer is the module built in §4.
7from .decoder_layer import TransformerDecoderLayer
8
9
10class SinusoidalPositionalEncoding(nn.Module):
11    """Copy of the ch04 PE module. Included for completeness."""
12
13    def __init__(self, d_model: int, max_len: int = 5000):
14        super().__init__()
15        pe = torch.zeros(max_len, d_model)
16        pos = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
17        div = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
18        pe[:, 0::2] = torch.sin(pos * div)
19        pe[:, 1::2] = torch.cos(pos * div)
20        self.register_buffer("pe", pe.unsqueeze(0))  # (1, max_len, d_model)
21
22    def forward(self, x: torch.Tensor) -> torch.Tensor:
23        return x + self.pe[:, : x.size(1)]
24
25
26class TransformerDecoder(nn.Module):
27    """Full Transformer decoder: embedding + PE + N layers + output projection."""
28
29    def __init__(
30        self,
31        num_layers: int,
32        d_model: int,
33        num_heads: int,
34        d_ff: int,
35        vocab_size: int,
36        max_len: int = 5000,
37        dropout: float = 0.1,
38        tie_weights: bool = True,
39    ):
40        super().__init__()
41        self.d_model = d_model
42        self.tok_emb = nn.Embedding(vocab_size, d_model)
43        self.pos_enc = SinusoidalPositionalEncoding(d_model, max_len)
44        self.drop = nn.Dropout(dropout)
45        self.layers = nn.ModuleList([
46            TransformerDecoderLayer(d_model, num_heads, d_ff, dropout)
47            for _ in range(num_layers)
48        ])
49        self.norm = nn.LayerNorm(d_model)
50        self.out_proj = nn.Linear(d_model, vocab_size, bias=False)
51        if tie_weights:
52            self.out_proj.weight = self.tok_emb.weight  # W_out = Embed
53
54    def forward(
55        self,
56        tgt: torch.Tensor,
57        memory: torch.Tensor,
58        tgt_mask: Optional[torch.Tensor] = None,
59        memory_mask: Optional[torch.Tensor] = None,
60    ) -> torch.Tensor:
61        h = self.tok_emb(tgt) * math.sqrt(self.d_model)
62        h = self.drop(self.pos_enc(h))
63        for layer in self.layers:
64            h = layer(h, memory, tgt_mask=tgt_mask, memory_mask=memory_mask)
65        h = self.norm(h)
66        logits = self.out_proj(h)
67        return logits

Note on KV caching: at inference time we want to feed one token at a time and reuse the keys/values the earlier layers already produced — this is called KV caching, and it becomes critical once your stack is 96 layers deep. We cover it in §9/05.

Parameter Counting

Let's tally parameters for a small demo configuration: $d_{model} = 128$, $H = 4$, $d_{ff} = 512$, $N = 2$, $V = 1000$. Per-layer numbers assume bias on every linear layer (Vaswani 2017) and three LayerNorms per layer.

Weight tying saves $V \cdot d_{model} + V$ parameters — here 129,000 out of 786,152, roughly 16% of the model. At production scales (V=50k, d_model=4096), weight tying saves ~200M parameters.

Interactive Visualization

The visualizer below shows an N-deep decoder stack in 3D. Each layer is a column of three translucent planes — Self-Attn, Cross-Attn, FFN — with the encoder memory block sitting to the side. Press Step to push activations one sublayer at a time from the bottom of the stack upward, or Play for automatic stepping. Adjust N (1–6) and the grid resolution (a proxy for d_model) to see how the stack grows. When a cross-attention sublayer is active, a glowing line from the encoder memory to that layer indicates where K and V are sourced.

Accessibility notes: all controls are tab-reachable with Enter/Space activation; the prefers-reduced-motion setting disables the auto-advance and you must use Step manually; on viewports under 768px a static SVG diagram is shown instead of the 3D canvas.

In Modern Systems

Depth across real models

Scaling laws

Kaplan et al. 2020 showed cross-entropy loss scales roughly as a power law in model size, data, and compute — with model-size exponents near $N^{-0.076}$ in the regime they studied. Hoffmann et al. 2022 (Chinchilla) revised this picture: for a fixed compute budget, parameters and training tokens should scale in roughly equal ratio. Doubling N without doubling data is wasteful. This reframed many post-2022 training runs (e.g., Llama-2 trained on 2T tokens at 7–70B parameters).

Why decoders got very deep, very fast

In encoder-decoder translation models, depth tapered around N = 6–12 because most quality gains came from wider models and more data. Decoder-only LMs flipped that: Transformer weights are the only place a language model stores knowledge, so deep stacks directly increase model capacity. GPT-3's jump from 48 to 96 layers (and d_model from 1600 to 12288) bought dramatic gains in in-context learning. This is why modern frontier models are much deeper than the original Vaswani decoder.

Summary

• A decoder is embedding + PE + N decoder layers + final norm + output projection.
• Stack forward pass: $h_0 = \mathrm{Dropout}(\mathrm{Embed}(\mathrm{tgt}) \cdot \sqrt{d_{model}} + \mathrm{PE})$; $h_\ell = \mathrm{DecoderLayer}_\ell(h_{\ell-1}, \mathrm{memory}, \ldots)$; $\mathrm{logits} = h_N W_{out}^\top$.
• The $\sqrt{d_{model}}$ factor balances embedding magnitude against the positional encoding amplitude.
• Weight tying (Press & Wolf 2017) sets $W_{out} = \mathrm{Embed}$; use it when input and output vocabularies match.
• Across the stack, hidden states are genuinely refined — not just rescaled. Shallow layers handle local structure, deep layers handle semantics and discourse.
• Real decoders range from N = 6 (Vaswani) to N = 96 (GPT-3) to N = 80 (Llama-65B).
• KV caching (§9/05) will become essential once you run inference on deep stacks.

Exercises

1. (Easy) Modify the plain-Python walkthrough to run N = 4 layers instead of 2. Print the first row of each $h_\ell$ and plot how a single feature (e.g., $h_\ell[0, 0]$) evolves across $\ell = 0, 1, 2, 3, 4$. Is the change monotone, or does it oscillate?
2. (Medium) Extend TransformerDecoder with a tie_weights=False flag and benchmark parameter counts against the tied version for $V \in \{1000, 10000, 50000\}$ at $d_{model}=512$. Confirm that tying saves exactly $V \cdot d_{model}$ parameters (plus $V$ if bias is included — which we disabled).
3. (Hard) Replace the final nn.LayerNorm with nn.Identity and retrain on a small language-modeling task. Track the norm of $h_N$ across training and describe what goes wrong. Why does pre-norm require a final LayerNorm while post-norm does not?

Next Section Preview

We now have a full encoder (ch07) and a full decoder (this section). Section 6 glues them into one Transformer module, walks through training-time teacher forcing vs inference-time autoregressive decoding, discusses shared-vs-separate vocabularies, and finishes with a parameter-count sanity check for the small Multi30k translation model we'll train in ch13.