ICML2022

Exploring the Gap between Collapsed & Whitened Features in Self-Supervised Learning

Bobby He, Mete Ozay

31 citations

Abstract

Avoiding feature collapse, when a Neural Network (NN) encoder maps all inputs to a constant vector, is a shared implicit desideratum of various methodological advances in self-supervised learning (SSL). To that end, whitened features have been proposed as an explicit objective to ensure uncollapsed features (Zbontar et al., 2021; Ermolov et al., 2021; Hua et al., 2021; Bardes et al., 2022) . We identify power law behaviour in eigenvalue decay, parameterised by exponent β≥0, as a spectrum that bridges between the collapsed & whitened feature extremes. We provide theoretical & empirical evidence highlighting the factors in SSL, like projection layers & regularisation strength, that influence eigenvalue decay rate, & demonstrate that the degree of feature whitening affects generalisation, particularly in label scarce regimes. We use our insights to motivate a novel method, Post-hoc Manipulation of the Principal Axes & Trace (PostMan-Pat), which efficiently post-processes a pretrained encoder to enforce eigenvalue decay rate with power law exponent β, & find that PostMan-Pat delivers improved label efficiency and transferability across a range of SSL methods and encoder architectures.