Contrastive Loss (InfoNCE / SimCLR-style)
Representation Learning, Advanced Theory & Miscellaneous DS practice problem on Onlearn.
Difficulty: medium.
Topics: Contrastive Loss (InfoNCE / SimCLR-style), Temperature Scaling, Positive Pair Sampling, Negative Sampling, Cosine Similarity, Log-Sum-Exp Trick, Representation Learning, Optimization Theory, Information Theory, Statistical Learning, Deep Learning Architectures, Metric Learning, Stochastic Gradient Descent, Entropy Estimation, Self-Supervised Learning, Embedding Space Geometry.
Implement the NT Xent (Normalized Temperature scaled Cross Entropy) contrastive loss function, commonly known as the SimCLR contrastive loss or InfoNCE loss. This loss is the backbone of self supervised contrastive learning frameworks. Given a batch of 2N embedding vectors where each consecutive pair (2i, 2i+1) represents two augmented views of the same input sample, the goal is to push representations of the same sample closer together while pushing representations of different samples apart. Your function receives: embeddings: a numpy array of shape (2N, d) where N is the number of original samples and d is the embedding dimension. Embeddings at indices (0,1), (2,3), (4,5), ... form positive pairs. temperature: a positive float scaling parameter (tau). The function should: 1. Compute pairwise cosine similarities between all L2 normalized embeddings. 2. For each of the 2N anchors, treat its paired view as the positive and all other 2N 2 samples (excluding itself) as negatives. 3. Compute the contrastive loss for each anchor and return the mean loss over all 2N anchors. Use the log sum exp trick for numerical stability when computing the log of the sum of exponentials.