PaperlayerView exact revision on arXiv

Source-linked AI summary

Whitening for Self-Supervised Representation Learning

Aleksandr Ermolov, Aliaksandr Siarohin, Enver Sangineto, Nicu Sebe

arXiv:2007.06346v5cs.LGcs.CVstat.ML

TL;DR

Contrastive SSL often depends on many negative samples, creating computational demands. The paper introduces Whitening MSE, which whitens latent features and compares positive pairs without negatives; it reports competitive performance and improved results from multiple positive pairs.

  • Problem

    Contrastive SSL methods require many negative samples to be competitive, making large-batch training computationally demanding.

  • Method

    Whitening MSE whitens latent-space features into a spherical distribution and computes an MSE loss over positive pairs without using negatives.

  • Results

    W-MSE outperforms the commonly adopted contrastive loss and is competitive with state-of-the-art SSL methods.

  • Takeaways & Limitations

    Because W-MSE does not need many negatives, it can include multiple positive pairs from one image, which the paper demonstrates improves performance.

  • Takeaways & Limitations

    The number of positive samples per image may vary, trading off batch diversity against the amount of training.

Abstract

from arXiv · show

Most of the current self-supervised representation learning (SSL) methods are based on the contrastive loss and the instance-discrimination task, where augmented versions of the same image instance ("positives") are contrasted with instances extracted from other images ("negatives"). For the learning to be effective, many negatives should be compared with a positive pair, which is computationally demanding. In this paper, we propose a different direction and a new loss function for SSL, which is based on the whitening of the latent-space features. The whitening operation has a "scattering" effect on the batch samples, avoiding degenerate solutions where all the sample representations collapse to a single point. Our solution does not require asymmetric networks and it is conceptually simple. Moreover, since negatives are not needed, we can extract multiple positive pairs from the same image instance. The source code of the method and of all the experiments is available at: https://github.com/htdt/self-supervised.

1. Introduction

SSL avoids manual annotation by learning from automatically constructed prediction tasks, but contrastive approaches depend on many negatives. The paper proposes Whitening MSE, which uses whitening to prevent collapse while learning from positive pairs alone.

  • Motivation: SSL replaces manual labels with prediction problems derived from context or pretext tasks.Examples include temporal consistency, spatial patch order, and image transformations.
  • Motivation: Large numbers of negative samples make contrastive SSL computationally demanding, especially with high-resolution images.Contrastive losses require large negative sets to be competitive, motivating alternatives that reduce batch-size demands.
  • Motivation: BYOL uses only positives and two networks, but the mechanism preventing representation collapse remains unclear.The paper identifies batch normalization as an important implicit ingredient and generalizes this idea through whitening.
  • Proposed approach: Whitening MSE first scatters batch representations into a spherical distribution, then penalizes distant positive pairs using normalized features.This removes the need to contrast positives against negatives to avoid collapse.
  • Proposed approach: W-MSE constrains batch samples to a spherical distribution as an alternative to positive-negative instance contrasting methods.The spherical-distribution assumption means zero mean and identity-matrix covariance.
  • Consequences: Multiple positive pairs extracted from one image improve performance without requiring a large number of negatives.The approach is presented as avoiding asymmetric networks and competing with state-of-the-art SSL methods.

2. Background and Related Work

This section reviews pretext tasks, contrastive and triplet losses, positive-only methods, clustering approaches, and feature whitening. It emphasizes the computational cost of large negative sets and the role of batch normalization in preventing collapse.

  • Overview: SSL combines a pretext task that extracts supervision from data with a loss function.The review covers both components and feature-whitening methods.
  • Pretext Tasks: Video SSL commonly exploits temporal consistency, including nearby-frame similarity, temporal order, and temporal cycle consistency.These methods use temporal structure as automatically available supervision.
  • Pretext Tasks: Still-image instance discrimination creates two augmented views of an image and distinguishes them from views of other samples.Common transformations include cropping, rotation, color jittering, and Sobel filtering.
  • Loss functions: Contrastive loss compares positive representations against negative pairs, with temperature τ and a denominator containing K −1 negatives.The formulation typically uses the current batch, while some methods maintain negatives beyond the batch.
  • Loss functions: Contrastive loss requires many negatives to be competitive, making large-batch solutions computationally demanding.Earlier methods precompute representations of old samples, whereas SimCLR uses large batches.
  • Positive-only methods: BYOL and SimSiam use positive pairs without conventional negative comparisons, while incorporating asymmetric prediction components and batch normalization.Removing batch-normalization layers causes dramatic performance degradation in SimSiam, and batch normalization helps BYOL avoid constant representations.
  • Clustering: SwAV uses cross-entropy over image-to-cluster prototype assignments and includes multiple positives among its ingredients.The paper positions W-MSE as exploiting multiple crops more efficiently.
  • Feature Whitening: The paper adopts a Cholesky-based whitening transform to project latent vectors into a spherical distribution.The cited whitening transform is distinguished from whitening used for multivariate batch normalization in intermediate layers.

