Understanding Deep Contrastive Learning via Coordinate-wise Optimization

Yuandong Tian

Keywords: contrastive learning, self-supervised learning, representation learning, principal component analysis, loss design, landscape analysis, deep linear network

TL;DR: We propose a unified formulation for a broad family of contrastive losses, including InfoNCE, propose novel losses, and show contrastive learning with deep linear network can be equivalent to PCA.

Abstract: We show that Contrastive Learning (CL) under a broad family of loss functions (including InfoNCE) has a unified formulation of coordinate-wise optimization on the network parameter $\vtheta$ and pairwise importance $\alpha$, where the \emph{max player} $\vtheta$ learns representation for contrastiveness, and the \emph{min player} $\alpha$ puts more weights on pairs of distinct samples that share similar representations. The resulting formulation, called \boldmethod{}, unifies not only various existing contrastive losses, which differ by how sample-pair importance $\alpha$ is constructed, but also is able to extrapolate to give novel contrastive losses beyond popular ones, opening a new avenue of contrastive loss design. These novel losses yield comparable (or better) performance on CIFAR10, STL-10 and CIFAR-100 than classic InfoNCE. Furthermore, we also analyze the max player in detail: we prove that with fixed $\alpha$, max player is equivalent to Principal Component Analysis (PCA) for deep linear network, and almost all local minima are global and rank-1, recovering optimal PCA solutions. Finally, we extend our analysis on max player to 2-layer ReLU networks, showing that its fixed points can have higher ranks. Codes are available in https://github.com/facebookresearch/luckmatters/tree/main/ssl/real-dataset.