[Paper Review] Understanding Deep Contrastive Learning via Coordinate-wise Optimization

Yuandong Tian|arXiv (Cornell University)|Jan 29, 2022

Domain Adaptation and Few-Shot Learning4 citations

TL;DR

This paper introduces \alpha-CL, a unified coordinate-wise optimization framework for contrastive learning that treats representation learning as a min-max game between network weights (max player) and sample-pair importance weights (min player). It proves that with fixed \alpha, the max player's optimization is equivalent to PCA in deep linear networks, and extends this to 2-layer ReLU networks, showing non-rank-1 solutions emerge under orthogonal mixture data, with novel contrastive losses outperforming InfoNCE on CIFAR-10, STL-10, and CIFAR-100.

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 $\boldsymbolθ$ and pairwise importance $α$, where the \emph{max player} $\boldsymbolθ$ learns representation for contrastiveness, and the \emph{min player} $α$ puts more weights on pairs of distinct samples that share similar representations. The resulting formulation, called $α$-CL, unifies not only various existing contrastive losses, which differ by how sample-pair importance $α$ 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 $α$, 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.

Motivation & Objective

• To unify diverse contrastive loss functions under a single optimization framework.
• To analyze the training dynamics of the representation learner (max player) in deep linear and ReLU networks.
• To establish a theoretical connection between contrastive learning and Principal Component Analysis (PCA).
• To design novel contrastive losses via the \alpha-CL framework that outperform existing methods.
• To understand how nonlinearities like ReLU affect the rank and structure of learned representations in contrastive learning.

Proposed method

• Proposes a min-max coordinate-wise optimization formulation where the max player optimizes network parameters \theta to maximize contrastiveness, and the min player optimizes pairwise importance weights \alpha to emphasize similar representations of distinct samples.
• Reformulates contrastive learning as optimizing an energy function \mathcal{E}_\alpha(\bm{\theta}) minus a regularizer \mathcal{R}(\alpha), linking the loss to a game-theoretic perspective.
• Proves that for deep linear networks, the max player's objective under fixed \alpha is equivalent to PCA, with all local minima being global and rank-1.
• Introduces a sticky weight rule for ReLU networks that preserves non-negativity of first-layer weights, enabling analysis under orthogonal mixture data.
• Derives conditions under which the max player in ReLU networks can achieve higher-rank solutions than rank-1 PCA.
• Empirically validates the framework by constructing novel contrastive losses via different regularizers on \alpha, achieving state-of-the-art performance on CIFAR-10, STL-10, and CIFAR-100.

Experimental results

Research questions

• RQ1 Can contrastive learning be unified under a single coordinate-wise optimization framework involving network weights and sample-pair importance?
• RQ2 Does the representation learning objective in contrastive learning correspond to PCA in deep linear networks?
• RQ3 How do nonlinear activations like ReLU alter the solution space compared to linear networks in contrastive learning?
• RQ4 Can the min-player optimization over \alpha generate novel, effective contrastive losses beyond InfoNCE?
• RQ5 What is the rank structure of optimal solutions in ReLU networks under the \alpha-CL framework?

Key findings

• The max player in deep linear networks under fixed \alpha is equivalent to PCA, with all local minima being global and rank-1, achieving the optimal PCA solution.
• For 2-layer ReLU networks with orthogonal mixture data, the solution is not necessarily rank-1, and the network can learn higher-rank representations.
• The framework enables the design of novel contrastive losses via different regularizers on \alpha, which achieve comparable or better performance than InfoNCE on CIFAR-10, STL-10, and CIFAR-100.
• In the ReLU case, a sticky weight rule ensures that first-layer weights remain non-negative and do not collapse to zero, preserving structural properties.
• The maximal eigenvector of the weighted covariance matrix under \alpha-CL always contains at least one negative entry when multiple data modes are present, indicating non-trivial solution structure.
• Empirical results show that the proposed \alpha-CL framework with new regularizers yields state-of-the-art performance on standard benchmarks, validating the theoretical insights.