Lesson 49: Vision-Language Models¶
Every classifier so far has predicted one of a small, fixed set of classes baked in at training time (Lesson 37). CLIP ("Learning Transferable Visual Models From Natural Language Supervision", Radford et al., 2021★) breaks that constraint by training an image encoder and a text encoder jointly, with a contrastive loss (Lesson 47) that pulls each image's embedding toward its caption's embedding and pushes it away from every other caption in the batch. The payoff: classification becomes a matter of writing down the class names as text prompts and checking similarity — no classifier head, no retraining, works on categories the model never saw labeled examples of.
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
Images paired with captions¶
A tiny synthetic "language" (8 words) generates simple captions like "a photo of a large plus" or "a photo of a small circle" for each image. This is obviously not real natural language, but it exercises exactly the same mechanism CLIP uses on real image-caption pairs scraped from the web: every training example is an (image, text) pair, and nothing else.
VOCAB = ['a', 'photo', 'of', 'plus', 'circle', 'small', 'large', 'noisy']
WORD_TO_ID = {w: i for i, w in enumerate(VOCAB)}
def make_image(shape_type, cx, cy, size, rng):
img = np.zeros((size, size), dtype=np.float32)
if shape_type == 'plus':
r = size // 5
img[cy-1:cy+2, cx-r:cx+r+1] = 1.0
img[cy-r:cy+r+1, cx-1:cx+2] = 1.0
else:
r = size // 4
yy, xx = np.mgrid[0:size, 0:size]
img[((xx-cx)**2 + (yy-cy)**2) <= r**2] = 1.0
return np.clip(img + rng.normal(0, 0.1, img.shape), 0, 1).astype(np.float32)
def caption_for(shape_type, size_desc):
return [WORD_TO_ID[w] for w in ['a', 'photo', 'of', 'a', size_desc, shape_type]]
def make_dataset(rng_local, n, image_size=16):
imgs, captions = [], []
for _ in range(n):
shape_type = rng_local.choice(['plus', 'circle'])
size_desc = rng_local.choice(['small', 'large'])
cx, cy = image_size // 2, image_size // 2
img = make_image(shape_type, cx, cy, image_size, rng_local)
if size_desc == 'small':
small = img[4:12, 4:12]
img = np.zeros_like(img)
img[4:12, 4:12] = small
imgs.append(img)
captions.append(caption_for(shape_type, size_desc))
return np.array(imgs, dtype=np.float32), captions
rng = np.random.default_rng(7)
X_train, cap_train = make_dataset(rng, 400)
X_test, cap_test = make_dataset(rng, 150)
fig, axes = plt.subplots(1, 4, figsize=(9, 2.5))
for ax, im, cap in zip(axes, X_train[:4], cap_train[:4]):
ax.imshow(im, cmap='gray'); ax.axis('off')
ax.set_title(' '.join(VOCAB[w] for w in cap), fontsize=7)
plt.show()
Two encoders, one shared embedding space¶
An image encoder (Lesson 33's CNN pattern) and a text encoder (a word-embedding lookup, mean-pooled over the caption — the simplest possible text encoder; real CLIP uses a Transformer, Lesson 45) both map their very different inputs into the same fixed-size vector space. Training pulls a batch's matching (image, caption) pairs together and pushes every mismatched pair apart — precisely Lesson 47's nt_xent_loss, just applied across two different modalities instead of two augmented views of one image.
class ImageEncoder(nn.Module):
def __init__(self, out_dim=16):
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(1, 16, 5, padding=2), nn.ReLU(), nn.MaxPool2d(2),
nn.Conv2d(16, 32, 5, padding=2), nn.ReLU(), nn.AdaptiveMaxPool2d(1),
)
self.proj = nn.Linear(32, out_dim)
def forward(self, x):
return self.proj(self.conv(x).flatten(1))
class TextEncoder(nn.Module):
def __init__(self, vocab_size, out_dim=16, embed_dim=8):
super().__init__()
self.embed = nn.Embedding(vocab_size, embed_dim)
self.proj = nn.Linear(embed_dim, out_dim)
def forward(self, token_ids):
return self.proj(self.embed(token_ids).mean(dim=1)) # bag-of-words text encoder
def clip_loss(img_emb, txt_emb, temperature=0.1):
img_emb = F.normalize(img_emb, dim=1)
txt_emb = F.normalize(txt_emb, dim=1)
logits = img_emb @ txt_emb.T / temperature # (batch, batch) similarity matrix
targets = torch.arange(img_emb.shape[0]) # the diagonal is the true pairing
loss_i2t = F.cross_entropy(logits, targets) # "which caption matches this image?"
loss_t2i = F.cross_entropy(logits.T, targets) # "which image matches this caption?"
return (loss_i2t + loss_t2i) / 2
def train_clip(seed, epochs=300, lr=0.01, batch_size=64):
torch.manual_seed(seed)
img_enc = ImageEncoder()
txt_enc = TextEncoder(len(VOCAB))
opt = torch.optim.Adam(list(img_enc.parameters()) + list(txt_enc.parameters()), lr=lr)
caps_tensor = torch.tensor(cap_train)
n = len(X_train)
for epoch in range(epochs):
idx = np.random.default_rng(epoch).permutation(n)[:batch_size]
imgs = torch.tensor(X_train[idx]).unsqueeze(1)
caps = caps_tensor[idx]
loss = clip_loss(img_enc(imgs), txt_enc(caps))
opt.zero_grad()
loss.backward()
opt.step()
return img_enc, txt_enc, loss.item()
img_enc, txt_enc, final_loss = train_clip(seed=0)
print(f'final CLIP loss: {final_loss:.3f}')
Zero-shot classification¶
No classifier head was ever trained. To classify a test image, embed a text prompt for each candidate class ("a photo of a large plus", "a photo of a large circle") and pick whichever prompt's embedding is most similar to the image's embedding — the same nearest-neighbor-in-embedding-space idea as Lesson 44's attention retrieval demo, just across modalities.
prompts = {'plus': caption_for('plus', 'large'), 'circle': caption_for('circle', 'large')}
plus_id, circle_id = WORD_TO_ID['plus'], WORD_TO_ID['circle']
true_labels = [c[-1] for c in cap_test] # last word id encodes the true shape
X_test_t = torch.tensor(X_test).unsqueeze(1)
def zero_shot_eval(img_encoder, txt_encoder):
with torch.no_grad():
prompt_embs = {k: F.normalize(txt_encoder(torch.tensor([v])), dim=1) for k, v in prompts.items()}
img_embs = F.normalize(img_encoder(X_test_t), dim=1)
sims_plus = (img_embs @ prompt_embs['plus'].T).squeeze(-1)
sims_circle = (img_embs @ prompt_embs['circle'].T).squeeze(-1)
preds = torch.where(sims_plus > sims_circle, plus_id, circle_id)
return sum(p.item() == t for p, t in zip(preds, true_labels)) / len(true_labels)
acc_trained = zero_shot_eval(img_enc, txt_enc)
torch.manual_seed(0) # same init recipe, but never trained -- the baseline
untrained_img_enc = ImageEncoder()
untrained_txt_enc = TextEncoder(len(VOCAB))
acc_untrained = zero_shot_eval(untrained_img_enc, untrained_txt_enc)
print(f'zero-shot accuracy, trained encoders: {acc_trained:.1%}')
print(f'zero-shot accuracy, untrained encoders: {acc_untrained:.1%} (baseline)')
Trained encoders solve zero-shot classification perfectly; untrained ones are at chance — the contrastive image-text objective, not the specific classifier, is doing all the work. This is the actual mechanism behind real CLIP's headline result: it can classify into any set of categories describable in words, including ones that never appeared as a labeled class during training, simply by changing the text prompts at inference time. It's also the standard way to build a vision-language model (VLM)'s visual front end — a frozen CLIP-style image encoder feeding into a language model is the basis for systems that can answer questions about images, and the same paired image-text contrastive objective, applied to region proposals instead of whole images, is one route to open-vocabulary detection and segmentation (Lesson 52).
Exercise¶
- Add a third shape category (e.g. "square", following Lesson 37's pattern) to
make_dataset, retrain, and add a'square'prompt to zero-shot evaluation. Does 3-way zero-shot classification still work as well as the 2-way case? - Try a prompt using a word combination never seen together in training — e.g. if training captions only ever paired
'noisy'with neither shape, constructcaption_forvariants that use'noisy'and check whether the model's similarity ranking still makes sense. What does this test about whether the model learned compositional word meaning versus caption memorization? - This lesson's text encoder mean-pools word embeddings, ignoring word order entirely (
'a large plus'and'plus large a'produce the identical embedding). Would that make a difference for these particular captions? Sketch a change (hint: Lesson 45) that would make the text encoder order-sensitive, and describe a caption pair where order-sensitivity would actually matter.