Training Dynamics of Nonlinear Contrastive Learning Model in the High Dimensional Limit
This work provides theoretical insights into contrastive learning dynamics, which is incremental but foundational for understanding complex large models in machine learning.
The authors analyzed the training dynamics of a single-layer nonlinear contrastive learning model in high dimensions, showing that the weight distribution converges to a deterministic measure governed by a PDE, which reduces to ODEs under L2 regularization, revealing that only the second moment of hidden variables affects feature learnability and that negatively correlated noise in data augmentation can improve performance by reducing gradient variance.
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.