3. The Whitening MSE Loss

W-MSE whitens batch features into a zero-centered, identity-covariance distribution, then minimizes normalized MSE across positive pairs. The method uses multiple augmented samples per image and batch slicing to estimate whitening matrices over sub-batches.

  • Loss formulation: W-MSE requires embeddings to be non-degenerate while clustering positive image pairs closely.The identity-covariance constraint prevents all representations from collapsing and encourages distinct embedding dimensions.
  • Loss formulation: Normalized-vector MSE implements cosine similarity, penalizing positive pairs that remain far apart after whitening.The loss shrinks inter-positive distances while the spherical constraint rearranges other samples.
  • Positive samples: d augmented samples from one image yield d(d −1)/2 positive comparisons, with experiments using d = 2 and d = 4.Larger d increases training signal and diversity but can cause instability through its relationship with sub-batch size.
  • Training pipeline: The encoder and projection head produce features that whitening maps to the spherical space used by W-MSE.A standard ResNet encoder is followed by a nonlinear projection head.
  • Loss formulation: Whitening transforms batch features into a zero-centered distribution with covariance matrix I before positive-pair comparisons.The resulting whitened vectors are used in the loss after normalization.
  • Batch slicing: Batch slicing partitions and permutes samples, then computes independent whitening matrices and centroids for sub-batches.Using all sub-batch elements in the loss and averaging repeated samplings helps alleviate differences among whitening matrices.
  • Discussion: Unlike instance-contrastive losses, W-MSE imposes a soft spherical constraint rather than forcing every instance pair far apart.This avoids treating semantically similar samples from different images as necessarily dissimilar.
  • Discussion: Batch normalization alone produced 0 loss after 50 epochs but only 68.15% linear classification accuracy on CIFAR-10.The authors use this result to distinguish whitening from feature standardization without whitening.

4. Experiments

The experiments compare W-MSE with other self-supervised losses across datasets using controlled encoder settings and linear evaluation. W-MSE performs competitively, while whitening alone does not improve contrastive learning and can be harmful when combined with negative contrasting.

  • Experimental setup: Experiments use CIFAR-10, CIFAR-100, Tiny ImageNet, STL-10, ImageNet-100, and ImageNet to evaluate self-supervised representations.Datasets range from small 32 × 32 images to large-scale ImageNet benchmarks.
  • Experimental setup: The comparisons control architectural effects by using ResNet-18 for small and medium datasets and ResNet-50 for ImageNet-based experiments.Methods generally share training schedules and use L2-normalized latent features.
  • Comparison with the state of the art: W-MSE 4 is best on CIFAR-10 and CIFAR-100, while BYOL leads on STL-10 and Tiny ImageNet with a marginal gap.For W-MSE, four samples per image are generally better than two, and contrastive loss performs worst in most cases.
  • Comparison with the state of the art: After 20 epochs, W-MSE 4 reaches 79.98% accuracy versus 73.24% for BYOL, while the methods perform similarly during most later training.The authors report that W-MSE performs well from the beginning, whereas BYOL requires a warmup period.
  • Comparison with the state of the art: On ImageNet-100, both W-MSE variants significantly outperform the compared methods despite using a lower-capacity ResNet-18 encoder.W-MSE 2 remains highly competitive without multi-cropping.
  • Comparison with the state of the art: On ImageNet, W-MSE 4 is state of the art at 100 epochs and close to the 400-epoch state of the art.The comparison uses ResNet-50 encoders, although some configuration details differ across methods.
  • Contrastive loss with whitening: Whitening alone does not improve self-supervised performance and can be harmful when combined with negative contrasting.Whitened and normalized contrastive features produced unstable training, while unnormalized whitening converged but underperformed standard contrastive loss.

5. Conclusion

