PaperlayerView exact revision on arXiv

Source-linked AI summary

Understanding Dimensional Collapse in Contrastive Self-supervised Learning

Li Jing, Pascal Vincent, Yann LeCun, Yuandong Tian

arXiv:2110.09348v3cs.CVcs.AIcs.LG

TL;DR

Contrastive self-supervised learning can still produce dimensional collapse, with embeddings spanning only a lower-dimensional subspace. The paper explains this through strong augmentation and implicit regularization, then proposes DirectCLR, which outperforms SimCLR with a linear projector on ImageNet.

  • Problem

    Contrastive self-supervised learning was known to prevent complete collapse, but evidence was needed that it can also suffer dimensional collapse into a lower-dimensional subspace.

  • Method

    The paper theoretically analyzes dimensional collapse from strong augmentation and implicit regularization, and proposes DirectCLR to optimize representations without an explicit trainable projector.

  • Results

    DirectCLR outperforms SimCLR with a linear trainable projector on ImageNet.

  • Takeaways & Limitations

    Contrastive learning is not immune to dimensional collapse, and DirectCLR provides a projector-free alternative that improves over a linear-projector SimCLR baseline on ImageNet.

  • Takeaways & Limitations

    The theory does not fully explain why nonlinear projectors prevent dimensional collapse, and DirectCLR still relies on a nonlinear-projector mechanism in the backbone.

Abstract

from arXiv · show

Self-supervised visual representation learning aims to learn useful representations without relying on human annotations. Joint embedding approach bases on maximizing the agreement between embedding vectors from different views of the same image. Various methods have been proposed to solve the collapsing problem where all embedding vectors collapse to a trivial constant solution. Among these methods, contrastive learning prevents collapse via negative sample pairs. It has been shown that non-contrastive methods suffer from a lesser collapse problem of a different nature: dimensional collapse, whereby the embedding vectors end up spanning a lower-dimensional subspace instead of the entire available embedding space. Here, we show that dimensional collapse also happens in contrastive learning. In this paper, we shed light on the dynamics at play in contrastive learning that leads to dimensional collapse. Inspired by our theory, we propose a novel contrastive learning method, called DirectCLR, which directly optimizes the representation space without relying on an explicit trainable projector. Experiments show that DirectCLR outperforms SimCLR with a trainable linear projector on ImageNet.

1 INTRODUCTION

The introduction establishes that dimensional collapse affects contrastive as well as non-contrastive self-supervised learning, despite contrastive methods’ negative pairs. It identifies two mechanisms behind this collapse and introduces DirectCLR, which optimizes representations without an explicit trainable projector and outperforms a linear-projector SimCLR on ImageNet.

  • Self-supervised visual representation learning seeks useful input representations without human annotations through joint embedding methods.
  • Contrastive methods use positive and negative sample pairs to avoid complete collapse, whereas non-contrastive methods use stop-gradient and an extra predictor.
  • Dimensional collapse occurs when embedding vectors span a lower-dimensional subspace rather than the entire available embedding space, and it affects contrastive learning too.
  • Two mechanisms drive contrastive dimensional collapse: strong augmentation along feature dimensions and implicit regularization toward low-rank solutions.
  • DirectCLR directly optimizes the encoder representation space without an explicit trainable projector and outperforms SimCLR with a linear trainable projector on ImageNet.

2 RELATED WORKS

Related work positions contrastive learning within joint-embedding self-supervision, alongside theoretical efforts to explain self-supervised learning and implicit regularization toward low-rank solutions.

  • Self-supervised Learning Methods: Joint-embedding methods match representations of augmented views, while contrastive methods compare training samples using losses such as InfoNCE.Contrastive methods effectively treat each sample as its own class.
  • Theoretical Understanding of Self-supervised Learning: The dynamics underlying successful self-supervised learning methods remain poorly understood, motivating several theoretical analyses.These methods have outperformed supervised counterparts on several downstream transfer-learning benchmarks.
  • Implicit Regularization: Theoretical work on implicit regularization shows that gradient descent can produce low-rank solutions in deep linear networks.This line of work builds on aligned adjacent matrices and minimal nuclear-norm solutions.

3 DIMENSIONAL COLLAPSE

