Understanding Dimensional Collapse in Contrastive Self-supervised Learning

In plain words

Contrastive learning trains a network so that two augmented views of the same image end up close together, while views of different images end up far apart. A common assumption is that the "push apart" part also stops the learned vectors from squeezing into a small corner of the space. This paper checks that assumption and finds it does not hold. The vectors still fill only a low-dimensional slice of the space available to them. The authors call this dimensional collapse.

They offer two explanations. The first is that augmentations which change an image too much can make the network drop some directions. The second is that stacked layers tend to find low-rank solutions even when the augmentations are mild.

Their method, DirectCLR, removes the extra trainable projection layer that SimCLR places between the encoder and the loss. The loss instead reads a fixed slice of the encoder's own output. The paper reports that this beats SimCLR with a single linear projection layer on ImageNet, but not SimCLR with a deeper projector. The theory covers only linear networks, so the explanation for nonlinear models is partly empirical.

TL;DR

• Contrastive learning avoids complete collapse, but the authors show that its embeddings still suffer dimensional collapse §1 §3.
• Two causes are identified. Strong augmentation drives collapse in a linear model §4 §4.2. Implicit regularisation in over-parameterised (deep) linear networks drives it when augmentation is weak §5 §5.3.
• DirectCLR computes InfoNCE on the first d0 channels of the ResNet-50 output, with no projector §6.2.
• ImageNet linear probe, ResNet-50, 100 epochs: DirectCLR 62.7%, SimCLR with a 1-layer linear projector 61.1%, SimCLR with a 2-layer projector 66.5% tab. 1. All numbers are author-reported.
• The theory is proven only for linear networks §4.2. The paper says it cannot yet explain why a nonlinear projector prevents collapse §6.2.

Problem

Joint-embedding self-supervised methods must avoid a trivial solution where all embeddings collapse to one point. Contrastive methods avoid this with negative pairs §1. The paper argues that non-contrastive methods can still suffer a milder collapse, in which the embeddings span fewer dimensions than available §1. It then shows that contrastive methods also suffer this dimensional collapse §1 §3.

Evidence: a SimCLR model with a two-layer MLP projector (ResNet-50, ImageNet, 100 epochs, 128-dimensional embedding) has several singular values of the embedding covariance that fall to zero §3 fig. 2 eq. 1.

The paper does not give a threshold for what counts as collapsed §3.

Method

The paper has four parts: a measurement, two analyses of collapse, and a method. Equation numbers and section numbers follow the compiled paper.

1. Measuring collapse §3. Collect embeddings z_i on the ImageNet validation set and form the covariance matrix C with N samples and mean z̄:

C = \frac{1}{N}\sum_{i=1}^{N} (z_i - \bar{z})(z_i - \bar{z})^{T} eq. 1

Take the singular values of C. Collapse shows up as values near zero §3 fig. 2.

2. Strong augmentation, linear model §4.1 §4.2. The encoder is one linear layer, z = Wx. The augmentation is additive noise. The loss is InfoNCE on unit-norm vectors, so the distance term reduces to an inner product §4.1.

L = -\sum_{i=1}^{N}\log\frac{\exp(-\lVert z_i - z_i'\rVert^2/2)}{\sum_{j\neq i}\exp(-\lVert z_i - z_j\rVert^2/2)+\exp(-\lVert z_i - z_i'\rVert^2/2)} eq. 2

Here z_i and z_i' come from the two views and z_j are in-batch negatives §4.1.

• Under gradient flow (infinitesimal step size), the weight obeys Ẇ = −G with G = −WX §4.2.
• X splits into a difference of two positive semi-definite matrices: a weighted data covariance minus a weighted augmentation covariance §4.2.
• If augmentation is strong enough that X has negative eigenvalues, W(t) = W(0) exp(Xt) loses rank. The embedding covariance is then low-rank §4.2.
• The paper calls this "strong augmentation" and states the theory holds only for linear networks §4.2.
• A simulation with isotropic Gaussian inputs, noise confined to an 8×8 block of a 16×16 matrix, and varying amplitude k confirms the rank loss §4.2 fig. 3.

3. Implicit regularisation, two-layer linear model §5.1 §5.2 §5.3. The encoder is z = W2 W1 x, with no bias and no normalisation before the loss §5.1.

