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

1 Introduction

• Contrastive learning achieves strong downstream vision performance, sometimes surpassing supervised representations, motivating analysis of how its representations are learned.It uses augmented views of unlabeled images to construct inputs and “labels,” while the required augmentation strength remains theoretically unclear.Source: §1 · p. 2

Fundamental Questions

• The paper asks how contrastive networks learn representations, why training collapses without augmentation, and how standard augmentations prevent collapse.It addresses these questions through theory on a simplified model and empirical evidence.Source: Fundamental Questions · p. 2, §1 · p. 2

1.1 Our Contributions

• Under a sparse-coding model, unaugmented contrastive learning emphasizes large-norm spurious dense noise, whereas RandomMask-style augmentation makes ReLU networks learn sparse signals efficiently with SGD.The model assumes x = Mz + ξ, where z is sparse and ξ is dense noise with norm much larger than the sparse signal.Source: §1.1 · pp. 2–3

Feature Decoupling

• Feature decoupling explains that augmentations reduce correlations of spurious features between positive samples while preserving augmentation-resistant sparse features, producing representations resembling supervised features.The supervised-learning similarity requires uncorrupted sparse signals and labels that mostly depend on those signals; empirically, color distortion improves sparse edge and shape features beyond crop-resize alone.Source: Feature Decoupling · pp. 3–4, §1.1 · p. 3

1.2 Related Work

• Prior theoretical work studies contrastive objectives, alignment, uniformity, information, spectra, or optimal solutions, but does not explain neural-network feature learning or augmentation effects.The paper also distinguishes its analysis from NTK approaches, which use prescribed feature maps and therefore cannot explain feature learning.Source: §1.2 · p. 5

2.1 Data Distribution.

• The sparse coding model generates x = Mz + ξ, where M is column-orthonormal, z ∈ {−1, 0, 1}^d is the sparse signal, and ξ is spurious dense Gaussian noise.The ambient dimension satisfies d1 = poly(d), and the sparse and dense components are treated as desired and undesired structure, respectively.Source: §2.1 · pp. 6–7
• Sparse features remain detectable because nonzero latent coordinates yield Ω(1) correlations with dictionary columns, whereas dense noise has only eO(1/√d) correlation with them.Although dense noise can have much larger norm than Mz, it cannot corrupt the sparse signal under the stated high-probability bounds.Source: §2.1 · p. 7

2.2 Learner Network and Contrastive Learning Algorithm

• The learner is a single-layer network f: Rd1 → Rm with symmetrized ReLU neurons, whose weights are interpreted as learned features aligned with either sparse dictionary directions or spurious dense directions.The desired sparse features are M, while M⊥ represents dense features associated only with ξ.Source: §2.2 · pp. 7–8
• RandomMask roughly removes half the coordinates of positive samples, making contrastive learning match disjoint coordinate sets while preserving sparse-signal correlations and removing undesired dense correlations.The method is analyzed with stochastic gradient descent both with augmentation and without augmentation; negative samples are unaugmented for theoretical simplicity.Source: §2.2 · pp. 8–9

3 Main Results

• The paper evaluates representations using linear learners on regression labels y = ⟨w⋆, z⟩ and classification labels y = sign(⟨w⋆, z⟩), arguing that learned features can be more efficient than raw inputs.Linear regression over raw inputs can require sample complexity at least √d1σξ, potentially as large as d100, while raw-input linear models can fail with constant probability.Source: §3 · p. 10

3.1 Contrastive Learning Without Augmentations

• Without augmentation, the learned representation is dominated by spurious dense signal, so downstream linear evaluation fails to achieve meaningful ℓ2-loss or accuracy regardless of labeled-data quantity.The result applies to the regression and classification tasks defined in the section.Source: §3.1 · pp. 10–11

3.2 Contrastive Learning With Augmentation

• With augmentation, contrastive learning removes spurious dense noise from the representation and performs well on both regression and classification downstream tasks.For both tasks, the required sample complexity is at most eO(d^1.001).Source: §3.2 · p. 11

