Symmetric CLIP Loss with a Learnable Temperature
Problem Statement
The full CLIP loss is the average of two cross-entropy terms — image-to-text and text-to-image — over the scaled similarity matrix. Implement it from a matrix of cosine similarities and a learnable temperature.
Background
Given cosine similarities sim (N x N, row = image, col = text), CLIP scales them by a temperature and treats each row (and each column) as an N-way classification whose correct label is the diagonal:
logits=τsim,L=21(CErows+CEcols)
with targets 0..N-1. Cross-entropy of a row is -log softmax(row)[label]. In practice 1/tau is a learned "logit scale" clamped to at most 100; here you are given tau directly. Use a numerically stable log-softmax (subtract the max).
Your Task
Implement:
def clip_loss(sim, tau):
Return the scalar loss rounded to 4 decimals.
Input Format
- sim: N x N nested list of cosine similarities.
- tau (float): temperature, tau > 0.
Output Format
- A float rounded to 4 decimals.
Sample
print(clip_loss([[1.0, 0.0], [0.0, 1.0]], 1.0))
Output:
0.3133
Example:
print(clip_loss([[1.0, 0.0], [0.0, 1.0]], 1.0))
0.3133
Each row's logits are [1,0] or [0,1]; the CE of the correct class is -log(softmax)[label] = 0.3133, symmetric for columns, so the average is 0.3133.
Constraints:
1 <= N <= 512,tau > 0.- Average the row-wise and column-wise cross-entropies.
- Diagonal entries are the correct labels.
- Use a stable log-softmax; round the final scalar to 4 decimals.
1. Background Knowledge
CLIP (Contrastive Language-Image Pre-training) learns joint image-text embeddings by maximizing similarity between matched pairs while minimizing similarity between mismatched pairs. The core loss is a symmetric cross-entropy: for an N×N cosine-similarity matrix, each row is treated as an N-way classification problem where the correct label is the diagonal entry (the matched text), and each column is treated similarly (the matched image). The final loss is the average of these two directional losses.
The temperature τ (or its reciprocal, the logit scale 1/τ) controls the sharpness of the softmax distribution. A small τ makes the softmax more peaked, increasing the penalty for incorrect pairs. In practice, 1/τ is a learnable parameter often clamped to at most 100 to prevent numerical overflow. Here, τ is given directly, so you divide the similarity matrix by τ to obtain the logits.
Numerical stability is critical: computing log(softmax(x)) naively as log(exp(xi)/∑jexp(xj)) can overflow for large xi. The standard fix is the log-sum-exp trick: subtract the row maximum before exponentiating, since softmax is shift-invariant.
2. Algorithm Approach
The problem reduces to computing two row-wise and column-wise cross-entropy losses on a scaled similarity matrix:
- Scale the similarity matrix by 1/τ to get logits.
- For each row, compute the cross-entropy with the diagonal index as the target.
- For each column, compute the cross-entropy with the diagonal index as the target.
- Average the two sums and divide by N (or equivalently, average the per-row and per-column losses).
The key sub-routine is a stable log-softmax followed by indexing the target position.
3. Step-by-Step Strategy
Continue the full explanation
You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.
Editor locked
The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.