Encoder-Decoder Transformers

The second half of week 3. Reading a multi-digit number needs an output whose length the model decides, which is what a decoder is for.

The second half of week 3. The classifier looks at one digit and picks one of ten labels. This one reads a whole number.

The difference between the halves of that diagram is larger than it looks. A classifier produces a fixed-size answer: ten scores, pick the biggest. Reading produces an ordered sequence whose length the model has to decide, which means something has to say where it stops.

The encoder is unchanged from the classifier. Everything new is on the right-hand side.

What a Decoder Adds

A start token. The decoder generates one position at a time, each conditioned on what came before. At the first step there is no “before”, so you feed it a token that means begin.

A causal mask. During training the decoder sees the whole target sequence at once, for speed. But position 2 must not see position 3. Otherwise the model learns to read the answer it is being asked to predict.

At inference the future genuinely is not there, so a model trained that way produces nothing useful. The mask stops it: each position sees only itself and everything to its left.

Cross-attention. The decoder needs to look at the image. Self-attention lets the output positions look at each other; cross-attention lets each output position query the encoder’s patch tokens and pull in what it needs.

And then an end token, so the model can say it’s done. That’s what makes the output length the model’s decision rather than a hyperparameter.

import sys

sys.path.insert(0, "..")
import matplotlib.pyplot as plt
import numpy as np
import torch
from _style import COLOURS, MUTED

L = 3
labels = ["<start>", "digit 1", "digit 2", "digit 3"]
mask = np.tril(np.ones((L + 1, L + 1)))

fig, ax = plt.subplots(figsize=(4.4, 4.0))
ax.imshow(mask, cmap="Blues", vmin=0, vmax=1.7)
ax.set_xticks(range(L + 1), labels, rotation=45, ha="right")
ax.set_yticks(range(L + 1), labels)
ax.set_xlabel("can attend to", color=MUTED, fontsize=9)
ax.set_ylabel("generating", color=MUTED, fontsize=9)
ax.tick_params(colors=MUTED, labelsize=9, length=0)
for s in ax.spines.values():
    s.set_visible(False)
fig.tight_layout()

Figure 1: The causal mask. A filled cell means that output position is allowed to attend to that one; the blank upper triangle is the future, hidden.

Teacher Forcing

There’s a subtlety in how this gets trained that took me a while to appreciate.

During training the decoder is fed the true previous digits at every position. That’s called teacher forcing, and it’s what lets the whole sequence be computed in one pass instead of three.

At inference there are no true previous digits. The model is fed its own previous outputs. So if it gets digit 1 wrong, digit 2 is being predicted from a prefix that never occurred during training.

So the training loss systematically flatters the model. The honest measure is generating the whole sequence the way you actually would.

import torch.nn as nn
from huggingface_hub import hf_hub_download

REVISION = "f705fed08827ff6c36e3b5329495c943a5e544e8"
data = np.load(hf_hub_download("roshbeed/ai-residency-blog-data", "mnist/mnist-small.npz",
                               repo_type="dataset", revision=REVISION))

x_train = torch.from_numpy(data["x_train"]).float() / 255.0
y_train = torch.from_numpy(data["y_train"]).long()
x_test = torch.from_numpy(data["x_test"]).float() / 255.0
y_test = torch.from_numpy(data["y_test"]).long()

START, END = 10, 11  # two extra vocabulary entries alongside the ten digits


def compose(images, labels, n, seed):
    """Glue three random digits side by side into one 28x84 image."""
    rng = np.random.default_rng(seed)
    pick = rng.integers(0, len(images), (n, L))
    wide = torch.cat([images[pick[:, i]] for i in range(L)], dim=2)
    return wide, torch.stack([labels[pick[:, i]] for i in range(L)], dim=1)


X, Y = compose(x_train, y_train, 20_000, seed=0)
X_test, Y_test = compose(x_test, y_test, 2_000, seed=1)
print(f"{len(X):,} composites of shape {tuple(X.shape[1:])}, targets of length {L}")
print(f"guessing three digits uniformly: {1 / 1000:.4f} exact-sequence accuracy")

20,000 composites of shape (28, 84), targets of length 3
guessing three digits uniformly: 0.0010 exact-sequence accuracy