• Gradient flow gives Ẇ1 = −W2ᵀG and Ẇ2 = −G W1ᵀ §5.2.
• With small augmentation, X is positive definite §5.2.
• Theorem (weight alignment): if W2 W1 is never zero, X is positive definite, and the final singular values are distinct, then the alignment matrix A = V2ᵀ U1 tends to the identity §5.3 fig. 5.
• Once aligned, the singular values of the two matrices evolve as σ̇1 = σ1 σ2² c and σ̇2 = σ2 σ1² c, where c = v1ᵀ X v1 ≥ 0 §5.3 fig. 6.
• Each pair therefore grows in proportion to itself. The smallest singular values grow most slowly, so they stay small §5.3.
• Corollary: with small augmentation and over-parameterised linear networks, the embedding covariance becomes low-rank §5.3.
• The paper says the alignment is perfect only when singular values are distinct. With random initialisation, it observes block-diagonal alignment instead §5.3.

4. DirectCLR §6.1 §6.2 §6.3. The authors propose two propositions about a linear projector W2 placed after a linear encoder W1 §6.1:

• Proposition 1: the projector only needs to be diagonal, because the orthogonal part is absorbed by the encoder §6.1.
• Proposition 2: the projector only needs to be low-rank, because the theory says its weight converges to low rank anyway §6.1.

Combining them gives a fixed low-rank diagonal projector, which in practice means keeping a fixed set of channels §6.1 §6.3.

DirectCLR uses the encoder output r directly. It keeps the first d0 channels as z = r[0:d0] and normalises z to unit length. It then applies InfoNCE to these normalised sub-vectors, with no projector §6.2. The paper states the loss without a temperature term §6.2:

# DirectCLR loss for one batch  [§6.2] [eq. 2]
# Input:  reps_a, reps_b (B, 2048): encoder outputs of the two views
#         d0: number of leading channels kept
# Output: scalar loss

def directclr_loss(reps_a, reps_b, d0):
    # Keep the first d0 channels as the embedding; no projector  [§6.2]
    sub_a = reps_a[:, :d0]                      # (B, d0)
    sub_b = reps_b[:, :d0]                      # (B, d0)

    # Scale each embedding to unit length  [§6.2]
    z_a = sub_a / norm(sub_a, axis=1, keepdims=True)  # (B, d0)
    z_b = sub_b / norm(sub_b, axis=1, keepdims=True)  # (B, d0)

    # Similarity of each view-a row with every view-b row  [§6.2]
    sims = z_a @ z_b.T                          # (B, B)

    # The positive for row i is row i; all other rows are negatives  [§6.2]
    log_probs = log_softmax(sims, axis=1)       # (B, B)

    # Negative log-likelihood of the positive pairs, summed over i  [§6.2]
    return -sum(diagonal(log_probs))

The paper uses one direction of the loss, and the sum over i §6.2. The reference code differs in several ways; see Paper vs code.

Training recipe §6.2 app. D: a standard SimCLR recipe for 100 epochs on ImageNet, with a ResNet-50 encoder whose output has 2048 channels. The two views are generated by the augmentation pipeline listed in Implementation Detail app. D. The paper does not state the augmentation probabilities or parameters.

Why the unused channels still carry information §6.2 (gradient-dynamics figure in that section): the gradient reaching the representation is non-zero only in the first d0 channels. It becomes full-rank once it passes back through the last convolutional block, and the residual connection feeds that block's output to the representation. Appendix F reports that a linear probe on the sub-vector alone reaches 47.9% app. F.

Hyperparameters

Results

Author-reported. Setup: ImageNet, ResNet-50 encoder, 100 pretraining epochs with the standard SimCLR recipe §6.2, linear probe top-1 accuracy. Number of seeds and probe protocol are not stated in the paper tab. 1.

Ablation (SimCLR with different projector settings, ImageNet top-1) tab. 2:

Other reported results (author-reported):

• Linear probe on the sub-vector alone: 47.9% app. F. The sub-vector length is not stated in that passage.
• Random d0 features per iteration instead of a fixed sub-vector, with no projector: 43.0% app. F.
• d0 sweep: accuracy drops as d0 approaches 0, and at d0 = 2048 the model reduces to SimCLR without a projector app. E. The sweep is shown only in a figure; its values are not stated in the compiled paper.
• Deeper or nonlinear encoders: the spectrum figures show more collapsed dimensions with more layers, and ReLU gives a similar effect app. C. These results are qualitative.

