Normalized Temperature-Scaled Cross-Entropy (NT-Xent)

A Simple Framework for Contrastive Learning of Visual Representations (SimCLR) — Chen, Kornblith, Norouzi, Hinton

The work in this paper is motivated by the InfoNCE Loss.

They go on to propose a NT-Xent Loss aka Normalized Temperature-Scaled Cross-Entropy:

$${\mathcal{L}}_{SimCLR}=-\log \frac{\exp (\text{sim}(z_i,z_j)/\tau )}{{\sum }_{i=1}^{2N}{1}_{k\ne i}\exp (\text{sim}(z_i,z_k)/\tau )}$$

where the similarity function is pairwise cosine similarity. As opposed to a dot product similarity, the cosine similarity (which has l2 normalization) leads to a worse contrastive task accuracy, but a better representation learning.

Tau here is the temperature scaling hyperparameter, which scales the logits (similarity scores) before applying the softmax function. A small tau sharpens the probability distribution, making the model highly sensitive to the relative distnaces between positive and negative pairs and concentrating learning effort on the “hard” negative examples.

The SimCLR framework consists of four major components:

1. Stochastic data augmentation that augments a given data example randomly, resulting in two correlated views ${\tilde{x}}_i,{\tilde{x}}_j$ of the same example. The data augmentations include random cropping, random color distortions, and random Gaussian blur.
2. A base encoder neural network $f(\cdot )$ (ResNet-50) then extracts representation vectors from augmented data samples.
3. A small neural network projection $g(\cdot )$ head (MLP with one hidden layer) maps representations to the space where contrastive loss is applied.
4. A contrastive loss function, where the two augmented images are a positive contrastive pair, and all others are negative.

As loss gets propagated backwards throughout the network and weights are updated, the base encoder learns representation vectors of the same input image more similar to each other, and more different from the representation vector of other images.

References