CLIP Contrastive (InfoNCE) Loss

Problem Statement

Given a batch of image embeddings and their paired text embeddings, compute CLIP's symmetric InfoNCE loss.

Background

CLIP trains two encoders so that the i-th image matches the i-th caption and nothing else in the batch. Both embeddings are first L2-normalised, so their dot product is a cosine similarity in [-1, 1], then scaled by a learned logit_scale (the inverse temperature, exp(t), roughly 100 at convergence):

$\text{logits} = s \cdot \hat{I}\hat{T}^\top, \qquad \hat{I}_i = \frac{I_i}{\lVert I_i \rVert_2}$

The diagonal holds the n positive pairs; every off-diagonal entry is a negative. The loss is cross-entropy against the labels 0, 1, ..., n-1 computed in both directions and averaged:

$\mathcal{L} = \tfrac{1}{2}\Big( \underbrace{\tfrac{1}{n}\sum_i -\log \tfrac{e^{s\,\hat I_i\cdot\hat T_i}}{\sum_j e^{s\,\hat I_i\cdot\hat T_j}}}_{\text{image}\to\text{text}} + \underbrace{\tfrac{1}{n}\sum_j -\log \tfrac{e^{s\,\hat I_j\cdot\hat T_j}}{\sum_i e^{s\,\hat I_i\cdot\hat T_j}}}_{\text{text}\to\text{image}} \Big)$

The two directions differ: the first normalises over each row of the logit matrix, the second over each column. Averaging them is what makes CLIP usable for both image-to-text and text-to-image retrieval.

Because logit_scale is large, exp overflows on raw logits. Use the log-sum-exp trick: subtract each row's max before exponentiating.

Your Task

Implement:

def clip_loss(image_emb, text_emb, logit_scale):

image_emb and text_emb are n x d nested lists, row i of each being a positive pair. Return the scalar loss rounded to 4 decimals.

Input Format

• image_emb, text_emb - lists of n lists of d floats, NOT pre-normalised
• logit_scale - a positive float

Output Format

A single float rounded to 4 decimals.

Sample

img = [[1.0, 0.0], [0.0, 1.0]]
txt = [[0.9, 0.1], [0.2, 1.0]]
print(clip_loss(img, txt, 10.0))

Output:

0.0003