🤖 AI Summary
This work investigates why contrastive learning yields effective representations under simple image augmentations. By analytically characterizing the optimal solution of the contrastive loss, the study theoretically establishes—for the first time—that, under specific augmentations, the optimal first-layer filters are sinusoidal functions. This insight leads to a derived CNN architecture comprising sinusoidal filters, pointwise nonlinearities, global average pooling, and a partially whitening linear layer. The authors further introduce a water-filling algorithm based on the data’s power spectrum to compute the frequencies and weights of these filters. Experiments across multiple image datasets and augmentation strategies confirm that the first layer indeed learns sinusoidal filters and performs partial whitening, thereby revealing the underlying mechanism by which contrastive learning shapes its representations.
📝 Abstract
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 a 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