Training Dynamics of Nonlinear Contrastive Learning Model in the High Dimensional Limit

TL;DR

This paper provides a high-dimensional theoretical analysis of nonlinear contrastive learning, revealing how weight distributions evolve and how various factors influence feature learning and model performance.

Contribution

It introduces a PDE-based framework for analyzing training dynamics in contrastive learning models, simplifying to ODEs under regularization, and uncovers key factors affecting feature learnability and stability.

Findings

Second moment of hidden variables influences feature learnability.

Higher moments affect feature selection probability via attraction regions.

Correlated noise in data augmentation improves performance by reducing gradient variance.

Abstract

This letter presents a high-dimensional analysis of the training dynamics for a single-layer nonlinear contrastive learning model. The empirical distribution of the model weights converges to a deterministic measure governed by a McKean-Vlasov nonlinear partial differential equation (PDE). Under L2 regularization, this PDE reduces to a closed set of low-dimensional ordinary differential equations (ODEs), reflecting the evolution of the model performance during the training process. We analyze the fixed point locations and their stability of the ODEs unveiling several interesting findings. First, only the hidden variable's second moment affects feature learnability at the state with uninformative initialization. Second, higher moments influence the probability of feature selection by controlling the attraction region, rather than affecting local stability. Finally, independent noises added in the data argumentation degrade performance but negatively correlated noise can reduces the variance of gradient estimation yielding better performance. Despite of the simplicity of the analyzed model, it exhibits a rich phenomena of training dynamics, paving a way to understand more complex mechanism behind practical large models.