A Theory of Contrastive Learning with Natural Images

Why does contrastive learning with simple images and augmentations yield useful representations for downstream tasks? We address this question by analytically computing the optimal representation in terms of the contrastive loss for a range of basic augmentations and any image dataset with stationary statistics. We show that for certain augmentations the optimum can be attained by a CNN whose first layer filters are sinusoids, followed by a pointwise nonlinearity, global average pooling, and a final linear layer that performs partial whitening. We also show that the optimal weights in such CNNs for more complicated augmentations are still sinusoids. The frequencies of the sinusoids and their weights can be computed using a simple "waterfilling" algorithm given the dataset's expected power spectrum. Experiments with different image datasets and augmentations show that such CNNs trained with SGD empirically learn sinusoids in their first layer and to perform partial whitening. Contrastive learning (CL) is a remarkably successful method for learning useful image representations without labeled data. While numerous variants have been suggested, almost all of them follow the recipe suggested by (Chen et al., 2020) . For each training image, an augmentation is applied, and the goal of learning is to find a representation where two augmentations of the same image are close, while augmentations of different images are far away. Although conceptually simple, when applied to large scale datasets, these approaches have paved the way for image representations that can then be used to solve a large number of computer vision tasks without requiring additional representation learning (e.g.