Paper vs code

Repository: https://github.com/facebookresearch/directclr at commit a86f64c9b0815ff19cfea9929cb7dc6bc293b5d4, linked from the paper's Reproducibility Statement (unnumbered section, after the conclusions). The repository licence is CC BY-NC 4.0 (non-commercial) LICENSE. Only the citations and paraphrases above are used; no code is quoted.

Differences between the paper's formulation and the reference code:

• Temperature. The paper's loss has no temperature term and, for unit-norm vectors, the distance form implies a scale of 1 §4.1 §6.2. The code divides the logits by 0.1 before the softmax main.py.
• Two directions. The paper's loss is one-directional §6.2. The code averages the loss over both views, in each direction main.py main.py.
• Reduction. The paper sums over i §6.2. The code averages over the batch, via cross-entropy's default reduction main.py.
• Weight decay. Not stated in the paper. The code uses 1e-6 and excludes 1-D parameters main.py main.py.
• Warm-up and final LR. The paper states warm-up and cosine decay app. D. The code starts warm-up at 0 and decays to 0.001 × base LR utils.py.
• Online classifier. The paper does not mention it. The code trains a 2048 → 1000 linear classifier on the detached encoder output and adds its cross-entropy to the total loss main.py main.py. Because the input is detached, the classifier does not change the encoder gradients. It does change the total loss value that is logged.
• SimCLR baseline widths. The paper says "two-layer MLP" §3. The code uses 2048 → 2048 → 128 with BatchNorm and ReLU main.py.
• Augmentation parameters. The paper lists the operations app. D. The code sets probabilities and jitter strengths utils.py.
• Default d0. The paper leaves the value in text unstated §6.2. The code default and the README's best-model command use 360 main.py README.md. The source also contains a commented-out table of d0 values, which is not in the compiled paper.
• Mixed precision and SyncBatchNorm. The code uses both main.py main.py. The paper does not mention them.

No other discrepancies found in the parts checked: encoder choice, DirectCLR sub-vector selection main.py, batch size, base LR, epochs, and LARS are consistent with the paper.

Pitfalls

• Temperature and loss form. The paper's loss, as written, has no temperature. The code uses 0.1 and a symmetric average. Matching the reported numbers requires the code's choices, which the paper does not state. Decide explicitly which one you implement.
• Choice of d0. The paper never states the value used for the main result in text. The sweep is only in a figure app. E. Too small a d0 starves the loss of gradient; d0 = 2048 reduces to SimCLR without a projector app. E.
• Fixed sub-vector versus random subset. A fixed choice is needed. Random feature dropout reaches 43.0% against 62.7% for the fixed sub-vector app. F.
• Theory covers linear networks. The paper states the strong-augmentation theory is limited to linear settings §4.2. For nonlinear networks it says the collapse condition "will still depend on strong augmentation" but is interpreted differently, without a formal result §4.2.
• Projector remains unexplained. The authors say their theory cannot explain why a nonlinear projector prevents collapse, and that DirectCLR still relies on the last backbone block for this §6.2.
• Gradient-flow and distinct-singular-value assumptions. The alignment and collapse results assume infinitesimal steps, positive-definite X, and distinct singular values §5.3. With random initialisation the paper only observes block-diagonal alignment §5.3.
• Positive definiteness of E[X]. The argument that E[X] is positive definite under isotropic augmentation exists in the source only as commented-out text in the appendix. The compiled paper does not prove it, so treat it as an assumption.
• Scale of the loss. Sum versus mean changes the loss scale. The paper does not analyse its effect under LARS.
• Compute. Pretraining batch 4096 needs 32 GPUs app. D. The README reports about 7 hours on 32 V100 GPUs README.md.
• Author-reported only. The ablation and the probe results come from one set of runs. Variance across seeds is not reported tab. 1 tab. 2.

Lineage

• Builds on: Tian et al., 2021, "Understanding Self-Supervised Learning Dynamics without Contrastive Pairs", whose explanation of non-contrastive collapse prevention the paper cites §2.
• Compared against: SimCLR with a 2-layer projector, SimCLR with a 1-layer linear projector, and SimCLR without a projector tab. 1. The paper also contrasts its gradient analysis with Tian et al., 2020 (arXiv:2010.00578) §4.2. That paper has no card.
• Superseded by: not known from content in this wiki.