Lesson 55: Vision-Language Models (VLMs)¶
Every classifier so far has predicted one of a small, fixed set of classes baked in at training time (Lesson 36). 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 50) 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 47's ViT pattern: patchify, embed, a small Transformer encoder) and a text encoder (a small Transformer encoder over the caption's tokens, Lesson 46) both map their very different inputs into the same fixed-size vector space — matching real CLIP's own Transformer-based design. (Real CLIP has also shipped CNN/ResNet image-encoder variants, which work too; Transformer is used here to stay consistent with this part of the course.) Training pulls a batch's matching (image, caption) pairs together and pushes every mismatched pair apart — precisely Lesson 50's nt_xent_loss, just applied across two different modalities instead of two augmented views of one image.
def patchify(img, patch_size):
B, C, H, W = img.shape
P = patch_size
patches = img.unfold(2, P, P).unfold(3, P, P)
patches = patches.contiguous().view(B, C, -1, P, P).permute(0, 2, 1, 3, 4)
return patches.reshape(B, -1, C * P * P)
def positional_encoding(T, D):
pos = torch.arange(T).unsqueeze(1).float()
i = torch.arange(D).unsqueeze(0).float()
angle_rates = 1.0 / (10000 ** (2 * (i // 2) / D))
angles = pos * angle_rates
pe = torch.zeros(T, D)
pe[:, 0::2] = torch.sin(angles[:, 0::2])
pe[:, 1::2] = torch.cos(angles[:, 1::2])
return pe
class ImageEncoder(nn.Module):
"""A small ViT (Lesson 47's pattern): patchify, linearly
embed, prepend a [CLS] token, run a Transformer encoder, read the [CLS] token's output."""
def __init__(self, img_size=16, patch_size=4, out_dim=16, embed_dim=24, n_heads=4, n_layers=2):
super().__init__()
self.patch_size = patch_size
n_patches = (img_size // patch_size) ** 2
patch_dim = patch_size * patch_size
self.embed = nn.Linear(patch_dim, embed_dim)
self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim) * 0.02)
self.register_buffer('pos_enc', positional_encoding(n_patches + 1, embed_dim))
layer = nn.TransformerEncoderLayer(embed_dim, n_heads, dim_feedforward=embed_dim * 2,
batch_first=True, dropout=0.0)
self.encoder = nn.TransformerEncoder(layer, num_layers=n_layers)
self.proj = nn.Linear(embed_dim, out_dim)
def forward(self, x):
p = patchify(x, self.patch_size)
tok = self.embed(p)
cls = self.cls_token.expand(x.shape[0], -1, -1)
tok = torch.cat([cls, tok], dim=1) + self.pos_enc
out = self.encoder(tok)
return self.proj(out[:, 0])
class TextEncoder(nn.Module):
"""A small bidirectional Transformer encoder over the caption's word tokens, pooled by
averaging the final token representations -- unlike a bag-of-words mean, these token
representations have already mixed in positional and attention information, so the
pooled result is sensitive to word order."""
def __init__(self, vocab_size, out_dim=16, embed_dim=16, n_heads=4, n_layers=1, max_len=6):
super().__init__()
self.embed = nn.Embedding(vocab_size, embed_dim)
self.register_buffer('pos_enc', positional_encoding(max_len, embed_dim))
layer = nn.TransformerEncoderLayer(embed_dim, n_heads, dim_feedforward=embed_dim * 2,
batch_first=True, dropout=0.0)
self.encoder = nn.TransformerEncoder(layer, num_layers=n_layers)
self.proj = nn.Linear(embed_dim, out_dim)
def forward(self, token_ids):
x = self.embed(token_ids) + self.pos_enc[:token_ids.shape[1]]
x = self.encoder(x)
return self.proj(x.mean(dim=1))
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 45'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 56).
Exercises¶
- Add a third shape category (e.g. "square", following Lesson 36'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? - Unlike a bag-of-words mean, this lesson's
TextEncoderhas both positional encoding and self-attention, so'a large plus'and'plus large a'are not guaranteed to produce the same embedding. Verify this directly: feed the encoder the same caption's words in their normal order and in a scrambled order, and compare the two embeddings. Then check whether this order-sensitivity actually changes zero-shot classification on this dataset — since every training caption always follows the exact same template ('a photo of a <size> <shape>'), does the model ever have a reason to use word order, even though it technically could?