NT-Xent Contrastive Loss
Medium · Solved in PyTorch · Machine learning coding practice
Problem
Two independently augmented views of the same image form a positive pair. Implement SimCLR's normalized temperature-scaled cross-entropy (NT-Xent) loss, introduced in A Simple Framework for Contrastive Learning of Visual Representations (Chen et al., 2020); all views of other images serve as negatives.
Rows with the same index in z_a and z_b are paired. In the ordering Z=[z_a;z_b], the positive index is p(i)=(i+N)\bmod 2N. Let s_{ik} be cosine similarity, with each L2 norm clamped below by 10^{-12}.
Topics: contrastive-learning, loss-functions, pytorch
The full statement, worked examples, hints and the test suite are available once you sign in. You can then solve NT-Xent Contrastive Loss in the browser and run it against the tests.