Vision-Language Pretraining

Vision-language pretraining trains one model on two objectives at once: ranking and generation. Ranking means that given a caption, the model finds the right image among many. Generation means that given an image, it writes a caption. Train on one skill alone and you get half a system. CLIP nailed ranking but cannot caption. GPT-4V can caption but uses a separate retrieval head for ranking. Multi-objective pretraining gets both in one pass.

InfoNCE handles the ranking half. For a batch of N pairs, the model treats the N matching pairs as positives and the N^2 - N mismatched pairs as negatives, then runs a cross-entropy loss on the resulting (N, N) similarity matrix. The LM loss handles the generation half: standard next-token prediction conditioned on the image. Both losses are differentiable and can share the vision encoder.

InfoNCE in one paragraph

InfoNCE is a cross-entropy loss over a similarity matrix: each image must pick out its own caption from the batch, and each caption its own image. Stack the N image embeddings as rows and the N text embeddings as rows, then L2-normalize both. Compute the N x N matrix S = exp(log_tau) * I T^T, which is the usual I T^T / tau with 1/tau = exp(log_tau) learned in log space. The diagonal entries are the matching pairs, and the off-diagonal entries are negatives. Apply cross-entropy with the target argmax running down the diagonal, so row i should have its highest entry in column i. Do the same symmetrically along columns and average the two. That is the CLIP loss in eight lines.

Temperature matters

The temperature tau controls how peaked the softmax is. Set it too small (say tau = 0.01) and the gradient comes only from the very hardest negative, which makes training noisy. Set it too large and the softmax flattens and the gradient vanishes. CLIP learns the temperature as a parameter, and the demo here does the same. In the code the parameter is log_tau (initialized to ln(1/0.07), so the scale starts near 14.3), and the tau printed during training is that multiplier, so a larger printed value means a sharper softmax.

Language modeling loss

The decoder consumes image memory tokens via cross-attention and predicts the next text token at every position. The loss is standard cross-entropy against the next-position target, with padding positions masked out.

Combining the losses

total = contrastive + lm_weight * lm, where lm_weight is a scalar (often 1.0). Both losses send gradient into the vision encoder. The projector and the text-side encoder get gradient only from the contrastive loss, and the decoder only from the LM loss, because the decoder cross-attends to the encoder's tokens directly rather than to the projected embedding. CoCa, BLIP, and SigLIP-style models all use this multi-task recipe, with various weightings.

Why 50 steps is enough for a demo

The mock corpus is a synthetic 200-pair set of noisy images, each with a brightness offset set by its index, and deterministic caption ids computed from the index. After 50 Adam steps with batch size 16, both losses drop visibly, even if the absolute values stay above what a real-data model would reach. The demo has two jobs: confirm the gradient plumbing works end to end, and confirm that adding the LM loss does not destabilize the contrastive objective.

code/main.py implements:

• MultimodalModel, combining a small ViT encoder, the MLP projector, a tiny text-side encoder (mean-pool over embedded ids), and the cross-attention decoder from lesson 61.
• info_nce_loss(image_emb, text_emb, log_tau), the bidirectional CLIP-style contrastive loss, returning (loss, sim).
• lm_loss(logits, target_ids, padding_id), masked next-token cross-entropy.
• make_mock_corpus(seed, n_pairs, vocab_size, max_len), returning 200 deterministic (image, caption_ids) pairs.
• A training loop running 50 steps with batch size 16, Adam optimizer, and a learned log-temperature parameter. Both losses are printed every 5 steps.

Run it:

python3 code/main.py

The run is seeded (seed=0), so every run prints the same numbers. On CPU the contrastive loss goes from 3.05 at step 0 to 2.36 at step 49, and the LM loss goes from 6.50 to 4.70. Both start above their uniform-guess baselines, ln(16) = 2.77 for picking one of 16 captions and ln(512) ≈ 6.24 for picking one of 512 tokens, because the random init produces small random similarities that the temperature scale of about 14 amplifies, so the first predictions are slightly worse than guessing. The contrastive curve is noisy (it ticks up at step 10 and again at step 30), so read the endpoints, not every row. The script prints ok: both losses decreased when both endpoints fall. Real models train for millions of steps, and the dynamics are the same.

The same loss recipe ships in these models:

• CLIP (2021). Image-text contrastive only, with a separate frozen-encoder caption probe.
• CoCa (2022). Image-text contrastive plus image-captioning LM loss in one model. This is the pattern you build here.
• BLIP (2022) and BLIP-2. Contrastive plus LM plus an image-text matching head: three losses combined.
• SigLIP (2023). Swaps InfoNCE for a sigmoid pair loss. The contrastive role is the same, the functional form differs.
• LLaVA family. Two-stage training. Stage one is alignment (cosine on a frozen LM), and stage two adds LM loss with an unfrozen LM. Lesson 60 maps to stage one; this lesson maps to stage two.

code/test_main.py has 7 tests:

• InfoNCE loss is symmetric across image/text rows
• InfoNCE loss is near 0 (below 1e-3) when the embeddings are orthonormal and the scale is large, so the similarity matrix is a sharp diagonal
• InfoNCE raises ValueError when the image and text batch sizes differ
• LM loss correctly masks padding positions
• model forward pass produces both losses without errors
• a 10-step training loop ends with a combined loss no higher than the first step (within 1e-3)
• the mock corpus has the right shapes and token ids stay inside the vocab

Run them:

python3 -m unittest code/test_main.py

1. Replace InfoNCE with a SigLIP-style sigmoid pair loss and compare convergence on the mock corpus.

2. Add a hard-negative mining step: every other batch, select the hardest off-diagonal pair from the previous batch and append it. Train and check whether the contrastive loss drops faster.

3. Add an image-text matching binary head on top of the joint embedding (true/false: do these match?) as a third loss, replicating BLIP's three-head setup.

4. Replace the mock corpus with caption-id sequences drawn from a Markov chain whose transition matrix is conditioned on image hash. There is now real learnable signal, so the captioning loss should drop further.

5. Train the same model with lm_weight = 0 and again with lm_weight = 1, then compare the contrastive loss. The LM loss should not regress the ranking objective.

• CLIP paper for the original contrastive recipe.
• CoCa paper for contrastive plus captioning in one model.
• SigLIP paper for the sigmoid pair-loss variant and why it scales better.