[Paper Review] Whitening for Self-Supervised Representation Learning

Aleksandr Ermolov, Aliaksandr Siarohin+2 more|arXiv (Cornell University)|Jul 13, 2020

Domain Adaptation and Few-Shot Learning61 references43 citations

TL;DR

Proposes a self-supervised loss (W-MSE) that uses full feature whitening to enforce a spherical latent distribution and rely solely on positive pairs, removing the need for negatives or asymmetric architectures.

ABSTRACT

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.

Motivation & Objective

• Motivate reducing reliance on large negative sets in self-supervised learning (SSL).
• Introduce a whitening-based loss that yields nondegenerate, spherical latent representations.
• Enable multiple positive pairs from a single image to improve SSL signal.
• Demonstrate that W-MSE is competitive with state-of-the-art SSL methods across standard benchmarks.

Proposed method

• Define a whitening MSE (W-MSE) loss that minimizes inter-positive distances after whitening the batch features.
• Project backbone features to a latent space, apply a whitening transform to obtain a spherical distribution (zero mean, identity covariance), and compute MSE over normalized positive pairs.
• Use a batch-level whitening with mean-centering and a whitening matrix W that satisfies W^T W = Sigma^{-1} (Sigma is the batch covariance).
• Collect multiple positive samples from the same image via standard augmentations; compute all positive pairs among d positives.
• Employ batch slicing to stabilize whitening by computing whitening in sub-batches and averaging results across iterations.

Experimental results

Research questions

• RQ1 Can whitening of latent features avoid representation collapse without using negatives or asymmetric networks?
• RQ2 How does a whitening-based SSL loss compare to contrastive and other state-of-the-art SSL losses on standard benchmarks?
• RQ3 Does leveraging multiple positive samples from the same image improve SSL performance with W-MSE?
• RQ4 What are the practical considerations (e.g., batch slicing) to stabilize whitening in SSL training?

Key findings

• W-MSE consistently performs competitively with state-of-the-art SSL methods across CIFAR-10, CIFAR-100, STL-10, Tiny ImageNet, ImageNet-100, and ImageNet under various settings.
• Using four positives (W-MSE 4) generally yields better accuracy than using two positives (W-MSE 2).
• The approach achieves strong linear and k-NN classification performance without negative samples or momentum networks, and it warms up faster than some competing methods like BYOL.
• Whitening alone does not improve SSL when paired with a contrastive loss; negatives are still needed for effective contrastive learning.
• W-MSE 4 reaches state-of-the-art performance on 100-epoch ImageNet-style evaluations, closely matching longer training regimes with greater simplicity.