PaperlayerView exact revision on arXiv

Source-linked AI summary

Toward Understanding the Feature Learning Process of Self-supervised Contrastive Learning

Zixin Wen, Yuanzhi Li

arXiv:2105.15134v3cs.LGcs.CVstat.ML

TL;DR

The paper asks how contrastive learning extracts useful features from unlabeled data and why it depends on strong augmentations. It analyzes this process with ReLU networks under a sparse-coding model, introducing feature decoupling to explain augmentation effects. The theory and experiments show that proper augmentations help preserve desired sparse features while suppressing spurious dense ones.

  • Problem

    The paper investigates how contrastive learning learns feature representations from unlabeled data and why its success depends on data augmentations.

  • Method

    The paper analyzes one-hidden-layer ReLU networks trained by contrastive learning on sparse signals mixed with spurious dense noise, using feature decoupling to characterize augmentation effects.

  • Results

    Contrastive learning with proper augmentations learns desired sparse features, avoids dense noise, and can produce features similar to supervised learning under stated conditions.

  • Takeaways & Limitations

    Augmentations are essential in this model because they remove correlations associated with spurious dense features while retaining correlations of sparse signals.

  • Takeaways & Limitations

    The theory uses a simple linear sparse-coding model, whereas realistic image data may require hierarchical sparse coding and deeper networks.

Abstract

from arXiv · show

How can neural networks trained by contrastive learning extract features from the unlabeled data? Why does contrastive learning usually need much stronger data augmentations than supervised learning to ensure good representations? These questions involve both the optimization and statistical aspects of deep learning, but can hardly be answered by analyzing supervised learning, where the target functions are the highest pursuit. Indeed, in self-supervised learning, it is inevitable to relate to the optimization/generalization of neural networks to how they can encode the latent structures in the data, which we refer to as the feature learning process. In this work, we formally study how contrastive learning learns the feature representations for neural networks by analyzing its feature learning process. We consider the case where our data are comprised of two types of features: the more semantically aligned sparse features which we want to learn from, and the other dense features we want to avoid. Theoretically, we prove that contrastive learning using $\mathbf{ReLU}$ networks provably learns the desired sparse features if proper augmentations are adopted. We present an underlying principle called $\textbf{feature decoupling}$ to explain the effects of augmentations, where we theoretically characterize how augmentations can reduce the correlations of dense features between positive samples while keeping the correlations of sparse features intact, thereby forcing the neural networks to learn from the self-supervision of sparse features. Empirically, we verified that the feature decoupling principle matches the underlying mechanism of contrastive learning in practice.

1 Introduction

The paper studies how contrastive learning extracts representations from unlabeled data and why augmentations are crucial. Its theory and experiments indicate that augmentations suppress dense spurious features while preserving sparse semantic features, making learned representations resemble supervised features.

  • Motivation: Contrastive learning can learn strong image representations, but understanding how it learns features and why it needs stronger augmentations remains an open problem.Prior work showed that stronger augmentations improve feature quality, with representation quality following a U-shaped relationship as augmentation strength increases.
  • Setting: The analysis models data as semantically aligned sparse signals mixed with large-norm spurious dense noise.The paper uses a sparse coding model and studies one-hidden-layer ReLU networks trained with contrastive learning.
  • Theoretical findings: Without augmentation, contrastive learning emphasizes dense noise, whereas RandomMask-style augmentation enables networks to learn sparse signals and avoid dense features.The authors attribute this difference to feature decoupling and show that the features can be learned efficiently with a variant of SGD.
  • Comparison with supervised learning: Under sparse coding assumptions, augmented contrastive learning learns features similar to supervised learning when augmentations preserve sparse signals and labels mainly depend on them.This similarity is also reflected empirically: both methods learn shape features, while contrastive learning learns fewer color features.
  • Theoretical findings: Feature decoupling reduces correlations of spurious features across positive samples while preserving correlations of features that resist augmentation.The resulting contrastive representations are predicted to retain desired sparse features and discard augmentation-sensitive dense features.
  • Empirical evidence: Experiments show a progression from dense or color-dominated representations to sparser, higher-quality edge and shape features as suitable augmentations are added.Crop-resize alone leaves dense color mixtures, while color distortion removes some of those mixtures; t-SNE and neuron activations provide supporting evidence.

2 Problem Setup

