Understanding the Role of Nonlinearity in Training Dynamics of Contrastive Learning

TL;DR

This paper investigates how nonlinearity influences the training dynamics of contrastive learning, revealing that nonlinearity introduces multiple local optima and enhances the model's ability to learn diverse data patterns, unlike linear models.

Contribution

It provides the first theoretical analysis of nonlinearity effects in contrastive learning, showing how it enables learning multiple patterns and the importance of nonlinearity in 2-layer networks.

Findings

Nonlinearity causes multiple local optima in 1-layer networks.

Linear activation cannot learn diverse patterns in 2-layer networks.

Global modulation prioritizes patterns discriminative at the global level.

Abstract

While the empirical success of self-supervised learning (SSL) heavily relies on the usage of deep nonlinear models, existing theoretical works on SSL understanding still focus on linear ones. In this paper, we study the role of nonlinearity in the training dynamics of contrastive learning (CL) on one and two-layer nonlinear networks with homogeneous activation $h(x) = h'(x)x$. We have two major theoretical discoveries. First, the presence of nonlinearity can lead to many local optima even in 1-layer setting, each corresponding to certain patterns from the data distribution, while with linear activation, only one major pattern can be learned. This suggests that models with lots of parameters can be regarded as a \emph{brute-force} way to find these local optima induced by nonlinearity. Second, in the 2-layer case, linear activation is proven not capable of learning specialized weights into diverse patterns, demonstrating the importance of nonlinearity. In addition, for 2-layer setting, we also discover \emph{global modulation}: those local patterns discriminative from the perspective of global-level patterns are prioritized to learn, further characterizing the learning process. Simulation verifies our theoretical findings.