CLIP InfoNCE Loss

Background

The CLIP training objective is a symmetric InfoNCE loss over a batch of paired image and text embeddings. For each image, the matched text is the positive; the other $N-1$ texts in the batch are negatives. Symmetric: do it once with image as query, once with text as query, average. Foundation of every modern multimodal retrieval system.

Problem statement

Implement clip_infonce(image_embs, text_embs, temperature=0.07). Both inputs are L2-normalised, shape $(N, d)$. Compute the scaled similarity matrix and the symmetric cross-entropy:

$$\text{logits} = \frac{Z_{\text{img}} Z_{\text{text}}^\top}{\tau},\qquad \mathcal L = \tfrac{1}{2}\bigl(\mathcal L_{\text{img}\to\text{text}} + \mathcal L_{\text{text}\to\text{img}}\bigr)$$

where each direction is a softmax cross-entropy with labels = arange(N) (the diagonal).

Input

• image_embs — np.ndarray shape $(N, d)$, L2-normalised.
• text_embs — np.ndarray same shape, L2-normalised.
• temperature — float > 0, softmax temperature (default 0.07, CLIP's published value).

Output

Returns a Python float >= 0.

Examples

Example 1 — identical paired embeddings → near-zero loss

Input:  image_embs == text_embs (random unit-norm)
Output: small positive value (entropy of the softmax)

Example 2 — shape mismatch raises

Input:  image (4, 8) vs text (5, 8)
Output: ValueError

Constraints

• Image and text must have matching shape (same $N$ and same $d$).
• Use numerically stable log-softmax (subtract max before exp).
• Return a plain Python float.

Notes

• Temperature matters. Too high → soft assignment, weak gradient. Too low → degenerate hard assignment. CLIP learns $\tau$ jointly; many derivatives fix it at $0.07$ or $0.1$.
• Batch size. InfoNCE quality scales with the number of negatives. Production CLIP uses batches of 32K+ via gradient accumulation or distributed all-gather.