The paper models inputs as sparse signals plus spurious dense noise, then studies a ReLU contrastive learner trained with and without augmentations. RandomMask is designed to preserve sparse-feature correlations while reducing dense-feature correlations between positive views.

  • Data distribution: Inputs follow a sparse coding model x = Mz + ξ, where Mz is the sparse signal and ξ is spurious dense noise.The sparse latent vector z is sparse, while the dense noise can have much larger norm than Mz.
  • Data distribution: Sparse features are shared across samples, whereas dense noise is sample-specific and therefore poorly suited to forming separable clusters.The model treats M as the desired sparse-feature dictionary and M⊥ as the undesired dense-feature subspace.
  • Learner network: The learner is a single-layer ReLU network whose neuron weights are interpreted as learned features.The network uses a symmetrized ReLU form, and the weights are analyzed according to whether they align with sparse or dense features.
  • Contrastive objective: Contrastive learning uses positive views and negative samples, with similarity between representations forming the basis of the loss.Training is analyzed with stochastic gradient descent both with RandomMask augmentation and without augmentation.
  • RandomMask augmentation: RandomMask masks roughly half the coordinates, aiming to preserve sparse-signal correlations across views while removing undesired dense-feature correlations.The theory applies augmentation to positive samples; negative-sample augmentation is omitted for simplicity and is not expected to add further decoupling insight.

3 Main Results

The paper studies contrastive representation learning under a sparse-signal-plus-dense-noise model, showing that augmentation determines whether learned features support downstream tasks. Without augmentation, dense noise overwhelms sparse signals; with augmentation, the learned representation removes this noise and supports regression and classification.

  • Problem setting: The analysis models inputs as x = Mz + ξ, where sparse signals Mz encode task-relevant structure and dense noise ξ has much larger norm.The downstream labels depend on the latent sparse code z, while dense noise can make linear learning from raw inputs ineffective.
  • Without augmentation: Without augmentation, contrastive learning learns representations in which sparse signals are overwhelmed by spurious dense signals.The theorem and corollary imply that these representations fail to achieve meaningful downstream ℓ2-loss or accuracy, regardless of labeled-data quantity for downstream linear evaluation.
  • Without augmentation: Without augmentation, downstream linear evaluation fails to achieve meaningful ℓ2-loss or accuracy on the defined regression and classification tasks.This conclusion concerns the downstream tasks whose labels are functions of the latent variables z.
  • With augmentation: With augmentation, contrastive learning provably learns features associated with sparse signals while avoiding dense-noise features under the stated network and training conditions.The theorem uses m = d^1.01 neurons, τ = polylog(d), and poly(d) negative samples, with successful learning during iterations T ∈ [T3, T4].
  • With augmentation: For augmented representations, the learned feature map removes spurious dense noise with probability at least 1 − 1/polylog(d) and performs well on both downstream tasks.Regression and classification can each be solved using linear models with sample complexity at most eO(d^1.001).

4 Proof Intuition: The Feature Decoupling Principle

The feature decoupling principle explains why augmentation changes what contrastive learning extracts. Augmentations reduce correlations of dense noise between positive samples while preserving correlations of sparse signals, directing learning toward the sparse features.

  • Contrastive objective: Contrastive objectives emphasize similarity between positive samples and dissimilarity between positive and negative samples.The feature-learning analysis separates these two parts to explain the effects of augmentation.
  • Without augmentation: Without augmentation, positive-sample correlations are dominated by dense noise because its norm is larger, causing the learner to emphasize dense features instead of sparse ones.Sparse signals can have stronger cross-sample correlations, but the dense component dominates the relevant representation geometry in this model.
  • Feature decoupling: Feature decoupling requires augmentations to make dense signals different between positive samples while preserving correlations of sparse signals.This is the paper’s stated principle for explaining how augmentations remove spurious dense noise.
  • Dense-feature decoupling: Opposite random masks D and I − D make the dense components of positive samples independent, substantially reducing dense-feature gradients.The construction generates positive samples from complementary masked views of the same input.
  • Sparse-feature preservation: Sparse signals remain correlated across the masked views because present dictionary atoms can be recovered from both complementary masked samples, except in rare complete-masking cases.These preserved correlations are subsequently reinforced along the neural network’s SGD trajectory.

5 Conclusion and Discussion

The paper develops a theory and empirical account of how augmentations help contrastive learning retain desired features and suppress undesired ones. It identifies more realistic data distributions and deeper networks as important directions for improving the results.

  • Feature decoupling explains how augmentations help contrastive learning learn desired features while removing undesired ones.The paper supports this explanation with theoretical results and empirical evidence.
  • More realistic data distributions could significantly improve the presented results.The paper gives hierarchical sparse coding, rather than the current simple linear sparse coding model, as an example.
  • Studying contrastive learning with hierarchical sparse coding models and deeper networks remains an important open direction.

