CLIP Contrastive Loss

The problem

CLIP (Contrastive Language-Image Pre-Training) learns a shared space for images and text from pairs of an image and its caption. An image encoder and a text encoder each produce an embedding vector; training pulls the embedding of an image close to that of its own caption, and pushes it away from the captions of the other images of the batch. You will write the loss, following the pseudocode in Figure 3 of Radford et al. (2021), and the zero-shot classifier built on it.

The loss

Take a batch of $N$ matching pairs: image $i$ goes with caption $i$.

1. Normalize each image embedding and each text embedding to unit length (divide by the L2 norm: the usual length of a vector, over the last axis). Now the dot product of two embeddings is their cosine similarity: 1 when they point the same way, 0 when unrelated.
2. Logits: img @ txt.T / temperature, an $N \times N$ table. Entry $(i, j)$ compares image $i$ with caption $j$; the correct pairs are on the diagonal. The temperature, a small number, makes the differences bigger.
3. Two cross-entropies (section 2) with the diagonal as the target: loss_i over each row (image $i$ must pick its caption among $N$), and loss_t over each column (caption $j$ must pick its image).
4. Return (loss_i + loss_t) / 2.

Example. Two images with embeddings $(1, 0)$ and $(0, 1)$, two captions with embeddings $(1, 0.2)$ and $(0.1, 1)$.

• Cosine similarities: $\begin{pmatrix} 0.98 & 0.10 \\ 0.20 & 0.995 \end{pmatrix}$: the diagonal is the largest in each row and column.
• With temperature = 0.5 the logits are these times 2; the row cross-entropies are small, and the loss is $0.17$. With temperature = 1: $0.36$. With the learned CLIP temperature $0.07$: almost 0 ($8 \times 10^{-6}$).
• If the embeddings carried no information, so that every caption looked the same for every image, the loss would be $\log 2 = 0.69$: two equally likely choices.

The column cross-entropy is the row cross-entropy of logits.T.

Zero-shot classification

Without training for the task, a CLIP model can classify an image: write one caption per class ("a photo of a dog", "a photo of a cat"), embed the captions, and for each image pick the class whose caption embedding has the highest cosine similarity.

Example. With the embeddings above, treat the two captions as two classes. Image $(1, 0)$ is most similar to caption 0 (0.98 against 0.10): class 0. Image $(0, 1)$: class 1.

What to implement

• clip_loss(img_emb, txt_emb, temperature): the symmetric loss, a float.
• zero_shot_predict(img_emb, class_txt_emb): for each image, the index of the class text embedding with the highest cosine similarity.

Requirements

• Stable log-softmax (the temperature 0.07 makes logits large).
• No loops.