• Contrastive learning comprises positive-sample similarity and positive–negative dissimilarity terms, motivating separate analyses with and without augmentations.Source: §4 · p. 12
• Without augmentation, contrastive learning emphasizes spurious dense noise because its larger ℓ2-norm overwhelms the stronger cross-sample correlations of sparse signals.Source: §4 · p. 12
• Feature decoupling uses opposite random masks D and I −D to make dense signals independent across positive samples, substantially reducing their gradients.The augmented views are generated as 2Dx_p and 2(I −D)x_p, separating dense noise between positives.Source: §4 · p. 12
• Sparse signals resist masking because they appear in both augmented views and retain their correlations, except when a signal is completely masked.Threshold decoding can recover a present sparse coefficient z_j from both masked views, after which SGD reinforces its representation.Source: §4 · p. 12

Conclusion and Discussion

• The feature decoupling principle explains how augmentations support contrastive learning, with empirical evidence that they are necessary to learn desired features and remove undesired ones.The theory aims to clarify how neural networks perform representation learning in self-supervised settings.Source: §5 · p. 13
• The results could be improved using hierarchical sparse coding models for real-life images, requiring study of contrastive learning with deeper networks.Extending the analysis beyond the current simple linear sparse coding model is identified as an important open direction.Source: §5 · p. 13

A Proof Overview

• The proof overview introduces the preliminaries before presenting the complete proof.Source: A Proof Overview · p. 14

A.1 Preliminaries and Notations

• Masked augmentations generate positive views by sampling a random mask and renormalizing the retained data by a factor of 2.Source: §A.1 · p. 14
• The population contrastive objective uses StopGrad-based similarity between positive views and incorporates independently sampled negative examples in each batch.Source: §A.1 · p. 14
• Training applies stochastic gradient descent to the contrastive objective with an L2 weight-regularization term over K = poly(d) independently sampled batches.Source: §A.1 · p. 14
• The analysis assumes population sampling rather than a finite dataset, motivated by the relative abundance and lower acquisition cost of unlabeled data.Source: §A.1 · p. 14
• Bernstein concentration yields high-probability approximation of population gradients by empirical gradients when the network weight norm satisfies ∥w(t)∥2_F ≤ poly(d).Source: §A.1 · p. 15
• The appendix defines positive and negative logits, empirical gradients, and feature-subtraction notation for original and masked inputs used throughout the proof.Source: §A.1 · p. 15
• Global conventions define notation, specify Cz as the constant in Pr(|zj| = 1) = Θ(log log d d ) = Cz log log d, and assume explicit polylog(d) factors exceed hidden eO(·) factors.Source: §A.1 · p. 15

Initial stage

• Before T1, RandomMask enables neurons to focus on learning the entire sparse-feature subspace M through feature decoupling.This initial emphasis is attributed to the RandomMask augmentation and the feature decoupling principle.Source: §A.2 · p. 15

Sparse-feature growth

• Sparse-feature components exhibit exponential growth during the initial stage, continuing until bias effects or gradient cancellation halt it.The analysis considers the growth of each neuron’s weights along sparse-feature directions.Source: §A.2 · p. 16

Dense-feature decoupling

• RandomMask suppresses growth along spurious dense-feature directions; without it, dense-feature growth would collapse the learned features.The controlled dense-feature behavior is possible specifically because RandomMask is used.Source: §A.2 · p. 16

Transition to the second stage

• After sufficient initial iterations, neuron weights consist mostly of sparse rather than dense features, enabling a second stage that tunes bias to simulate sparsification.The initial stage occupies relatively few iterations compared with total training time.Source: §A.2 · p. 16

A.3 The Second Stage of Training: Singletons Emerge

