Image Captioning with Frozen Models

Week 4 of the residency. Take a vision model and a language model, freeze both, and train only the small adapter that connects them.

Week 4’s title slide reads multimodal transfer learning, and its subtitle is the part that matters: merge models for “novel” tasks.

Not train a model for a new task. Merge existing models, so that between them they do something neither was trained for.

Week 3 built a vision transformer from scratch and trained it end to end.

That works when the task is small and cheap to supervise. Here is the new one.

There’s no dataset big enough to teach vision and language from scratch on a residency budget. There’s also no need. Models that understand images already exist. So do models that write English.

The task is to connect them.

Where to Join Them

This is the decision the week turns on, and both options are already familiar.

xAtt, cross-attention. A dedicated attention layer inside the decoder that looks at the encoder’s output at every generation step. This is exactly what the multi-digit reader did last week.

sAtt, self-attention. Project the image into the language model’s embedding space and put it in the sequence as if it were a token. Ordinary self-attention then does the rest, and the language model is not modified at all.

The second is cheaper in every way that matters here.

• No new attention weights to initialise
• No architectural change to a model you did not train
• The only new parameters sit between the two models’ widths

LLaVA and PaLiGemma both take this path. So did I.

What CLIP Already Knows

CLIP is trained by pushing an image and its caption to the same place in a shared space, over 400 million pairs. That means its image vectors are already organised by what the picture is about, in a space that was built alongside text.

So the adapter is not being asked to teach a language model to see. It is being asked to translate between two coordinate systems that were each built separately and happen to describe overlapping things.

The week’s actual assignment was:

• Learn how to use the hidden state from ViT or CLIP
• Code the decoder from scratch
• Create synthetic datasets and save them on Hugging Face
• Train and experiment with datasets and alignment
• Use a pretrained Qwen model as the decoder

Building One

The real service uses CLIP and Qwen3-0.6B with a 1.58M-parameter adapter between them. Neither fits in a page build.

So the rest of this post builds the same arrangement out of two tiny models. Both are trained here, on separate tasks, so they genuinely have never met.

The vision side learns to classify digits. The language side is a character-level model trained on one sentence pattern, the digit is <word>, and nothing else. It knows the template and the ten words, and it has never seen an image.

First, the data and the sentence the language model will learn.

import sys

sys.path.insert(0, "..")
import numpy as np
import torch
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()

WORDS = ["zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine"]
characters = sorted(set("".join(WORDS) + " the digit is "))
VOCAB = ["<bos>", "<eos>"] + characters
INDEX = {c: i for i, c in enumerate(VOCAB)}


def encode(sentence):
    return [INDEX["<bos>"]] + [INDEX[c] for c in sentence] + [INDEX["<eos>"]]


captions = [encode(f"the digit is {w}") for w in WORDS]
LENGTH = max(len(c) for c in captions)
CAPTIONS = torch.tensor([c + [INDEX["<eos>"]] * (LENGTH - len(c)) for c in captions])

print(f"{len(VOCAB)} vocabulary entries, captions {LENGTH} tokens long")
print(f"example: {''.join(VOCAB[i] for i in CAPTIONS[7] if VOCAB[i] not in ('<bos>', '<eos>'))}")

19 vocabulary entries, captions 20 tokens long
example: the digit is seven

Now the vision model. It’s a two-layer network trained to classify digits. The number to watch is its accuracy, because it’s the ceiling for everything that follows.

DIM, PREFIX = 64, 8


class VisionModel(nn.Module):
    """Trained to classify digits. Its hidden layer is the image embedding."""

    def __init__(self):
        super().__init__()
        self.trunk = nn.Sequential(nn.Flatten(), nn.Linear(784, 128), nn.ReLU(),
                                   nn.Linear(128, DIM), nn.ReLU())
        self.classifier = nn.Linear(DIM, 10)

    def forward(self, images):
        return self.trunk(images)


torch.manual_seed(0)
vision = VisionModel()
optimiser = torch.optim.Adam(vision.parameters(), lr=1e-3)
generator = torch.Generator().manual_seed(0)

