Katkat
Back

Implement CLIP Contrastive Loss

UberMachine Learning Engineer · Technical Screen
MediumMachine Learning EngineerTechnical ScreenMachine Learning

Last updated

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.

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.