A Proof Overview

The proof overview defines a masked-augmentation contrastive-learning objective, its stochastic-gradient training procedure, and the notation used to analyze population and empirical gradients. The analysis assumes polynomially bounded dimensions and gradient-approximation conditions.

  • Preliminaries and Notations: The overview introduces notation for neuron weights, network outputs, feature-subtracted inputs, augmented inputs, and gradient terms.These definitions support the later feature-growth analysis.
  • Data Preparation and Loss Objective: Random masks generate augmented positive samples, while independently sampled data form the negative batch for the population contrastive objective.The augmentation includes a 2-factor to renormalize the data, and similarity uses a stop-gradient operation.
  • The Gradient of Weights: SGD optimizes the contrastive objective with weight regularization over K = poly(d) independently sampled batches.Each positive sample is augmented before the batches are regrouped for updates.
  • Population and Empirical Gradients: The proof analyzes population gradients through empirical gradients under bounded-weight conditions and K = poly(d) samples.The approximation holds with high probability for all iterations under the stated polynomial bounds.

A.2 The Initial Stage of Training: Initial Feature Decoupling

During the initial training stage, RandomMask and feature decoupling drive neurons toward the sparse-feature subspace while limiting growth in spurious dense directions. This establishes the representation structure before later bias tuning produces singleton features.

  • Initial Feature Decoupling: Before T1, training emphasizes the entire sparse-feature subspace through RandomMask and the feature decoupling principle.
  • Sparse Feature Growth: Sparse-feature weights grow exponentially when bias is absent, continuing until bias growth or positive-negative contrast cancels the gradient.
  • Dense Feature Suppression: RandomMask is necessary to prevent spurious dense-feature growth from collapsing the learned features.Without this augmentation, the dense-feature growth rate would be expected to increase instead.
  • Initial Feature Decoupling: After sufficient initial-stage iterations, neuron weights become dominated by sparse features rather than dense features.
  • Transition to Stage II: The second training stage tunes biases to simulate sparsification after the initial sparse-subspace learning stage.

A.3 The Second Stage of Training: Singletons Emerge

The second stage analyzes how bias growth turns neurons into singletons while accounting for correlations among masked sparse and dense signals. The proof leaves the exact bias-training mechanism open, and experiments show slower feature emergence than in supervised learning.

  • Singleton Emergence: Bias growth drives neurons toward singleton features, but it introduces chains of correlations that complicate gradient analysis.
  • Gradient Analysis: The proof uses a lemma incorporating the factors affecting gradients to control the resulting correlation chain.
  • Scope of the Analysis: The exact mechanism of bias growth is left open because positive-negative contrast makes its analysis substantially more complicated.
  • Correlation Structure: The analysis must account for positive correlations between masked sparse and dense signals in activations and gradients.
  • Sparse Feature Selection: Lucky sparse features retain favorable initialization conditions and obtain gradient growth that exceeds corresponding dense-feature effects.
  • Empirical Connection: Feature emergence in the first layer of AlexNet is slower during contrastive training than during supervised learning.

A.4 The Final Stage of Training: Convergence to Sparse Features

In the final training stage, sparse neuron activations and gradient dynamics drive the network toward sparse feature representations. The analysis argues that SGD converges to desired sparse solutions rather than remaining in undesired low-loss states.

  • Gradient dynamics: After T2, the negative gradient term begins canceling the positive gradient, causing the learning dynamics to converge.The convergence argument is developed from the stage-specific gradient behavior.
  • Neuron specialization: At t ≥ T2, each feature Mj is associated with neurons having substantial alignment, while other neurons have negligible alignment with Mj.These alignment properties organize neurons around individual sparse features.
  • Sparse activations: Neuron activations become approximately decomposable into O(1) signals plus a small mixture, reflecting sparse activation.The decomposition holds with high probability once neurons are truly sparse.
  • Convergence: SGD reaches the desired solution when its weight trajectory selects only good features, allowing the singletons of each sparse feature Mj to converge.The argument addresses how optimization avoids undesired solutions despite the initial non-convexity.

A.5 Without Augmentations, Dense Features Are Preferred