for _ in range(6):
    perm = torch.randperm(len(x_train), generator=generator)
    for i in range(0, len(perm) - 128, 128):
        b = perm[i:i + 128]
        loss = nn.functional.cross_entropy(vision.classifier(vision(x_train[b])), y_train[b])
        optimiser.zero_grad()
        loss.backward()
        optimiser.step()

with torch.no_grad():
    VISION_CEILING = (vision.classifier(vision(x_test)).argmax(1) == y_test).float().mean().item()
print(f"the vision model classifies digits at {VISION_CEILING:.4f}")

the vision model classifies digits at 0.9307

Then the language model, trained only on those ten sentences. One detail in it matters later: it’s pretrained with the image slots already present in the sequence, filled by a learned placeholder. Without that, inserting an image at inference would shift every text token one position along, into positional embeddings the frozen model was never trained with.

class LanguageModel(nn.Module):
    """A causal character model. It always has PREFIX slots in front of the text,
    filled by a learned null during pretraining and by the adapter afterwards."""

    def __init__(self):
        super().__init__()
        self.embed = nn.Embedding(len(VOCAB), DIM)
        self.positions = nn.Parameter(torch.randn(1, LENGTH + PREFIX, DIM) * 0.02)
        self.null = nn.Parameter(torch.zeros(1, PREFIX, DIM))
        layer = nn.TransformerEncoderLayer(DIM, 4, 4 * DIM, batch_first=True,
                                           norm_first=True, dropout=0.0)
        self.stack = nn.TransformerEncoder(layer, 2)
        self.out = nn.Linear(DIM, len(VOCAB))

    def forward(self, tokens, prefix=None):
        front = self.null.expand(len(tokens), -1, -1) if prefix is None else prefix
        h = torch.cat([front, self.embed(tokens)], dim=1) 
        h = h + self.positions[:, :h.shape[1]]
        mask = nn.Transformer.generate_square_subsequent_mask(h.shape[1])
        return self.out(self.stack(h, mask=mask, is_causal=True))


torch.manual_seed(1)
language = LanguageModel()
optimiser = torch.optim.Adam(language.parameters(), lr=3e-3)

for _ in range(1500):
    batch = CAPTIONS[torch.randint(0, 10, (64,))]
    logits = language(batch[:, :-1])[:, PREFIX:]
    loss = nn.functional.cross_entropy(logits.reshape(-1, len(VOCAB)), batch[:, 1:].reshape(-1))
    optimiser.zero_grad()
    loss.backward()
    optimiser.step()

print(f"the language model writes the sentence pattern, final loss {loss.item():.4f}")

the language model writes the sentence pattern, final loss 0.1206

Both are now frozen, and a small adapter goes between them. Only the adapter gets an optimiser. The two models either side never receive a gradient again.

for p in vision.parameters():
    p.requires_grad_(False)
for p in language.parameters():
    p.requires_grad_(False)

adapter = nn.Sequential(nn.Linear(DIM, DIM), nn.GELU(), nn.Linear(DIM, PREFIX * DIM))

trainable = sum(p.numel() for p in adapter.parameters())
frozen = sum(p.numel() for p in vision.parameters()) + sum(p.numel() for p in language.parameters())
print(f"trainable (adapter): {trainable:,}")
print(f"frozen (both models): {frozen:,}")

trainable (adapter): 37,440
frozen (both models): 214,109

Drawn out, the trainable part is small:

from _arch import diagram

diagram(adapter, input_shape=(1, DIM))

Figure 1: The entire trainable part. It takes the vision model’s 64-number summary and produces eight 64-number vectors for the language model to read as if they were text.

Before training the adapter, it’s worth seeing what the frozen pair does on its own. The language model will write a perfectly good sentence, because that’s all it knows how to do. It just has no idea which digit.

