AI for Everyone, Part 4 — Captioning Images With 1.37M Parameters, and Why Part 3 Bug Was Absent Here

• How captioning differs from VQA
• Data — giving images sentences
• Architecture — splitting encoder and decoder
• Why Part 3's bug was absent here
• Training log
• Results
• Summary
• 🧠 Comprehension quiz
• References

How captioning differs from VQA

Part 3's VQA took an image and a question and produced a short answer. This part has no question: the model looks at an image alone and writes a sentence.

[sneaker image]   ->  "a photo of a sneaker"
[trousers image]  ->  "trousers on white background"

The difference is output length. A VQA answer was six characters; a caption is forty. Longer output calls for a different structure. Part 3 concatenated image, question, and answer into one sequence; here we separate the encoder from the decoder.

Data — giving images sentences

Fashion-MNIST has no captions, only labels 0 through 9. So we built templates that turn a label into a sentence.

ARTICLES = ["t-shirt", "trousers", "a pullover", "a dress", "a coat",
            "a sandal", "a shirt", "a sneaker", "a bag", "an ankle boot"]
TEMPLATES = ["a photo of {}", "this looks like {}", "{} on white background"]

Each image gets one of three templates at random. This forces the model to learn the structure rather than memorise whole sentences: if it starts with "a photo of", an item name follows; if it starts with an item name, "on white background" follows.

The articles are deliberately irregular. t-shirt and trousers take none, and only an ankle boot takes an. We can check whether the model gets those right too.

Architecture — splitting encoder and decoder

class Captioner(nn.Module):
    def __init__(self):
        super().__init__()
        self.cnn = nn.Sequential(
            nn.Conv2d(1, 32, 3, 2, 1), nn.GELU(),
            nn.Conv2d(32, 64, 3, 2, 1), nn.GELU(),
            nn.Conv2d(64, 128, 3, 2, 1), nn.GELU(),
        )
        self.proj = nn.Linear(128, 160)
        self.emb = nn.Embedding(V, 160)
        self.pos = nn.Embedding(CAP, 160)
        layer = nn.TransformerDecoderLayer(160, 5, 640, batch_first=True, norm_first=True)
        self.dec = nn.TransformerDecoder(layer, 3)
        self.head = nn.Linear(160, V)
        self.register_buffer("m", torch.triu(torch.full((CAP, CAP), float("-inf")), 1))

    def forward(self, img, cap_in):
        mem = self.proj(self.cnn(img).flatten(2).transpose(1, 2))
        T = cap_in.shape[1]
        h = self.emb(cap_in) + self.pos(torch.arange(T, device=cap_in.device))
        return self.head(self.dec(h, mem, tgt_mask=self.m[:T, :T]))

The heart of it is the single line self.dec(h, mem, ...). nn.TransformerDecoder takes two inputs:

• h — the caption tokens generated so far
• mem — the 16 image vectors produced by the CNN (the memory)

Inside each decoder layer, two kinds of attention happen in order. First self-attention, where the caption looks at its own past — which is why tgt_mask is needed. Then cross-attention, where each caption position looks at the 16 image vectors.

That second one matters. When generating the word "sneaker", the model chooses for itself which region of the image to consult. Packing everything into one sequence, as in Part 3, makes image and text share the same attention; here the roles are separated and each is optimised on its own terms. The longer the caption, the more this separation pays off.

Why Part 3's bug was absent here

The culprit that dropped Part 3 to 7.5% was the function that builds target sequences. Let us place the two side by side.

Part 3 (buggy)

def tok(s, L):
    ids = [enc[c] for c in s][:L]
    return ids + [0] * (L - len(ids))

a_full = torch.tensor([tok(s, AL - 1) + [1] for s in ans])   # pad, then EOS

Part 4 (correct)

def tok(s):
    ids = [enc[c] for c in s][:CAP - 1] + [1]   # EOS first
    return ids + [0] * (CAP - len(ids))          # then pad

Feeding "even" through both:

One difference in ordering. Because Part 4 attaches EOS first and pads afterwards, EOS always lands where the answer ends. The model could learn to stop, and so never hit the bug.

The same person wrote both on the same day and only one was wrong. This class of mistake is hard to catch in review and invisible to metrics — as Part 3 showed, the loss actually looked better. There is no substitute for looking at the output.

Training log

[   12.6s] params {"params": 1367262, "cap_len": 40}
[   13.3s] step      0 loss=3.7440
[   14.2s] step    100 loss=0.1660
[  608.4s] SUMMARY {"steps": 79504, "final_loss": 0.0468, "caption_acc": 0.91}

Within 100 steps — barely a second — loss fell from 3.74 to 0.166. There are only 30 caption combinations (10 items × 3 templates), so the structure was grasped quickly.

Results

We scored 300 test images by whether the correct item name appears in the generated caption. The result is 91%.

[an ankle boot] -> an ankle boot on white background
[a pullover]    -> a photo of a pullover
[trousers]      -> trousers on white background
[trousers]      -> this looks like trousers
[a shirt]       -> a photo of a shirt
[a coat]        -> a photo of a coat

Several things are worth noting.

Templates vary. The same trousers produced "trousers on white background" once and "this looks like trousers" another time. It learned all three structures rather than memorising one.

Articles are correct. The an in an ankle boot, the a in a pullover, and the absence of an article before trousers are all right. Getting this from character-level generation means it learned the relationship between an item's first letter and its article.

The wrong 9% is mostly shirt-like items. In Fashion-MNIST, shirt, t-shirt, coat, and pullover are hard to tell apart at 28×28 in greyscale. People confuse them too, so this error is closer to a limit of the data than a defect in the model.

Summary

Encoder-decoder separation earns its keep as output grows longer. Cross-attention lets each word being generated consult only the part of the image it needs, and image representations stop competing with language representations.

The contrast with Part 3 is this article's other lesson. The same mistake was within reach and was avoided only because EOS happened to be attached first. To stop relying on luck, print the output of your target-building function once. Pulling one batch and calling print takes ten seconds.

The next part turns to output that is an image rather than text — colourising greyscale photographs. There we will see the choice of loss function change the character of the result.

🧠 Comprehension quiz

References

• Fashion-MNIST dataset
• Show and Tell — A Neural Image Caption Generator
• Show, Attend and Tell — Captioning with attention
• PyTorch nn.TransformerDecoder docs