Contrastive methods avoid complete collapse by pushing embeddings from different images apart, but their representations can still occupy a lower-dimensional subspace. SimCLR analysis confirms this dimensional collapse through zero singular values in the embedding covariance spectrum.

  • Dimensional collapse: Contrastive methods prevent complete collapse through a negative term that pushes embeddings of different input images apart, yet still experience dimensional collapse.The embeddings occupy a lower-dimensional subspace than their available dimension.
  • Experimental analysis: A SimCLR model with a two-layer MLP projector was trained on ImageNet for 100 epoch using 128-dimensional embeddings.Dimensionality was evaluated from validation-set embeddings by computing their covariance matrix.
  • Experimental analysis: Singular value decomposition of the embedding covariance matrix shows that a number of singular values collapse to zero, representing collapsed dimensions.The singular values are examined in sorted order on a logarithmic scale.
  • Experimental analysis: The SimCLR embedding spectrum contains singular values that drop to zero, indicating collapsed dimensions in the 128-dimensional embedding space.The spectrum is computed from validation-set embeddings of a pretrained SimCLR model on ImageNet.

4 DIMENSIONAL COLLAPSE CAUSED BY STRONG AUGMENTATION

The section analyzes dimensional collapse in a linear contrastive model when additive augmentation noise surpasses the input information. It shows that strong augmentation causes vanishing singular values in the weights and a low-rank embedding covariance.

  • Model and dynamics: The analysis considers a single-layer linear network, z = Wx, trained with InfoNCE and additive-noise augmentation.The dynamics are studied using basic stochastic gradient descent and, specifically, gradient flow.
  • Model and dynamics: The weight dynamics depend on X, a difference between weighted data-distribution and augmentation-distribution covariance matrices.The augmentation amplitude determines whether X is positive definite; the result applies to multiple negatives, varying pairwise weights, and finite batches.
  • Theoretical result: With fixed X having negative eigenvalues under strong augmentation, the weight matrix W develops vanishing singular values.This is the stated condition in Theorem 1 for dimensional collapse dynamics.
  • Theoretical result: Strong augmentation makes the embedding covariance matrix low-rank, indicating that embedding dimensions collapse.The covariance is formed from transformed centered inputs, so vanishing singular values of W make the covariance low-rank.
  • Numerical verification: Numerical simulation with isotropic Gaussian data and additive Gaussian augmentation verifies that increasing augmentation amplitude causes dimensional collapse in the linear embedding space.The augmentation covariance uses an 8x8 block structure scaled by amplitude k, and the weight singular-value spectrum is plotted for varying k.
  • Limitation: For nonlinear networks, the collapse condition still depends on strong augmentation but also on augmentation statistics, manifold structure, and network capacity.The linear-network theory does not directly cover these more complex settings.

5 DIMENSIONAL COLLAPSE · CAUSED BY IMPLICIT REGULARIZATION

The paper shows that dimensional collapse can arise in contrastive learning through implicit regularization: strong augmentation collapses a limited-capacity linear model, while deep over-parameterized networks can collapse even with small augmentation. In a two-layer linear model, weight alignment and uneven singular-value growth drive the embedding covariance toward low rank.

  • 5.1 TWO-LAYER LINEAR MODEL: Strong augmentation causes dimensional collapse under InfoNCE in linear models with limited capacity, whereas a single linear model avoids it with small augmentation and positive-semidefinite X.Deep networks nevertheless exhibit dimensional collapse in practice.
  • 5.1 TWO-LAYER LINEAR MODEL: The analyzed over-parameterized model is a bias-free two-layer linear MLP whose embeddings are z = W2W1x and are trained with InfoNCE without normalization.Both weight matrices are d×d, and augmentation is modeled as additive noise.
  • 5.3 WEIGHT ALIGNMENT: Under nonzero products, positive-definite X, and distinctive limiting singular values, the two weight matrices align through their corresponding singular vectors.Empirically, the absolute alignment matrix converges to the identity under InfoNCE.
  • 5.3 WEIGHT ALIGNMENT: Corresponding singular values receive gradients proportional to themselves, so the smallest singular values grow substantially more slowly.Degenerate singular values can prevent perfect alignment because their singular decomposition is not unique.
  • 5.3 WEIGHT ALIGNMENT: With small augmentation and over-parameterized linear networks, implicit regularization makes the embedding-space covariance matrix low-rank.This result is stated as the paper’s corollary on dimensional collapse caused by implicit regularization.
  • 5.3 WEIGHT ALIGNMENT: As W2W1 evolves toward low rank, the embedding covariance C also becomes low rank, indicating collapsed dimensions.Figure 6 experimentally verifies the low-rank covariance and persistently small lowest singular values.
  • 5.3 WEIGHT ALIGNMENT: The theory extends beyond the two-layer linear setting to multilayer networks and nonlinear models.The paper directs the broader extension to Appendix C.