PATCH, DIM, HEADS, LAYERS = 7, 96, 4, 2
GRID = (28 // PATCH, 84 // PATCH)          # 4 rows, 12 columns of patches
PATCHES = GRID[0] * GRID[1]


def patchify(images):
    batch = images.shape[0]
    return images.unfold(1, PATCH, PATCH).unfold(2, PATCH, PATCH).reshape(batch, PATCHES, PATCH**2)


class EncoderBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.attention = nn.MultiheadAttention(DIM, HEADS, batch_first=True)
        self.norm1, self.norm2 = nn.LayerNorm(DIM), nn.LayerNorm(DIM)
        self.ff = nn.Sequential(nn.Linear(DIM, 4 * DIM), nn.GELU(), nn.Linear(4 * DIM, DIM))

    def forward(self, x):
        x = self.norm1(x + self.attention(x, x, x, need_weights=False)[0])
        return self.norm2(x + self.ff(x))


class DecoderBlock(nn.Module):
    """Masked self-attention, then cross-attention onto the image, then a feed-forward."""

    def __init__(self):
        super().__init__()
        self.self_attention = nn.MultiheadAttention(DIM, HEADS, batch_first=True)
        self.cross_attention = nn.MultiheadAttention(DIM, HEADS, batch_first=True)
        self.norm1, self.norm2, self.norm3 = (nn.LayerNorm(DIM) for _ in range(3))
        self.ff = nn.Sequential(nn.Linear(DIM, 4 * DIM), nn.GELU(), nn.Linear(4 * DIM, DIM))

    def forward(self, x, memory, mask, want_weights=False):
        x = self.norm1(x + self.self_attention(x, x, x, attn_mask=mask, need_weights=False)[0])
        attended, weights = self.cross_attention(x, memory, memory,
                                                 need_weights=want_weights,
                                                 average_attn_weights=True)
        x = self.norm2(x + attended)
        return self.norm3(x + self.ff(x)), weights

class Reader(nn.Module):
    def __init__(self):
        super().__init__()
        self.project = nn.Linear(PATCH**2, DIM)
        self.image_positions = nn.Parameter(torch.randn(1, PATCHES, DIM) * 0.02)
        self.encoder = nn.ModuleList([EncoderBlock() for _ in range(LAYERS)])

        self.embed = nn.Embedding(12, DIM)   # ten digits, <start>, <end>
        self.token_positions = nn.Parameter(torch.randn(1, L + 1, DIM) * 0.02)
        self.decoder = nn.ModuleList([DecoderBlock() for _ in range(LAYERS)])
        self.out = nn.Linear(DIM, 12)

    def encode(self, images):
        h = self.project(patchify(images)) + self.image_positions
        for block in self.encoder:
            h = block(h)
        return h

    def decode(self, memory, tokens, want_weights=False):
        h = self.embed(tokens) + self.token_positions[:, :tokens.shape[1]]
        mask = torch.triu(torch.full((tokens.shape[1],) * 2, float("-inf")), diagonal=1)

        weights = None
        for i, block in enumerate(self.decoder):
            h, w = block(h, memory, mask, want_weights and i == LAYERS - 1)
            if w is not None:
                weights = w
        return self.out(h), weights

    def forward(self, images, tokens):
        return self.decode(self.encode(images), tokens)[0]

Both halves together. The image enters top left. The tokens enter bottom left through their embedding. The two streams meet in the decoder.

from _arch import diagram

diagram(Reader(), input_shape=((1, 28, 84), (1, L + 1)), style="flow",
        input_dtype=(torch.float32, torch.long))

Figure 2: The encoder-decoder. Two inputs, two stacks, joined by the decoder’s cross-attention.

So the model below is built with two ways to run it: one for training, and exact_sequence for measuring, which generates a digit at a time from the model’s own output the way inference has to.

torch.manual_seed(0)
model = Reader()
optimiser = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
generator = torch.Generator().manual_seed(0)

# teacher forcing: the decoder is fed <start> d1 d2 and must produce d1 d2 d3,
# then <end>.
decoder_input = torch.cat([torch.full((len(Y), 1), START), Y], dim=1)
decoder_target = torch.cat([Y, torch.full((len(Y), 1), END)], dim=1)


@torch.no_grad()
def exact_sequence():
    """Generate one digit at a time from the model's own output, as at inference."""
    tokens = torch.full((len(X_test), 1), START)
    memory = model.encode(X_test)
    for _ in range(L):
        nxt = model.decode(memory, tokens)[0][:, -1].argmax(-1, keepdim=True)
        tokens = torch.cat([tokens, nxt], dim=1)
    return (tokens[:, 1:] == Y_test).all(1).float().mean().item()


print(f"{sum(p.numel() for p in model.parameters()):,} parameters")

534,348 parameters

Eight epochs, reporting after each how often all three digits come out right.

history = []
for epoch in range(8):
    perm = torch.randperm(len(X), generator=generator)
    for i in range(0, len(perm) - 128, 128):
        b = perm[i:i + 128]
        logits = model(X[b], decoder_input[b])
        loss = nn.functional.cross_entropy(logits.reshape(-1, 12), decoder_target[b].reshape(-1))
        optimiser.zero_grad()
        loss.backward()
        optimiser.step()
    history.append(exact_sequence())
    print(f"epoch {epoch + 1}: all three digits correct on {history[-1]:.1%} of held-out images")

epoch 1: all three digits correct on 39.4% of held-out images
epoch 2: all three digits correct on 78.5% of held-out images
epoch 3: all three digits correct on 80.8% of held-out images
epoch 4: all three digits correct on 85.6% of held-out images
epoch 5: all three digits correct on 87.3% of held-out images
epoch 6: all three digits correct on 88.4% of held-out images
epoch 7: all three digits correct on 88.1% of held-out images
epoch 8: all three digits correct on 88.7% of held-out images

What It Learned on Its Own

Accuracy isn’t the interesting output here. What’s interesting is that nobody told this model the digits run left to right.

The encoder produces 48 patch tokens in a 4×12 grid. The decoder generates three outputs. Which patches does each output position attend to?

with torch.no_grad():
    memory = model.encode(X_test[:1])
    tokens = torch.cat([torch.full((1, 1), START), Y_test[:1]], dim=1)
    _, weights = model.decode(memory, tokens, want_weights=True)

attention = weights[0].reshape(-1, *GRID)   # one map per output position

fig, axes = plt.subplots(3, 1, figsize=(7.4, 4.2))
for position, ax in enumerate(axes):
    ax.imshow(X_test[0], cmap="gray_r")
    ax.imshow(attention[position], cmap="inferno", alpha=0.45,
              extent=(0, 84, 28, 0), interpolation="bilinear")
    ax.set_ylabel(f"digit {position + 1}", color=MUTED, fontsize=9, rotation=0,
                  ha="right", va="center", labelpad=12)
    ax.set_xticks([]); ax.set_yticks([])
    for s in ax.spines.values():
        s.set_visible(False)
fig.tight_layout()

Figure 3: Cross-attention from each output position onto the image patches. The alignment is learned, not specified.

Putting numbers on that: the share of each output position’s attention landing in each third of the image.

print("share of each output's attention falling on each third of the image\n")
print(f"{'':10} {'left':>8} {'middle':>8} {'right':>8}")
for position in range(L):
    by_column = attention[position].sum(0)
    thirds = [by_column[i:i + 4].sum() / by_column.sum() for i in (0, 4, 8)]
    print(f"digit {position + 1:<4} " + "".join(f"{t:>8.2f}" for t in thirds))

share of each output's attention falling on each third of the image

               left   middle    right
digit 1        0.95    0.05    0.00
digit 2        0.01    0.93    0.05
digit 3        0.00    0.03    0.97

Each output position learns to look at its own third of the image. The loss never mentions position. It only ever says the first digit is a 7. The alignment between output order and image geometry is something the model works out because it’s the only way to get the answer right.

The same mechanism makes translation work. There the alignment is between words in two languages rather than positions in an image. It’s why cross-attention replaced the fixed-size context vector that encoder-decoders used before it.

Conclusion

A decoder is not a bigger classifier. It is a different contract.

• The model commits to one token
• It then conditions on its own commitment
• Errors compound down the sequence
• Teacher forcing hides all of that during training

Which is why the number that matters has to be generated the slow way.

One consequence of generating a token at a time is that you can watch it happen. The service streams each digit the moment the decoder emits it, so the page shows the number being read rather than appearing at once. That is not a UI flourish, it is what the architecture does, made visible.

This project is at the stage where the plumbing is verified rather than the accuracy tuned. The numbers above show the mechanism works. They are not a result.

Where this goes

The week ends by pointing the same encoder-decoder at something that is not an image at all.

Swap the patch projection for a convolution over a waveform and the decoder now writes music notation. Nothing in the middle changed. That is the same point the Encoders slide made, and it is what week 5 is built on.

The full project is on GitHub.