Theoretical Analysis of Contrastive Learning under Imbalanced Data: From Training Dynamics to a Pruning Solution

TL;DR

This paper provides a theoretical analysis of contrastive learning with Transformer encoders on imbalanced data, revealing neuron dynamics and proposing pruning as a solution to improve representation quality.

Contribution

It introduces a novel theoretical framework for contrastive learning under data imbalance, detailing neuron training stages and proposing pruning to mitigate imbalance effects.

Findings

Neuron weights evolve through three training stages.

Minority features reduce representational capacity.

Pruning restores performance and improves feature separation.

Abstract

Contrastive learning has emerged as a powerful framework for learning generalizable representations, yet its theoretical understanding remains limited, particularly under imbalanced data distributions that are prevalent in real-world applications. Such an imbalance can degrade representation quality and induce biased model behavior, yet a rigorous characterization of these effects is lacking. In this work, we develop a theoretical framework to analyze the training dynamics of contrastive learning with Transformer-based encoders under imbalanced data. Our results reveal that neuron weights evolve through three distinct stages of training, with different dynamics for majority features, minority features, and noise. We further show that minority features reduce representational capacity, increase the need for more complex architectures, and hinder the separation of ground-truth features from noise. Inspired by these neuron-level behaviors, we show that pruning restores performance degraded by imbalance and enhances feature separation, offering both conceptual insights and practical guidance. Major theoretical findings are validated through numerical experiments.