Source-linked AI summary
On Feature Decorrelation in Self-Supervised Learning
Tianyu Hua, Wenxiao Wang, Zihui Xue, Sucheng Ren, Yue Wang, Hang Zhao
TL;DR
Self-supervised augmentation robustness does not by itself prevent collapsed representations. This paper studies a concise framework, identifies complete and dimensional collapse, and connects dimensional collapse to feature correlations; empirically, feature decorrelation improves utility within the reported settings.
Problem
Augmentation-robust representation learning does not by design exclude collapse, and dimensional collapse has been overlooked despite existing methods addressing complete collapse.
Method
The paper studies a concise framework and uses variance standardization, covariance standardization through DBN, and a further decorrelation add-on called Shuffled-DBN.
Results
The paper verifies complete collapse, finds dimensional collapse reachable, links it to strong axis correlations, and reports empirical gains from feature decorrelation.
Takeaways & Limitations
Feature decorrelation is presented as important and potentially useful for addressing collapse and improving self-supervised representation utility.
Takeaways & Limitations
On ImageNet, Shuffled-DBN does not achieve state-of-the-art performance, although it remains promising in a concise framework.
Abstract
from arXiv · showhide
In self-supervised representation learning, a common idea behind most of the state-of-the-art approaches is to enforce the robustness of the representations to predefined augmentations. A potential issue of this idea is the existence of completely collapsed solutions (i.e., constant features), which are typically avoided implicitly by carefully chosen implementation details. In this work, we study a relatively concise framework containing the most common components from recent approaches. We verify the existence of complete collapse and discover another reachable collapse pattern that is usually overlooked, namely dimensional collapse. We connect dimensional collapse with strong correlations between axes and consider such connection as a strong motivation for feature decorrelation (i.e., standardizing the covariance matrix). The gains from feature decorrelation are verified empirically to highlight the importance and the potential of this insight.
1. Introduction
Self-supervised learning commonly seeks representations robust to augmentations, but this objective does not by itself exclude collapse. The paper studies collapse patterns and motivates feature decorrelation as a remedy.
- Self-supervised methods commonly learn representations robust to predefined augmentations.
- Augmentation robustness alone does not rule out trivial collapsed solutions by design.
- Complete collapse produces constant representations and is typically mitigated through method-specific implementation details.
- Feature decorrelation is motivated by its value in discriminative and generative machine-learning tasks.
- The paper uses a concise framework containing common components of existing approaches to revisit collapse and feature decorrelation.
- The contributions include identifying dimensional collapse, linking it to strong axis correlations, and empirically evaluating covariance standardization.
2. Related Work
Related work spans contrastive, clustering-based, positive-pair, and normalization approaches to representation learning. The paper positions feature decorrelation and normalization as an alternative perspective on collapse prevention.
- Contrastive learning: Contrastive methods maximize agreement for positive pairs and disagreement for negative pairs, often requiring memory banks, queues, or large batches.
- Clustering: Clustering-based methods discriminate among groups of similar images but can require costly clustering and large batches.
- BYOL and SimSiam: BYOL and SimSiam prevent collapse through different combinations of target-network updates, predictor networks, and stop-gradient operations.
- Normalization: The paper argues that BN encounters dimensional collapse, whereas DBN effectively avoids all kinds of collapse in self-supervised learning.
- Concurrent methods corroborate feature decorrelation’s potential through whitening or cross-correlation regularization.
3. Main Results
The paper analyzes complete and dimensional collapse in a concise self-supervised framework. It associates the former with vanishing variance and the latter with strong correlations, then evaluates further decorrelation through Shuffled-DBN.
- Reachable Collapse Patterns: Complete collapse is associated with vanishing variances and can be addressed with Batch Normalization.
- Reachable Collapse Patterns: Dimensional collapse is an overlooked reachable pattern in which representations collapse into a lower-dimensional structure.
- Visualization: Figure 2 compares concise-framework projection spaces with SimCLR and supervised two-dimensional baselines on CIFAR-10.
- Feature Decorrelation: Dimensional collapse is linked to strong correlations between axes, motivating covariance standardization with DBN.
- Feature Decorrelation: Shuffled-DBN adds an easy decorrelation mechanism to DBN for empirical comparison.
3.1. Preliminary
The concise framework trains an encoder on pairs of augmented views using a similarity or distance objective. The encoder consists of a backbone and projector, and the framework includes common components of modern self-supervised approaches.
- The framework contains only the most common components of modern self-supervised approaches.
- Training optimizes an expected objective over data samples and two independently augmented views.
- The encoder fθ contains a backbone and projector, while x1 and x2 are transformed views and ℓ measures similarity or distance.
- Unless otherwise specified, the distance function is squared error.
3.2. Reachable Collapse Patterns and Their Indicators
The concise framework exhibits both complete collapse and dimensional collapse, which are associated with vanishing variances and strong correlations, respectively. Standardizing variance resolves complete collapse, while covariance standardization improves the dimensionally collapsed representation.
- Complete collapse: The baseline reaches complete collapse, producing a single-point projection space and only 28.56% linear-evaluation accuracy.Complete collapse is associated with vanishing variances and leaves almost no gradient to influence the learned representation.
- Complete collapse: Batch Normalization can mitigate complete collapse by standardizing feature variance.The paper links complete collapse to vanishing variances and motivates variance standardization as the corresponding remedy.
- Complete collapse: Adding a variance-standardizing BN layer raises linear-evaluation accuracy from 28.56% to 69.52%.The added BN layer has no affine parameter and uses ϵ = 0 in the reported realization.
- Dimensional collapse: The BN-treated representation still undergoes dimensional collapse, concentrating projected features on a low-dimensional manifold such as a single line.Dimensional collapse can harm utility and is associated by definition with strong correlations between axes.
Definition 3 (Decorrelated Batch Normalization [25])
Decorrelated Batch Normalization divides feature dimensions into groups and applies ZCA whitening within each group to standardize covariance. In the reported comparison, addressing collapse patterns improves representation utility, with larger gains at higher projection dimensions.
- Definition 3 (Decorrelated Batch Normalization [25]): DBN takes batches of D-dimensional vectors and applies ZCA whitening independently to feature groups of size G.The grouping strategy divides the D feature dimensions into groups of size G before whitening each group.
- Definition 3 (Decorrelated Batch Normalization [25]): ZCA whitening produces zero-mean outputs with identity covariance when the input covariance matrix is full-rank.The construction uses the covariance matrix's eigenvalues and orthonormal eigenvectors to whiten the features.
- Definition 3 (Decorrelated Batch Normalization [25]): DBN standardizes covariance matrices within dimension groups to alleviate dimensional-collapse issues.Because whitening is grouped, the operation targets covariances among dimensions in each group.
- Empirical effect: 72.45% linear-evaluation accuracy with DBN exceeds the 69.52% obtained from the dimensionally collapsed representation.The comparison uses a two-dimensional projection visualization with group size G = 2.
- Empirical effect: In 128-dimensional projections, vanilla training shows vanishing standard deviations while BN shows strong correlations, indicating complete and dimensional collapse, respectively.The reported comparison uses a two-layer MLP with 128 hidden neurons and 128-dimensional outputs.
- Empirical effect: Utility gains from addressing collapse patterns increase with projection dimension, supporting the potential importance of feature decorrelation.The figure caption states that collapsed variants can optimize the loss easily while yielding representations with degraded utility.
3.3. Further Decorrelation, Further Gains
DBN only decorrelates dimensions within fixed groups, so dimensional collapse can remain partially unresolved. Shuffled-DBN randomizes feature assignments before grouped whitening and reverses the permutation afterward to provide further decorrelation.
- Motivation: DBN can leave dimensional collapse partially unresolved because it standardizes covariances only within fixed dimension groups.This limitation motivates a variant that changes the grouping across applications.
- Method: Shuffled-DBN randomly permutes feature dimensions, applies DBN with group size G, and reverses the permutation for the output.The permutation is applied before DBN and its inverse restores the original feature ordering afterward.
- Method: Random regrouping makes each dimension whiten with G − 1 randomly chosen dimensions rather than a fixed set.The paper presents this as a way to standardize the covariance matrix better.
- Results: Shuffled-DBN yields lower correlation strength, less dimensional collapse, and better utility than the corresponding DBN approach.These empirical observations support the paper's claim that Shuffled-DBN provides further decorrelation.
3.4. Details that Matter
This section explains objective and normalization choices that affect feature decorrelation, collapse mitigation, and utility. It also motivates grouping in DBN through rank and computational constraints.
- The Choice of The Objective ℓ: Squared error is used as the default objective instead of cosine similarity because cosine normalization can introduce interference across projection dimensions.With squared error, each gradient component involves only the corresponding dimension; experimentally, Shuffled-DBN fails to decorrelate with cosine similarity.
- The Choice of The Objective ℓ: Shuffled-DBN achieves lower correlation, higher estimated rank, and higher accuracy than DBN under the same group size and projector configuration.Here, corr. measures average absolute off-diagonal correlation, rank estimates the span of projected features using singular values, and acc. measures utility.
- The Choice of The Objective ℓ: Unit-norm normalization may conflict with feature decorrelation, while squared error reduces to a cosine-based expression only when both vectors have unit L2-norm.The paper reports experimentally that Shuffled-DBN fails in decorrelation with cosine similarity as the objective.
- The Detailed Setup of BN: Learnable affine scaling in BN can nullify variance standardization because output variance scales linearly with γ.The paper reports that both learnable affine transformation and non-negligible ϵ are detrimental to eventual utility.
- The Detailed Setup of BN: For ϵ > 0, BN output variance increases monotonically with input variance, so vanishing variance remains a reachable trivial solution.The reported variance is σ_d^2/(σ_d^2 + ϵ) = 1 − ϵ/(σ_d^2 + ϵ).
- The Role of Grouping in DBN: ZCA whitening requires a full-rank covariance matrix, forcing B ≥ D + 1, whereas grouping requires batch size to scale only with group size G.Grouping also reduces one-pass cost from O(BD^2) to O(BDG).
4. Evalution
Experiments across multiple benchmarks show that further feature decorrelation improves representation utility, generalization, and robustness to projection-space and decorrelation-strength choices. Shuffled-DBN generally outperforms DBN, though ImageNet performance remains below state of the art.
- Shuffled-DBN outperforms DBN on all 4 benchmarks and remains competitive with the best evaluated methods.The comparison uses linear evaluation after 200-epoch pretraining.
- Shuffled-DBN generalizes better than DBN from Tiny ImageNet to both CIFAR-10 and CIFAR-100.The representations are pretrained on Tiny ImageNet and evaluated with linear classifiers on the two target datasets.
- Further decorrelation yields gains consistently across different projector output dimensions.The result is reported for CIFAR-10 with varying projection-space dimensionality.
- Utility improves with stronger decorrelation in both kNN classification and linear evaluation.In Shuffled-DBN, stronger decorrelation is obtained by increasing the group size G.
- On ImageNet, Shuffled-DBN does not achieve state-of-the-art performance but remains promising in the concise framework.The framework uses no predictor, momentum encoder, or other special implementation detail.
5. Conclusion
The paper concludes that self-supervised learning can reach both complete and dimensional collapse, and that feature decorrelation addresses the latter while improving utility. It frames these findings as motivation for diagnosing design flaws rather than relying mainly on trial and error.
- Complete collapse is verified and addressed by standardizing variance.
- Dimensional collapse is an overlooked but reachable failure pattern in self-supervised representation learning.
- The connection between dimensional collapse and strong axis correlations motivates feature decorrelation through covariance standardization.
- The authors present feature decorrelation as a basis for advancing self-supervised learning by addressing design flaws instead of relying mostly on trial and error.