CLIP Contrastive Loss (InfoNCE) - Representation Learning & Retrieval

Symmetric InfoNCE (CLIP Loss)

What you'll build

The contrastive loss CLIP is trained with: L2-normalize image and text embeddings, build the temperature-scaled cosine similarity matrix, and average the cross-entropy over both directions with the matching pairs on the diagonal.

Background

CLIP (Radford et al., 2021) learns one embedding space for images and text from 400M (image, caption) pairs, with no class labels. In a batch of N pairs, image i and caption i are a positive pair; the other N - 1 captions are negatives for that image, and vice versa. The same loss trains text embedding models for retrieval (with (query, passage) pairs), SimCLR (two augmentations of one image) and many more. More negatives make the task harder and the representation better, which is why these models are trained with very large batches.

Formula

img_n  = img / max(||img||, 1e-12)                    # per row, (N, D)
txt_n  = txt / max(||txt||, 1e-12)                    # per row, (N, D)
logits = img_n @ txt_n.T / tau                        # (N, N); row = image, column = text

CE(Z)    = mean_i [ logsumexp(Z[i, :]) - Z[i, i] ]    # correct class of row i is column i
loss_i2t = CE(logits)                                 # each image picks its caption among N
loss_t2i = CE(logits.T)                               # each caption picks its image among N
loss     = (loss_i2t + loss_t2i) / 2

A lower temperature tau sharpens the softmax; CLIP learns it (initialised to 0.07).

Inputs

• img_emb: float array (N, D); row i pairs with txt_emb row i
• txt_emb: float array (N, D)
• temperature: float > 0, default 0.07

Expected Return Format

A dictionary:

• 'logits': (N, N) array of scaled cosine similarities (row = image, column = text)
• 'loss_i2t': image-to-text cross-entropy (Python float)
• 'loss_t2i': text-to-image cross-entropy (Python float)
• 'loss': their mean (Python float)

Edge cases and conventions

• Normalize the rows (each embedding), not the columns; rescaling an embedding must not change the loss.
• Use a stable log-sum-exp: with tau = 0.07 logits reach about 14, and smaller temperatures go higher.
• The matrix is not symmetric in general, so loss_i2t != loss_t2i.

Constraints

NumPy only, vectorised (no loop over the batch).

Follow-up

• With N pairs, the loss of a random model is about log N. Why? What does that say about comparing losses across batch sizes?
• SigLIP replaces the softmax with independent sigmoid losses per pair. What does that buy for distributed training?

Input: clip_loss(np.eye(2), np.eye(2), temperature=1.0)
Output: {'logits': [[1.0, 0.0], [0.0, 1.0]], 'loss_i2t': 0.3133, 'loss_t2i': 0.3133, 'loss': 0.3133}

Input: clip_loss(np.array([[1.0, 0.0], [0.6, 0.8], [0.0, 1.0]]), np.array([[0.8, 0.6], [0.0, 1.0], [1.0, 1.0]]), temperature=0.5)
Output: {'loss_i2t': 1.1121, 'loss_t2i': 1.1433, 'loss': 1.1277, ...}

import numpy as np

def clip_loss(img_emb, txt_emb, temperature=0.07):
    """
    Symmetric InfoNCE loss used by CLIP.

    Args:
        img_emb: (N, D) image embeddings (row i pairs with txt_emb row i)
        txt_emb: (N, D) text embeddings
        temperature: softmax temperature tau

    Returns:
        dict with 'logits' (N, N) = cos / tau (row = image), and the Python
        floats 'loss_i2t', 'loss_t2i' and 'loss' (their mean)
    """
    # Your code here
    pass