Title: Linear Separability in Contrastive Learning via Neural Training Dynamics
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: contrastive learning, neural network, linear separability, training dynamics, variational analysis, gradient flow, neural tangent kernel

Abstract: The SimCLR method for contrastive learning of invariant visual representations has become extensively used in supervised, semi-supervised, and unsupervised settings, due to its ability to uncover patterns and structures in image data that are not directly present in the pixel representations. However, this success is still not well understood; neither the loss function nor invariance alone explains it. In this paper, we present a mathematical analysis that clarifies how the geometry of the learned latent distribution arises from SimCLR. Despite the nonconvex SimCLR loss and the presence of many undesirable local minimizers, we show that the training dynamics driven by gradient flow tend toward favorable representations. In particular, early training induces clustering in feature space. Under a structural assumption on the neural network, our main theorem proves that the learned features become linearly separable with respect to the ground-truth labels. To support our theoretical insights, we present numerical results that align with our theoretical predictions.

TL;DR: We present a novel theoretical result showing that SimCLR training dynamics lead to clustering and linear separability, despite nonconvex loss and poor local minima.