Understanding Dimensional Collapse in Contrastive Self-supervised Learning (2110.09348v3)

Abstract: 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.

• The paper identifies that both strong augmentation and implicit regularization drive dimensional collapse in contrastive learning.
• The study introduces DirectCLR, a novel method that applies contrastive loss directly to sub-vectors, eliminating the need for an explicit trainable projector.
• Empirical tests on ImageNet show that DirectCLR consistently outperforms SimCLR, offering improved accuracy and efficiency in representation learning.

Understanding Dimensional Collapse in Contrastive Self-Supervised Learning

Introduction

The phenomenon of embedding vectors occupying a lower-dimensional subspace, termed as dimensional collapse, presents a critical challenge in the field of contrastive self-supervised learning. Despite the efficacy of contrastive learning methods in leveraging negative sample pairs to mitigate complete collapse, they remain susceptible to dimensional collapse. This paper explores the dual mechanisms underlying dimensional collapse in contrastive learning and introduces DirectCLR, a novel contrastive learning approach designed to directly optimize the representation space without the necessity for an explicit trainable projector. Extensive experiments on the ImageNet dataset demonstrate DirectCLR's superiority over SimCLR when utilizing a linear trainable projector.

Mechanisms of Dimensional Collapse

Strong Augmentation

Dimensional collapse can be attributed to excessively strong augmentations that surpass the information present in input data. In the context of a simplified linear network, when the augmentation's magnitude eclipses the variance intrinsic to the data distribution, the weights of the network tend to collapse. The analysis within this paper establishes that strong augmentation triggers the embedding space covariance matrix to adopt a low-rank structure, indicative of dimensional collapse.

Implicit Regularization

Another contributing factor to dimensional collapse is implicit regularization, a phenomenon observed even without the presence of strong augmentation. This effect becomes pronounced in networks exceeding a single layer, driving the model towards low-rank solutions due to the interplay between weight matrices across different layers. Intriguingly, the paper illustrates how over-parametrization in deep networks fosters a more pronounced dimensional collapse through implicit regularization.

DirectCLR: A Novel Learning Method

Addressing the issue of dimensional collapse, the paper proposes DirectCLR, a method that eschews the use of an explicit trainable projector, a common fixture in many contrastive learning frameworks. DirectCLR innovates by applying the contrastive loss directly to a sub-vector of the representation vector, effectively preventing dimensional collapse while optimizing the representation space. Empirical validation on ImageNet reveals that DirectCLR outperforms a traditional SimCLR setup with a linear projector, achieving higher accuracy metrics.

Empirical Insights and Theoretical Contributions

The research meticulously verifies its proposed mechanisms through a series of controlled experiments, offering substantial theoretical insights into the dynamics of dimensional collapse. Findings concerning the role of strong augmentation and implicit regularization in causing dimensional collapse significantly advance the understanding of self-supervised learning models. Furthermore, the introduction of DirectCLR sheds light on the potential redundancy of projectors in certain self-supervised learning configurations.

Future Outlook and Implications

The identified mechanisms of dimensional collapse highlight the importance of careful augmentation strategies and model architecture design in self-supervised learning. The DirectCLR approach, remarkable for its simplicity and effectiveness, prompts a reevaluation of conventional practices in contrastive learning. Prospects for future research include exploring the applicability of DirectCLR's principles to a broader array of self-supervised tasks and delving deeper into the theoretical underpinnings of dimensional collapse in more complex, non-linear settings.

In conclusion, this paper's exploration into the roots of dimensional collapse in contrastive self-supervised learning and the introduction of DirectCLR mark significant advancements in the pursuit of more robust and efficient representation learning methods. The findings not only offer practical solutions to prevailing challenges but also prompt a rethinking of foundational aspects of contrastive learning methodologies.