• Bias growth causes sparse–dense correlations and chained dependencies among activations and gradients, making second-stage gradient bounds substantially harder than in the initial stage.The analysis accounts for correlations involving masked sparse and dense signals in both activations and gradients.Source: §A.3 · pp. 16–17
• The theory leaves the exact bias-training mechanism open and instead analyzes manually growing biases because positive-negative contrast effects would otherwise significantly complicate the analysis.This is an explicit limitation of the second-stage treatment.Source: §A.3 · p. 16
• Lucky sparse features can maintain favorable correlations into stage II, producing gradients that make their activations much larger than those of dense features.This establishes the feature-growth mechanism leading toward singleton neurons and stage III.Source: §A.3 · p. 17
• In practice, features emerge slowly in the first layer of AlexNet during training, unlike the faster emergence observed under supervised learning.This empirical behavior matches the theoretical feature-growth picture.Source: §A.3 · p. 17

• At iterations t ≥ T2, the negative gradient begins canceling the positive gradient, driving the learning process toward convergence by training time T3 = poly(d).Source: §A.4 · p. 17
• For every sparse feature Mj, a neuron set Mj has substantial alignment with it, while neurons outside that set have only negligible alignment.Source: §A.4 · pp. 17–18
• At this stage, neuron activations are approximately decompositions of O(1) signals with only a small mixture, reflecting genuinely sparse representations.Source: §A.4 · p. 18
• Provided SGD follows a trajectory that selects only good features, the final-stage gradient points toward the desired sparse solution rather than an undesired non-convex solution.Source: §A.4 · p. 18
• Gradients for other features, including unfavored sparse features and spurious dense features, are smaller, preserving the sparse representations.Source: §A.4 · p. 18

A.5 Without Augmentations, Dense Features Are Preferred

• Without augmentations, sparse signals are not emphasized during training, so dense features can be preferred even as the learning process converges to low loss.The passage states that sparse-signal growth is approximately unchanged without augmentations and that loss and gradient convergence can still occur.Source: §A.5 · p. 18

B.1 Characterization of Neurons

• At initialization, neuron sets have limited overlap: distinct feature-associated sets intersect in at most O(log d) neurons with high probability, and each neuron belongs to at most O(1) such sets.These properties support the later characterization of the SGD trajectory.Source: §B.1 · p. 19

B.2 Activation Size and Probability

• The analysis establishes correlation and pre-activation tail bounds for sparse and dense components under RandomMask augmentation, providing the probability controls used throughout the training analysis.The bounds cover augmented correlations, pre-activation sizes, sparse signals, dense signals, and related inner products.Source: §B.2 · pp. 20–21

C Stage I: Initial Feature Growth

• At T1 = Θ(d log d / (η log log d)), Theorem C.1 establishes initial feature decoupling, including controlled responses to dense coordinates after the initial growth stage.The theorem is stated at the transition time T1, with the supplied results including a bound for every dense coordinate j ∈ [d1] \ [d].Source: C Stage I: Initial Feature Growth · p. 22

C.1 Gradient Computations

• During Stage I, biases remain zero and the positive gradients are computed using independence and zero-mean properties of masked noise and sparse signals, together with the correlation and activation bounds.The gradient derivation also uses the independence of distinct sparse coordinates and the zero mean of I − 2D.Source: §C.1 · pp. 23–24

C.2 The Learning Process at Initial Stage

• The initial-stage induction proves that weights mostly ignore spurious features M⊥ and learn to emphasize M through T1 iterations, with the induction properties maintained throughout the stage.The analysis explicitly tracks the stochastic gradient toward dense features and verifies the claimed properties iteration by iteration until T1.Source: §C.2 · pp. 25, 26, 28

D Stage II: Singleton Emerge

• Stage II analyzes how individual features can be “won” by neuron subsets and establishes an induction hypothesis over iterations t ∈ (T1, T2].The hypothesis includes a bound on how many features can significantly align with each neuron.Source: D Stage II: Singleton Emerge · p. 29

D.1 Gradient Computations

• Under the Stage II induction hypothesis, Lemma D.3 derives gradient expressions and bounds for neurons associated with sparse features.The proof decomposes the relevant expectations into terms J1, J2, and J3 and uses symmetry and randomness of the data and augmentations.Source: §D.1 · pp. 29, 33
• Lemma D.4 analyzes gradients from dense signals, using RandomMask to make the masked noise components independent and bounding contributions for both feature subspaces.The dense-feature analysis treats M_j and M_j^⊥ separately and invokes the induction hypothesis and auxiliary lemmas to obtain the desired bounds.Source: §D.1 · pp. 34–35