Without augmentations, dense features are not suppressed during training, while the analysis establishes conditions under which sparse-feature learning and feature decoupling can occur with augmentation. The resulting characterization supports convergence and sparse representations under the augmented process.

  • No augmentation: Without augmentation, sparse signals are not emphasized during any training stage because their growth rate is comparable to that of dense signals.The analysis also states that a corresponding dense-feature condition holds at every iteration.
  • Neuron characterization: The analysis defines neuron sets Mj and M⋆j to characterize which neurons align with sparse features during SGD.These sets support the trajectory analysis from initialization onward.
  • Initialization: At initialization, different feature-associated neuron sets overlap by at most O(log d), and each neuron belongs to at most O(1) such sets.These properties constrain feature overlap before training begins.
  • Augmentation analysis: The augmentation analysis provides concentration and tail bounds for sparse and dense signals, pre-activations, and augmented noise.These bounds are used throughout the training-process analysis.
  • Initial feature decoupling: At iteration T1, Theorem C.1 establishes initial feature decoupling, while subsequent analysis states that neurons mostly ignore spurious features M⊥ and emphasize M.The initial-stage proof uses T1 as the transition point for this separation.
  • Gradient analysis: The initial-stage analysis derives positive-gradient expressions from the augmented data and combines them with concentration bounds to obtain the stated approximations.These calculations support the transition from initialization to feature decoupling.

D.2 The Learning Process at the Second Stage

The second training stage tracks feature-specific weight growth from T1 to T2. It establishes the emergence of singleton neurons that specialize in sparse features while dense-feature components remain controlled.

  • Stage definition: The second stage runs from T1 through T2 and analyzes the evolution of sparse and dense feature components under SGD.The stage is introduced as the interval t ≥ T1 and t ≤ T2.
  • Singleton emergence: At T2, Theorem D.5 establishes the emergence of singleton neurons and additional trajectory conditions.The theorem extends the induction properties maintained during the second stage.
  • Sparse-feature growth: For feature-associated neurons, alignment with the corresponding sparse feature reaches a constant-scale value Cj = Θ(1).This establishes specialization of neurons to individual sparse features.
  • Feature separation: During the stage, non-associated feature alignments remain bounded while the analysis separately tracks dense-feature dynamics.The induction argument compares growth across sparse and dense subspaces.
  • Iteration complexity: The characteristic transition time satisfies T2 = Θ(d log d / (η log log d)).The proof verifies this order by bounding the time required for a neuron to reach the relevant weight scale.
  • Proof strategy: The second-stage proof proceeds by induction from the properties at T1, using gradient calculations along sparse and dense feature subspaces.This induction establishes the conditions needed for Theorem D.5.

E.3 Learning Process at the Final Stage

The final-stage analysis shows that, once neurons are sparsely activated, their training dynamics converge toward the desired sparse feature map. It combines induction over gradient trajectories with a pseudo-objective to establish convergence and a loss guarantee.

  • Gradient decomposition: The analysis expands gradients into feature-specific terms and controls these terms using the quantities Φ(t) and Ψ(t).This reduction enables the induction proof for the final-stage trajectory.
  • Trajectory control: The induction proof maintains sparse-feature growth while bounding non-associated sparse and dense-feature alignments throughout the final stage.The update analysis treats associated and non-associated features separately.
  • Loss analysis: The paper defines an optimal learner and a pseudo-objective, then uses online-learning arguments to derive a loss guarantee.The pseudo-objective is convex in the auxiliary parameter because it is linear in that parameter.
  • Sparse convergence: For sufficiently long training, the network learns the desired features with sparse representations rather than merely remaining at a low-loss state.The proof explicitly distinguishes convergence to the desired solution from oscillation around it.
  • Final result: The final theorem combines the induction hypothesis with the training updates to establish convergence toward the desired sparse solutions.The result follows after controlling the trajectory and applying the pseudo-objective analysis.

F Results for Learning Without Augmentations

Without augmentations, contrastive learning can favor dense features in the presence of large dense signals, producing representations that do not support meaningful downstream prediction. The analysis contrasts this outcome with the intended sparse-feature representation.

  • The no-augmentation analysis sets the augmentation compensation term aside and tracks feature growth through gradient-based lemmas.The proof uses feature-gradient and norm arguments for the one-hidden-layer ReLU network.
  • Without augmentations, dense features can have a superiority in the contrastive-learning loss under the stated orthogonality conditions.The analysis considers dense directions orthogonal to the sparse coding features and derives their advantage in this setting.
  • Using the full sparse feature matrix and raising the bias above Ω(∥w_i∥2) removes the superior-loss property in this setting.
  • The resulting no-augmentation representations are unable to support meaningful linear regression or classification accuracy on the sparse coding data.The paper states that this follows with high probability over samples from the data distribution.
Loading 2105.15134v3…