CLIP contrastive loss for image-text pairs

Implement the symmetric contrastive loss used to train CLIP-style models that align images and text. You get a minibatch of paired embeddings: the i-th image matches the i-th text, and every other pairing in the batch is a negative. Compute a similarity matrix between every image and every text, take the cross-entropy in both directions, and average the two. Then explain the details that make it work: normalization, temperature and how the labels are built.

Interface

Unless you agree on a different interface with the interviewer, assume:

• clip_loss(image_emb, text_emb, temperature: float) -> float — both inputs have shape (batch_size, dim).
  1. L2-normalize each row of both inputs.
  2. logits = image_emb @ text_emb.T / temperature, shape (batch_size, batch_size).
  3. labels = [0, 1, ..., batch_size - 1].
  4. Image-to-text loss: cross-entropy of each row of logits against labels. Text-to-image loss: the same on logits.T.
  5. Return the mean of the two losses.

Use NumPy or PyTorch.

Example

image_emb = [[1, 0], [0, 1]]
text_emb  = [[1, 0], [0, 1]]
clip_loss(image_emb, text_emb, temperature=1.0)  ->  0.3133   # log(1 + e^-1)
clip_loss(image_emb, text_emb, temperature=0.5)  ->  0.1269   # log(1 + e^-2)

text_emb  = [[0, 1], [1, 0]]                    # pairs swapped
clip_loss(image_emb, text_emb, temperature=1.0)  ->  1.3133   # log(1 + e)

Constraints

• batch_size from 2 to tens of thousands; dim in the hundreds.
• The computation must be numerically stable for small temperatures (large logits).

Clarifying questions

• Is the temperature fixed or a learned parameter (CLIP learns its log)?
• Are the embeddings already normalized?
• Can the batch contain duplicate texts, such as two captions that are the same sentence?
• Should the loss be averaged over the batch or summed?

What a strong answer covers

• Normalization makes the dot product a cosine similarity, so the scale is controlled only by the temperature.
• Labels are the diagonal: each row's correct class is its own index, which is why the batch supplies its own negatives.
• A stable cross-entropy using log-sum-exp (subtract the row maximum), or the framework's built-in cross_entropy.
• Why both directions: the matrix is not symmetric in what each row normalizes over.
• The temperature's effect: lower values sharpen the softmax and weight hard negatives; learned as exp(t) with a clamp to avoid instability.
• Cost: O(B² · d) time and O(B²) memory for the logits.

Follow-ups

• The batch does not fit on one GPU. How do you compute the loss across devices (gathering embeddings from all workers)?
• The batch contains false negatives, such as two images of the same dog with similar captions. How do you handle them?
• Replace the softmax with a pairwise sigmoid loss. What changes and why might it scale better?
• How would you check the implementation is correct?