Freeze both ends and train the middle

Week 4 was image captioning the way LLaVA does it: take a vision model and a language model that have never met, freeze them both, and train only a small translator in between.

Week 4: image captioning. Look at a picture, write a sentence about it.

The obvious approach is to train one model on image–caption pairs. The approach that actually gets used is stranger and much cheaper, and it’s what LLaVA [1] and PaLiGemma popularised:

1. Take a vision encoder that already understands images. Freeze it.
2. Take a language model that already writes English. Freeze it.
3. Train a small adapter that turns the vision model’s output into something the language model will accept as input.

Nothing in either big model ever receives a gradient. The project version used CLIP on one side and Qwen3-0.6B on the other, and only ~1.58M adapter parameters were trainable.

The bet is that the two models already contain everything needed, and what’s missing is a translation between two representation spaces that were never trained together.

Why that bet is reasonable

A vision encoder’s output is a vector that means something — but it means something in the vision model’s coordinate system, which was shaped by a completely different training run.

The language model reads token embeddings, which live in their own coordinate system with their own geometry. It has no reason to interpret a CLIP vector as anything at all.

The adapter’s job is not to understand images or to write English. It’s to land the image’s vector somewhere in the language model’s embedding space that makes the language model say the right thing. That’s a much smaller job, which is why a small matrix can do it.

Building both ends from scratch, small

To show the recipe honestly I need two models that genuinely have never met. So I train both here, from scratch, on separate tasks, then freeze them.

The vision model learns to classify MNIST digits. Its hidden layer becomes the “image embedding”.

The language model is a character-level transformer trained on exactly one sentence pattern — the digit is <word> — and nothing else. It knows the template and it knows the ten words. It has never seen an image.

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

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 can classify digits at {VISION_CEILING:.4f}")

the vision model can classify digits at 0.9307

class LanguageModel(nn.Module):
    """A causal character LM. It always has PREFIX slots in front of the text —
    filled with a learned null during pretraining, and by the adapter later, so
    the positions never shift out from under the frozen weights."""

    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.1208

Now freeze both, and put a small adapter between them. Only the adapter gets an optimiser.

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:,}")
print(f"the adapter is {100 * trainable / (trainable + frozen):.0f}% of the whole thing")

trainable (adapter): 37,440
frozen (both models): 214,109
the adapter is 15% of the whole thing

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.

@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 still writes fluent English: {examples[0]!r}")
print("  it just has no idea which digit, so it guesses the same word every time")

before training the adapter: 0.0000
  it still writes fluent English: 'the digit is fo'
  it just has no idea which digit, so it guesses the same word every time

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()):
    print(f"  {text!r}   (actually a {WORDS[truth]})")

after training only the adapter: 0.9190

  '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 a eight)
  'the digit is nine'   (actually a nine)

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. It cannot go past it: an adapter translates information, it does not add any.

What the toy actually shows

The ceiling is the frozen encoder. The adapter gets close to the vision model’s own classification accuracy and stops there. It can’t do better, because it has no access to the image — only to what the vision model chose to keep. If the encoder threw something away, the adapter cannot recover it. That’s the real constraint of this recipe, and it’s why the choice of frozen encoder matters more than the adapter design.

Fluency is free and grounding is not. Before the adapter is trained, the captions are already perfect English in the right format. The entire training run is spent on which word, because that’s the only thing the language model couldn’t already do.

One prefix token wasn’t enough. My first version projected the image to a single embedding and it barely beat chance — the frozen language model couldn’t carry that one vector fifteen characters to where the word appears. Projecting to eight tokens fixed it. Real implementations use many image tokens, and now I know why rather than just that they do.

The positions have to line up. The language model is pretrained with the prefix slots already present, filled by a learned null. Without that, inserting the image shifts every text token one position to the right, into positional embeddings the frozen model was never trained with — and it quietly fails.

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

---

[1] Liu, Li, Wu, Lee. Visual Instruction Tuning. NeurIPS 2023.

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