CLIP Contrastive (InfoNCE) Loss

mediumFREEContrastive Objectives: CLIP and SigLIPVision-Language ModelsVLM 1: Vision-Language Models

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

Example

Input

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

Reasoning

Inputs: img holds two image vectors and txt two text vectors in $\mathbb{R}^2$ ; row $i$ of each pairs up. logit_scale is $s = 10$ .
Normalise: the image rows already have length 1. The text rows have lengths $\sqrt{0.82} \approx 0.9055$ and $\sqrt{1.04} \approx 1.0198$ , giving $\hat T_1 = (0.9939, 0.1104)$ and $\hat T_2 = (0.1961, 0.9806)$ .
Logits: entry $(i, j)$ is $s\,\hat I_i \cdot \hat T_j$ , so the $2 \times 2$ matrix is $\begin{pmatrix} 9.9388 & 1.9612 \\ 1.1043 & 9.8058 \end{pmatrix}$ . Positives sit on the diagonal.
Image to text (rows): row 1 costs $\log\big[1 + e^{1.9612 - 9.9388}\big] \approx 0.000343$ , row 2 costs $\log\big[1 + e^{1.1043 - 9.8058}\big] \approx 0.000166$ ; mean $\approx 0.000255$ .
Text to image (columns): column 1 compares $9.9388$ with $1.1043$ ( $\approx 0.000146$ ), column 2 compares $9.8058$ with $1.9612$ ( $\approx 0.000392$ ); mean $\approx 0.000269$ .
Output: $\tfrac12(0.000255 + 0.000269) \approx 0.000262$ , which rounds to 0.0003. Notice the directions differ because the off-diagonals are not symmetric, and the loss is tiny because each positive beats its rival by about 8 logits.

Constraints

1 <= n <= 64, 1 <= d <= 64; no row is the zero vector
Embeddings arrive unnormalised - you must L2-normalise the rows yourself
The loss is the AVERAGE of the image-to-text and text-to-image cross-entropies
Use a numerically stable log-sum-exp; logit_scale can be 100
Round the returned loss to 4 decimals

Learn this in

Latent Diffusion & Stable DiffusionDiffusion Study Plan Multimodal LLMsLLM Study Plan Advanced NLP TopicsNLP Study Plan Introduction to Vision-Language ModelsVLM Study Plan

More from VLM 1: Vision-Language Models

Visual Token Pruning by Attention Masshard CHAIR-i and CHAIR-s Hallucination Ratesmedium Merge a LoRA Update into the Base Weightmedium QLoRA Memory Footprint Estimatemedium POPE Precision, Recall, and F1medium

Concept primer

1. Background Knowledge

Contrastive Learning is a self-supervised learning paradigm where models learn by distinguishing between "positive" pairs (matching items) and "negative" pairs (non-matching items). In the context of Vision-Language Models (VLMs) like CLIP, the goal is to align image and text representations in a shared embedding space. The InfoNCE (Information Noise-Contrastive Estimation) loss is the standard objective used to achieve this alignment. It maximizes the similarity between positive pairs while minimizing the similarity between all other pairs in the batch.

The core mechanism involves computing a similarity matrix, often using cosine similarity. Since embeddings are L2-normalized, their dot product equals their cosine similarity. This matrix is scaled by a temperature parameter, logit_scale (or $\tau^{-1}$ ), which controls the sharpness of the probability distribution. A higher scale makes the model more confident in its predictions, effectively sharpening the contrast between positives and negatives.

Crucially, CLIP employs a symmetric loss. It computes the cross-entropy loss in two directions: image-to-text (treating images as queries and texts as keys) and text-to-image (treating texts as queries and images as keys). These two losses are averaged. This symmetry ensures that the model performs well for both image retrieval given text and text retrieval given images. Numerical stability is paramount because the logit_scale can be large (e.g., 100), causing exp() to overflow. The log-sum-exp trick is used to compute the softmax denominator safely by subtracting the maximum value in each row/column before exponentiation.

