Implement CLIP Contrastive Loss
Company: Uber
Role: Machine Learning Engineer
Category: Machine Learning
Difficulty: medium
Interview Round: Technical Screen
Given a minibatch of paired image and text embeddings, implement the symmetric contrastive loss used in CLIP-style image-text representation learning.
You are given:
- `image_embeddings`: a tensor of shape `(batch_size, embedding_dim)`.
- `text_embeddings`: a tensor of shape `(batch_size, embedding_dim)`.
- The `i`-th image corresponds to the `i`-th text.
Compute a similarity matrix between every image embedding and every text embedding. Then compute:
- Image-to-text loss: cross entropy over each image row, where the correct class for row `i` is text `i`.
- Text-to-image loss: cross entropy over each text row, equivalently cross entropy on the transposed similarity matrix, where the correct class for row `i` is image `i`.
- Final loss: the average of the two losses.
Implement this loss function and explain any important details such as normalization, temperature scaling, and label construction.
Overview: This question evaluates understanding and implementation of contrastive representation learning concepts—specifically similarity matrices, symmetric image-text contrastive loss, normalization, temperature scaling, and label construction—testing competency in building losses for embedding alignment in the Machine Learning domain.