6 DIRECTCLR

DirectCLR removes the explicit trainable projector by applying InfoNCE to a normalized sub-vector of the representation, motivated by the finding that effective projectors can be low-rank and diagonal. On ImageNet, it outperforms SimCLR with a trainable linear projector and prevents dimensional collapse, while its explanation of nonlinear projectors remains incomplete.

  • Motivation: Without a projector, SimCLR’s representation space undergoes dimensional collapse, whereas a projector substantially improves representation quality and downstream performance.This is shown through representation-spectrum comparisons of SimCLR models trained with and without a projector.
  • Theory: A linear projector need only be diagonal and low-rank, because adjacent trainable layers make its orthogonal component redundant and its behavior depend on singular values.The preceding encoder layer can align its orthogonal component with the projector, while the singular-value matrix naturally converges to low rank.
  • Method: DirectCLR sends a fixed sub-vector z = r[0 : d0] directly to a normalized standard InfoNCE loss, without an explicit trainable projector.The sub-vector dimension d0 is a hyperparameter, and the normalized vector is ẑ = z/|z|.
  • Experiments: DirectCLR outperforms SimCLR with a trainable linear projector on ImageNet after 100 epochs using a ResNet50 backbone and standard training recipe.Linear probe accuracies are reported in Table 1.
  • Results: DirectCLR prevents dimensional collapse similarly to SimCLR with projectors, while ablations support diagonal and low-rank projector formulations, with low-rank achieving the highest accuracy.The ablations find equivalence between orthogonal-constrained and non-projector settings, fixed low-rank and fixed diagonal projectors, and trainable linear and trainable diagonal projectors.
  • Limitations: DirectCLR’s theory replaces linear projectors but does not fully explain nonlinear projectors, whose collapse-prevention mechanism is still supplied by the backbone’s last block.Although only the first d0 representation channels receive direct gradient, nonlinear processing makes gradients full rank in the preceding hidden layer.

7 CONCLUSIONS … B.4 PROOF OF LEMMA 3

The paper concludes that contrastive self-supervised learning suffers dimensional collapse through strong augmentation and implicit regularization, motivating DirectCLR. The accompanying lemmas and proofs characterize singular-value, singular-vector, alignment, and rank-deficiency dynamics.

  • 7 CONCLUSIONS: Contrastive self-supervised learning suffers dimensional collapse, with embeddings spanning only a lower-dimensional subspace through strong augmentation and implicit regularization.
  • 7 CONCLUSIONS: DirectCLR is proposed to directly optimize the representation space without relying on an explicit trainable projector.
  • A USEFUL LEMMAS: The adapted lemmas describe how a matrix’s singular values evolve as σ̇_k = u_k^T Ẇv_k under matrix dynamics.
  • A USEFUL LEMMAS: The singular vectors U and V evolve through Hadamard products involving the matrix dynamics and a skew-symmetric matrix H.
  • A USEFUL LEMMAS: The alignment matrix A evolves according to coupled Hadamard-product terms involving A, F, H1, and H2.
  • B.1 PROOF OF LEMMA 1: The proof of Lemma 1 begins by deriving the gradient on W and denoting the gradients on paired representations by g_zi and g_z′i.
  • B.2 PROOF OF LEMMA 2: The proof of Lemma 2 uses the definition of X and reindexes pairwise sums over all i, j with j ≠ i.
  • B.3 PROOF OF THEOREM 1: For fixed X, negative eigenvalues make exp(Λt) rank deficient as t →∞, so W(∞) is rank deficient with vanishing singular values.

B.5 PROOF OF THEOREM 2 · B.6 PROOF OF THEOREM 3

