CLIP Contrastive (InfoNCE) Loss

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.

1. Normalization: First, normalize the input image and text embeddings to unit length. This converts dot products into cosine similarities.
2. Similarity Matrix: Compute the dot product matrix $S = I \times T^\top$. Scale this matrix by logit_scale.
3. 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.

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

3. Step-by-Step Strategy

1. 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).

2. 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.

3. 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.

4. 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.

5. 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.