Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Lineghuan Meng, Chuang Wang

TL;DR

This work addresses the training dynamics of a one-layer nonlinear contrastive learning model in the high-dimensional limit, deriving a McKean–Vlasov PDE that describes the evolution of the weight distribution. When regularization is limited to $L_2$, this PDE collapses to a low-dimensional system of ODEs that track model performance via the cosine-similarity metrics. The key findings are that the learnability of a feature at initialization depends solely on the second moment $\langle c^2 \rangle$, while higher moments influence feature selection probability through attraction regions; moreover, correlated additive noise can reduce gradient variance and improve performance, whereas independent noise generally harms learning. The results reveal a rich set of training-dynamics phenomena in a minimal model, offering insights toward understanding large-priority models and guiding design of data-augmentation schemes. Overall, the work provides a rigorous mean-field framework to connect data statistics, noise, and optimization dynamics in self-supervised contrastive learning.keywordsmean-field limitMcKean-Vlasov PDEone-layer nonlinear contrastive learningcosine similarity dynamicsdata augmentation noise correlationgradient variance reductionphase portrait feature selection

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.

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

TL;DR

This work addresses the training dynamics of a one-layer nonlinear contrastive learning model in the high-dimensional limit, deriving a McKean–Vlasov PDE that describes the evolution of the weight distribution. When regularization is limited to $L_2$, this PDE collapses to a low-dimensional system of ODEs that track model performance via the cosine-similarity metrics. The key findings are that the learnability of a feature at initialization depends solely on the second moment $\langle c^2 \rangle$, while higher moments influence feature selection probability through attraction regions; moreover, correlated additive noise can reduce gradient variance and improve performance, whereas independent noise generally harms learning. The results reveal a rich set of training-dynamics phenomena in a minimal model, offering insights toward understanding large-priority models and guiding design of data-augmentation schemes. Overall, the work provides a rigorous mean-field framework to connect data statistics, noise, and optimization dynamics in self-supervised contrastive learning.keywordsmean-field limitMcKean-Vlasov PDEone-layer nonlinear contrastive learningcosine similarity dynamicsdata augmentation noise correlationgradient variance reductionphase portrait feature selection

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.
Paper Structure (28 sections, 3 theorems, 95 equations, 11 figures)

This paper contains 28 sections, 3 theorems, 95 equations, 11 figures.

Table of Contents

  1. Introduction
  2. Models for Theoretical Analysis
  3. Data Model
  4. One-Layer Nonlinear Model for Contrastive Learning
  5. High-dimensional Training Dynamics
  6. Simplifications and Assumptions
  7. Performance metric and Order Parameters
  8. PDE for joint distribution of feature and network weights
  9. ODEs for the Cosine Similarity
  10. Analysis of the training dynamics
  11. Feature learnability depends on the second moment only
  12. Higher moments affect feature selection probability as well
  13. Correlated additive noise reduces gradient variance
  14. Conclusion
  15. Derivation of PDE in Theorem 1
  16. ...and 13 more sections

Key Result

Theorem 1

Under the assumptions in Section sec:ass, as $N \to \infty$, the sequence of random probability measures $\{\mu^N_t\}$ converges weakly to a deterministic measure $\mu_t$, of which the limiting density function $P_t(w,\boldsymbol{u})$ is the unique solution to the following nonlinear PDE where the drift and diffusion coefficients are respectively, and with the symbols $\boldsymbol{g}$, $f$, and

Figures (11)

  • Figure 1: Structure of the 1-layer nonlinear contrastive learning model. For a batch of data, two additive noises are injected to each sample as two argumented views, which are fed into the two branches of a one-layer nonlinear network respectively.
  • Figure 2: (a) Dynamics of $Q_t$ with the theoretical prediction by \ref{['eq:ODE1d']} (solid curves) and numerical simulations (errorbars) (b) The r.h.s of the ODE \ref{['eq:ODE1d']} with $h(Q)=\frac{\text{d}Q}{\text{d}t}/\tau$. It visualizes the locations of fixed points (where cross the horizontal axis) and their stability (stable if derivative of $g(Q)$ is negative).
  • Figure 3: (a) Phase portrait (b) Squared cosine similarities. The variance of hidden variables of the two features are $m_2^{(1)}=1.2$, $m_2^{(2)}=1.1$. Green- and red-shaddowed areas are attraction regions to the first and second features respectively.
  • Figure 4: stationary cosine similarity v.s. strength additive noises at the recovery state
  • Figure 5: Structure of the 1-layer nonlinear contrastive model analyzed in this work.
  • ...and 6 more figures

Theorems & Definitions (13)

  • Theorem 1
  • Theorem 2
  • Proposition 3
  • Proof : Proof of Theorem 2
  • Proof : Proof of Proposition 3
  • Claim 1
  • Proof : Proof of Claim 1
  • Claim 2
  • Proof : Proof of Claim 2
  • Claim 3
  • ...and 3 more