Under non-degenerate singular values, the proof establishes alignment of corresponding singular-vector directions and shows that the alignment matrix converges to the identity. It also derives divergence of the weight-matrix norms and notes that alignment is not guaranteed when singular values are degenerate.

  • B.5 PROOF OF THEOREM 2: Under non-degenerate singular values, the alignment matrix A = V_2^T U_1 converges to the identity matrix.The proof assumes singular values are non-degenerate and invokes uniqueness of the eigen-decomposition.
  • B.5 PROOF OF THEOREM 2: The proof uses Lemma 3 and singular-value decompositions of W_1 and W_2 to derive the alignment result.The derivation includes substituting the decompositions into the relevant matrix expression.
  • B.5 PROOF OF THEOREM 2: The proof shows that the Frobenius norm of each weight matrix grows without bound.This conclusion follows from the positive trace argument for the matrix B.
  • B.5 PROOF OF THEOREM 2: Because X is positive definite and W_2(t)W_1(t) ≠ 0, B is nonzero positive semidefinite.The trace is therefore positive because not all eigenvalues of B are zero.
  • B.5 PROOF OF THEOREM 2: If singular values are degenerate, the corresponding singular vectors are non-unique, so the associated dimensions do not become aligned.This is explicitly identified as a limitation of the non-degeneracy assumption.
  • B.6 PROOF OF THEOREM 3: For matching indices, the corresponding singular vectors u_k^2 and u_k^1 become aligned, expressed through their inner-product relation.The proof of Theorem 3 derives this alignment after inserting Eqn 9 and accounting for the alignment.

C EFFECT OF MORE LAYERS AND NONLINEARITY … E HYPERPARAMETER TUNING ON d0

The theory extends to multilayer and nonlinear networks, where additional layers and ReLU produce stronger dimensional collapse. ImageNet experiments specify the augmentation, ResNet-50, optimization setup, and d0-dependent accuracy behavior.

  • C EFFECT OF MORE LAYERS AND NONLINEARITY: More layers amplify dimensional collapse, whereas a single layer (L = 1) has no dimensional collapse in the embedding space.The amplification is attributed to alignment among adjacent matrices and collapse in their product matrix.
  • C EFFECT OF MORE LAYERS AND NONLINEARITY: Adding layers creates more collapsed dimensions, and adding nonlinearity produces a similar collapsing effect.Figure 11 uses 16x16 weight matrices for linear and nonlinear networks.
  • C EFFECT OF MORE LAYERS AND NONLINEARITY: Inserting ReLU between linear layers yields singular value collapse similar to the linear case.This nonlinear collapse is shown in Figure 11b.
  • D.1 AUGMENTATIONS: Each image produces two distorted views using random cropping, resizing to 224x224, flipping, color transformations, grayscale, blurring, and solarization.The two views are used for contrastive loss.
  • D.2 NETWORK: ImageNet experiments use a ResNet-50 encoder with a 2048-dimensional output representation vector.The encoder is identified as He et al. (2016).
  • D.3 OPTIMIZATION: Models use LARS for 100 epochs with batch size 4096 across 32 GPUs and learning rate 4.8, followed by cosine decay after 10 warmup epochs.The learning rate follows the SimCLR setup.
  • E HYPERPARAMETER TUNING ON d0: When d0 →0, insufficient loss-gradient information reduces performance, while d0 →2048 becomes standard SimCLR without a projector and suffers representation-space dimensional collapse.The d0 comparison is based on ImageNet accuracy.
  • E HYPERPARAMETER TUNING ON d0: Figure 12 evaluates d0 using ImageNet linear probe Top-1 accuracy.The figure reports hyperparameter tuning results for d0.

F ABLATION STUDY DETAIL

The ablations support using a diagonal, low-rank projector and show that DirectCLR’s fixed subvector preserves useful information beyond the dimensions directly optimized by the loss.

  • Projector structure: DirectCLR matches SimCLR with a fixed low-rank projector, achieving 62.3% linear probe accuracy.DirectCLR is equivalent to SimCLR with a fixed low-rank diagonal projector; the fixed projector additionally applies a fixed orthogonal transformation.
  • Projector structure: A trainable diagonal projector achieves 60.2% linear probe accuracy on ImageNet, close to SimCLR with a 1-layer linear projector.
  • Projector structure: The ablations verify that the SimCLR projector only needs to be diagonal and that low-rank settings consistently improve performance.
  • Subvector ablation: 47.9% accuracy from probing only DirectCLR’s subvector z shows that the remaining dimensions r retain useful information despite receiving no direct loss gradient.
  • Subvector ablation: 43.0% accuracy from randomly dropping features each iteration demonstrates that a fixed subvector is important for enabling the alignment effect.
Loading 2110.09348v3…