The paper presents W-MSE as a simpler alternative to common SSL losses, using whitening to avoid collapsed representations without asymmetric networks. It reports classification accuracy that is generally comparable or superior to state-of-the-art methods.

  • Conclusion: W-MSE computes inter-positive distances while using whitening to avoid degenerate representations.Unlike triplet and contrastive losses, it does not compare positive pairs against other samples.
  • Conclusion: W-MSE avoids the asymmetric network architectures and training protocols used by BYOL and SimSiam.The paper characterizes its solution for avoiding collapse as simpler than asymmetry-based approaches.
  • Conclusion: ImageNet evaluation uses top-1 classification accuracy from a linear classifier with a ResNet-50 encoder.The supplied table caption identifies the evaluation setting and notes that most comparison results come from Chen and He (2020).
  • Conclusion: A separate CIFAR-10 experiment evaluates contrastive loss with whitened features after 200 training epochs.The supplied table caption identifies this experiment but does not provide its accuracy value.
  • Conclusion: Classification accuracy is, in most cases, comparable or superior to state-of-the-art methods.The conclusion reports this outcome without specifying a single benchmark value in the supplied passages.

A. Training Dynamics

The training-dynamics figures compare the methods listed in Table 1 and smooth all curves with a 0.3 moving average for readability.

  • Training Dynamics: Figures 4 and 5 show the training dynamics of methods compared in Table 1.The passage identifies the figures’ subject but does not state a performance outcome.
  • Training Dynamics: All training-dynamics curves are smoothed with a 0.3 moving average.The smoothing is applied to improve readability.
  • Training Dynamics: The unsmoothed curves remain visible as semi-transparent lines.This preserves access to the original dynamics alongside the smoothed curves.

B. Cholesky Whitening and Backprogation

The whitening matrix is computed with a differentiable Cholesky decomposition of the batch covariance matrix, then inverted to transform the features.

  • Cholesky Whitening: The covariance matrix is factorized as Σ_V = LL^⊤ using two triangular matrices.L is the lower triangular matrix in the Cholesky factorization.
  • Cholesky Whitening: The whitening matrix is obtained by computing the inverse of the Cholesky factor: W_V = L^-1.The passage defines W_V after L has been computed.
  • Cholesky Whitening: Cholesky decomposition is fully differentiable and implemented in major frameworks including PyTorch and TensorFlow.This supports differentiable use of the whitening operation during training.

B.1. Gradient Computation

The paper gives the backpropagation equations for whitening by differentiating through the whitened batch, covariance, and whitening matrix.

  • Gradient Computation: The whitened batch is defined as Z = W_V(V − μ_V).The gradient with respect to V is then computed from the whitening differentiation equations.
  • Gradient Computation: The gradient with respect to V is obtained by backpropagating through Z and differentiating the covariance contribution.The supplied passages state that ∂L/∂Z is backpropagated and ∂L/∂Σ is computed separately.
  • Gradient Computation: The covariance-gradient expression uses the Hadamard product, and the whitening-matrix gradient is given separately.The passage identifies ◦ as the Hadamard product in equation 10.

C. Training time complexity

The whitening transform has complexity O(k^3 + Mk^2), which reduces to O(Mk^2) when the embedding dimension is smaller than the sub-batch size. This cost is comparable to a fully connected forward pass and is not the dominant training expense.

  • Whitening costs O(k^3 + Mk^2), where k is the embedding dimension and M is the sub-batch size.The stated complexity follows the whitening-transform analysis.
  • When k < M, the whitening transform simplifies to O(Mk^2).This condition is the paper’s stated regime for the reduced complexity.
  • The resulting cost is comparable to a fully connected forward pass and is usually smaller than other architectural training costs.The paper contrasts whitening with more expensive architectural choices, including BYOL’s multiple network passes.

D. Euclidean distance

The Euclidean-distance experiment evaluates unnormalized embeddings on STL-10 using top-1 classification accuracy. Whitening preserves performance better than the other methods under this unnormalized evaluation, while the accompanying figures show training dynamics for linear and 5-nearest-neighbor classifiers.

  • D. Euclidean distance: Table 5 reports top-1 classification accuracy on STL-10 using Euclidean distance with unnormalized embeddings.The table caption specifies the dataset, metric, and embedding treatment.
  • D. Euclidean distance: W-MSE has a significantly smaller accuracy decrease than the other methods when embeddings are unnormalized.This comparison is made against the corresponding normalized-embedding results in Table 1.
  • D. Euclidean distance: The training-dynamics figures separately use linear-classifier and 5-nearest-neighbor evaluations on STL-10.Figure 4 uses linear classification, while Figure 5 uses a 5-nn classifier.
Loading 2007.06346v5…