Projector’s last layer controls representation quality in contrastive learning

On the Role of the Projector in Contrastive Self-Supervised Learning: Last-Layer Rank Dynamics Drive Representation Quality

Computer Vision and Pattern Recognition

Summary

Self-supervised learning methods can suffer from a problem where their internal data representations lose diversity, making them less useful. The authors studied a part of these methods called the projector, a small neural network that helps prevent this issue. They found that changes in the last layer of either the projector or the base model have the biggest impact on preserving useful information. By focusing on regularizing just the last layer’s weights, they improved the quality of learned representations across several datasets.

What this means in practice

  • For machine learning engineers: Improve image classification models by applying last-layer weight regularization to boost representation quality in contrastive pre-training.
  • For computer vision developers: Enhance self-supervised feature extractors used in visual recognition by focusing on the projector’s last layer, raising accuracy on various datasets.

Authors

Siladittya Manna, Priyangshu Mandal, Umapada Pal, Saumik Bhattacharya

Abstract

The dimensional collapse of representations in self-supervised contrastive learning is an ever-present issue. One notable technique to prevent such a collapse of representations is using a multi-layered perceptron network called Projector. In several works, the projector has been found to heavily influence the quality of representations learned in a self-supervised contrastive pre-training task. However, the question still lingers. What role does the projector play? Assuming the projector mitigates dimensional collapse, what prevents the terminal layer of the base encoder from functioning as the projector in the absence of an explicit multi-layer perceptron (MLP) head? In this work, we intend to study what happens inside the projector by examining the rank dynamics of the same and the encoder through empirical study and analysis. Through mathematical analysis, we observe that the effect of rank reduction predominantly occurs in the last layer. Motivated by this insight, we propose a weight regularization strategy applied specifically to the last layer. We demonstrate that this targeted approach yields better performance than applying orthogonal weight regularization across the entire network (WeRank), both with and without a projector. Our method improves Top-1 accuracy by more than 1% on SimCLR on the ImageNet100 dataset and consistently outperforms baseline SimCLR variants on CIFAR datasets, supporting our interpretation of the projector's role.