D.2 The Learning Process at the Second Stage

• At iteration T2, Theorem D.5 states the emergence of singleton neurons under Induction Hypothesis D.1 and additional conditions holding for each neuron.The second-stage proof considers gradients along each feature Mj and its orthogonal complement M⊥, assuming the stated properties hold for t < T2 and checking them at t + 1.Source: §D.2 · pp. 35–36

E.1 Gradient Upper and Lower Bounds for Ψ(t)

• E.1 expands each gradient projection into explicit terms and controls them with high-probability concentration, activation facts, feature-overlap bounds, and decorrelation arguments.The proof uses Fact E.3, bounds the overlap set Γj by O(log d), and decorrelates factors in the control of R1,1.Source: E Stage III: Convergence to Sparse… · p. 43, §E.1 · pp. 46–47
• E.1 establishes upper and lower bounds for Ψ(t) by decomposing gradient terms and controlling them under the induction hypotheses.The proof uses Newton–Leibniz decompositions, symmetry, Jensen’s inequality, and high-probability bounds to combine the component estimates.Source: §E.1 · pp. 49, 50, 52

E.2 Gradient Computations II

• E.2 bounds dense-feature and spurious-noise gradient terms by decomposing expectations, using feature-incidence sets of cardinality at most O(1), and applying high-probability activation bounds under Lemma B.2 conditions.For the spurious-noise terms, a different basis can yield an eO(1/√d1) factor in the second bound.Source: §E.2 · pp. 52–54

Gradient trajectory and induction

• The final-stage gradient trajectory is established by reducing the Φ(t) term to bounds on Ψ(t) under Induction Hypothesis E.2, then proving that the hypothesis persists across iterations.The induction begins at T2 using earlier lemmas and theorems, with the learning-rate and λ conditions used to maintain the required bounds.Source: §E.3 · pp. 54, 56, 57, 59

Optimal learner and loss guarantee

• The optimal learner is defined as a feature map whose performance is evaluated through a pseudo loss function, yielding the loss guarantee needed for Theorem E.1.The pseudo-objective is convex in the learner parameters because it is linear in them, and its gradient matches the network objective at the current weights.Source: §E.3 · pp. 59–61

Theorem E.1

• Theorem E.1 shows that continued training learns the desired features with sparse representations and converges to the desired solutions rather than merely bouncing around.The proof uses Induction Hypothesis E.2 together with online-learning tools and a telescoping argument to obtain the loss guarantee.Source: §E.3 · pp. 61–62

Corollary

• The corollary follows by linear regression analysis because independently generated support data form separable clusters according to differences in latent variables that determine labels.This separability holds with high probability over polynomially many support examples.Source: §E.3 · p. 62

• The feature-growth analysis is simpler without augmentation because no augmentation compensation term is needed, while feature growth remains exponential before convergence.The comparison between sparse and dense feature correlations therefore does not affect the growth-rate conclusion during the relevant training phase.Source: F Results for Learning Without… · p. 62
• Dense features have a loss advantage without augmentations under the stated orthogonality and norm conditions.The result is established using a Johnson–Lindenstrauss argument and assumes the dense-feature vectors are orthogonal to the sparse-feature directions.Source: F Results for Learning Without… · pp. 62–63
• Using the full feature vector and raising the bias above Ω(∥w_i∥2) eliminates the superior loss property in this setting.Thus, the dense-feature superiority depends on the specified feature choice and bias regime.Source: F Results for Learning Without… · p. 63
• Without augmentations, contrastive learning learns dense features in the presence of a large dense signal, producing representations unsuitable for downstream tasks.The analysis combines the preceding lemmas to characterize the learned network and concludes that meaningful downstream accuracy cannot be obtained from its feature map.Source: F Results for Learning Without… · p. 63