Whitening Consistently Improves Self-Supervised Learning

TL;DR

This paper demonstrates that incorporating ZCA whitening as the final layer in self-supervised learning models consistently enhances feature quality and improves downstream task performance across various methods and architectures.

Contribution

It introduces the universal application of whitening as a final layer in SSL models, independent of the specific SSL method or encoder architecture.

Findings

Whitening improves linear and k-NN accuracy by 1-5%.

Whitening is effective across multiple SSL methods and architectures.

Proposes metrics for analyzing learned feature quality.

Abstract

Self-supervised learning (SSL) has been shown to be a powerful approach for learning visual representations. In this study, we propose incorporating ZCA whitening as the final layer of the encoder in self-supervised learning to enhance the quality of learned features by normalizing and decorrelating them. Although whitening has been utilized in SSL in previous works, its potential to universally improve any SSL model has not been explored. We demonstrate that adding whitening as the last layer of SSL pretrained encoders is independent of the self-supervised learning method and encoder architecture, thus it improves performance for a wide range of SSL methods across multiple encoder architectures and datasets. Our experiments show that whitening is capable of improving linear and k-NN probing accuracy by 1-5%. Additionally, we propose metrics that allow for a comprehensive analysis of the learned features, provide insights into the quality of the representations and help identify collapse patterns.