2. Algorithm Approach

The problem requires implementing the InfoNCE loss with symmetric contrastive objectives. The approach involves matrix operations to compute similarities, followed by stable softmax calculations.

Normalization: First, normalize the input image and text embeddings to unit length. This converts dot products into cosine similarities.
Similarity Matrix: Compute the dot product matrix $S = I \times T^\top$ . Scale this matrix by logit_scale.
Symmetric Loss Calculation:

Image-to-Text: For each row (image), compute the cross-entropy against the diagonal element (its paired text). Use the log-sum-exp trick for numerical stability.
Text-to-Image: For each column (text), compute the cross-entropy against the diagonal element (its paired image). This is equivalent to computing the loss on the transpose of the similarity matrix.

Averaging: Average the two directional losses to get the final scalar loss.

3. Step-by-Step Strategy

L2 Normalization:

Iterate through each embedding vector in image_emb and text_emb.
Compute the L2 norm: $\|v\|_2 = \sqrt{\sum v_i^2}$ .
Divide each element by its norm. Handle zero-norm vectors if necessary (though rare in practice).

Compute Logits Matrix:

Calculate the dot product between every image embedding and every text embedding. This results in an $n \times n$ matrix where logits[i][j] is the similarity between image $i$ and text $j$ .
Multiply the entire matrix by logit_scale.

Image-to-Text Loss:

For each row $i$ in the logits matrix:
Find the maximum value $m_i = \max_j(\text{logits}[i][j])$ .
Compute the log-sum-exp: $\text{lse}_i = m_i + \log(\sum_j e^{\text{logits}[i][j] - m_i})$ .
The loss for this row is: $\text{lse}_i - \text{logits}[i][i]$ (since the target is the diagonal).
Average these losses over all $n$ rows.

Text-to-Image Loss:

Repeat the process for columns. Alternatively, transpose the logits matrix and apply the same row-wise logic.
For each column $j$ :
Find max $m_j = \max_i(\text{logits}[i][j])$ .
Compute $\text{lse}_j = m_j + \log(\sum_i e^{\text{logits}[i][j] - m_j})$ .
Loss for this column is: $\text{lse}_j - \text{logits}[j][j]$ .
Average these losses over all $n$ columns.

Final Loss:

Compute $\mathcal{L} = 0.5 \times (\text{loss}_{\text{img}\to\text{text}} + \text{loss}_{\text{text}\to\text{img}})$ .
Round the result to 4 decimal places.

4. Common Pitfalls

Numerical Overflow: Directly computing exp(logits) with large logit_scale will result in inf. Always use the log-sum-exp trick by subtracting the row/column maximum before exponentiating.
Incorrect Normalization: Forgetting to L2-normalize the embeddings before computing dot products will result in magnitudes affecting similarity, not just direction.
Symmetry Error: Implementing only one direction (e.g., image-to-text) instead of averaging both directions. CLIP's loss is explicitly symmetric.
Indexing Errors: When computing the text-to-image loss, ensure you are comparing against the correct diagonal elements. The positive pair for text $j$ is image $j$ , so the target logit is logits[j][j].
Division by Zero: If an embedding is all zeros, its norm is zero. While unlikely in trained models, handle this edge case to avoid division by zero during normalization.

5. Time & Space Complexity

Time Complexity: $O(n^2 d)$ , where $n$ is the batch size and $d$ is the embedding dimension. Normalization takes $O(nd)$ . Computing the $n \times n$ similarity matrix takes $O(n^2 d)$ . The loss computation iterates over the $n \times n$ matrix, taking $O(n^2)$ . Thus, the dominant term is $O(n^2 d)$ .
Space Complexity: $O(n^2)$ to store the similarity matrix. The normalized embeddings take $O(nd)$ , but the $n \times n$ matrix is typically larger for large batches.

solution.py

Python 3.13

Test results