@torch.no_grad()
def caption(n=1000):
    """Generate a caption one character at a time and check the word at the end."""
    prefix = adapter(vision(x_test[:n])).view(n, PREFIX, DIM)
    tokens = torch.full((n, 1), INDEX["<bos>"])
    for _ in range(LENGTH):
        nxt = language(tokens, prefix)[:, -1].argmax(-1, keepdim=True)
        tokens = torch.cat([tokens, nxt], dim=1)

    text = ["".join(VOCAB[i] for i in row[1:LENGTH] if VOCAB[i] not in ("<bos>", "<eos>"))
            for row in tokens]
    correct = sum(t.strip().endswith(WORDS[y]) for t, y in zip(text, y_test[:n].tolist()))
    return correct / n, text


before, examples = caption()
print(f"before training the adapter: {before:.4f}")
print(f"  it writes fluent English already: {examples[0]!r}")
print("  and guesses the same word every time")

before training the adapter: 0.0000
  it writes fluent English already: 'the digit is fo'
  and guesses the same word every time

Now train only the adapter, on image and caption pairs. The label at the image position is masked out. The model is never asked to predict the image, only to predict caption words given it.

optimiser = torch.optim.Adam(adapter.parameters(), lr=1e-3)
history = []

for epoch in range(24):
    perm = torch.randperm(len(x_train), generator=generator)
    for i in range(0, len(perm) - 128, 128):
        b = perm[i:i + 128]
        prefix = adapter(vision(x_train[b])).view(len(b), PREFIX, DIM)
        target = CAPTIONS[y_train[b]]
        logits = language(target[:, :-1], prefix)[:, PREFIX:]
        loss = nn.functional.cross_entropy(logits.reshape(-1, len(VOCAB)),
                                           target[:, 1:].reshape(-1))
        optimiser.zero_grad()
        loss.backward()
        optimiser.step()
    history.append(caption()[0])

accuracy, examples = caption()
print(f"after training only the adapter: {accuracy:.4f}\n")
for text, truth in zip(examples[:5], y_test[:5].tolist()):
    word = WORDS[truth]
    article = "an" if word[0] in "aeiou" else "a"
    print(f"  {text!r}   (actually {article} {word})")

after training only the adapter: 0.9220

  'the digit is six'   (actually a six)
  'the digit is three'   (actually a three)
  'the digit is three'   (actually a three)
  'the digit is eight'   (actually an eight)
  'the digit is nine'   (actually a nine)

Plotting that against the vision model’s own accuracy shows where it stops, and why.

from _style import COLOURS, MUTED, figure, style_axes

epochs = range(1, len(history) + 1)
fig, ax = figure(height=4.0)
ax.axhline(VISION_CEILING, color=COLOURS[2], linewidth=1.2, linestyle="--")
ax.axhline(0.1, color=MUTED, linewidth=1.2, linestyle="--")
ax.plot(epochs, history, color=COLOURS[0], linewidth=2)
ax.annotate("what the frozen vision model knows", xy=(1, VISION_CEILING), xytext=(2, 6),
            textcoords="offset points", fontsize=9, color=COLOURS[2])
ax.annotate("guessing one of ten words", xy=(1, 0.1), xytext=(2, 6),
            textcoords="offset points", fontsize=9, color=MUTED)
ax.set_ylim(0, 1.0)
style_axes(ax, "Epoch (adapter only)", "Captions with the right digit")
fig.tight_layout()

Figure 2: The adapter climbing toward the frozen vision model’s own accuracy, which it cannot pass.

Conclusion

Four things came out of building it.

The ceiling is the frozen encoder. The adapter climbs to the vision model’s own accuracy and stops. It never sees the image, only what the vision model kept. Anything the encoder discarded is gone. So picking the encoder matters more than designing the adapter.

Fluency comes free, grounding does not. Before any adapter training the captions are already correct English in the right format. The whole run is spent on which word. That is the only thing the language model could not already do.

One prefix token was not enough here, which surprised me, because the real project gets by with a single one. My first version projected the image to one embedding and it barely beat chance. CLIP’s pooled vector is a summary built from 400 million pairs being read by a model with real depth; mine was a small classifier’s hidden layer being read by two transformer layers. How many tokens you need depends on how good the summary is and how much model is reading it.

And the positions have to line up, which is the reason for those placeholder slots in the language model above.

The full project, with CLIP and Qwen3 in place of these two toys, is on GitHub.