Teaching a model to know when it has finished

Classifying an image needs one label. Reading a number off one needs a sequence, of a length the model has to decide for itself. That is what a decoder is for, and the alignment it learns is not something anybody specifies.

The second half of week 3. The classifier looks at one digit and picks one of ten labels. This one looks at three digits side by side and has to read the number.

That sounds like a small step and it is a different problem. A classifier produces a fixed-size answer. Reading produces an ordered sequence of unknown length — and something has to decide where it ends.

Everything below is standard encoder–decoder transformer, from the original attention paper [1]. The reason to build it rather than read about it is that three separate mechanisms have to be in place before a single digit comes out right, and they’re easy to nod along to and hard to get correct.

The three things 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 be allowed to look at position 3, or the model learns to read the answer it is being asked to predict, and then produces nothing useful at inference when the future genuinely isn’t there. The mask enforces that: each position attends only to 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, and the gap it leaves

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.

Which means the training loss systematically flatters the model, and the honest measure is generating the whole sequence the way you actually would — which is what exact_sequence below does.

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 and the tokens enter bottom left through their embedding, and 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.

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

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 77.2% of held-out images
epoch 3: all three digits correct on 83.0% of held-out images
epoch 4: all three digits correct on 85.8% of held-out images
epoch 5: all three digits correct on 85.4% of held-out images
epoch 6: all three digits correct on 86.3% of held-out images
epoch 7: all three digits correct on 87.1% of held-out images
epoch 8: all three digits correct on 87.3% of held-out images

The thing worth looking at

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.

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.96    0.04    0.00
digit 2        0.02    0.96    0.02
digit 3        0.00    0.02    0.98

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.

This is the same mechanism that makes translation work, where the alignment is between words in two languages rather than positions in an image, and it’s the reason cross-attention replaced the fixed-size context vector that encoder–decoders used before it.

What I’d take from this one

The decoder is not a bigger classifier. It’s a different contract: the model commits to one token, then conditions on its own commitment, and the errors compound. Teacher forcing hides that during training, which is why the number that matters is generated the slow way.

This project is at the stage where the plumbing is verified rather than the accuracy tuned, so I’d treat the numbers above as a demonstration that the mechanism works, not as a result.

The full project is on GitHub.

---

[1] Vaswani et al. Attention Is All You Need. NeurIPS 2017.

Built:      2026-09-18 02:19:27 UTC
Python:     3.13.15
matplotlib: 3.11.2
numpy:      2.5.3
torch:      2.14.0+cpu