AdaDim: Dimensionality Adaptation for SSL Representational Dynamics

A key factor in effective Self-Supervised Learning (SSL) is preventing dimensional collapse, where higher-dimensional representation spaces span a lower-dimensional subspace. Therefore, SSL optimization strategies involve guiding a model to produce a representation space with a higher dimensionality through objectives that encourage feature decorrelation or sample uniformity. A higher dimensionality indicates that the representation space has greater feature diversity which is useful for generalization to downstream tasks. In addition to dimensionality optimization, SSL algorithms also utilize a projection head that maps the representation space to an embedding space where the optimization objective is applied. Recent work has characterized the projection head as a filter of noisy or irrelevant features by reducing the mutual information between the representation and embedding spaces. Therefore, the current literature views a good SSL representation space as one with a high dimensionality and a low mutual information between each respective space. However, we directly oppose this perspective and instead assert that the best performing SSL model arrives at an ideal balance between both dimensionality and mutual information. This conclusion is supported through a novel training dynamics analysis that reveals that increases in dimensionality due to feature decorrelation at the start of training lead to correspondingly higher mutual information, while increases in dimensionality due to samples distributing uniformly in a high-dimensional space at the end of training cause mutual information to plateau or decrease. To leverage these dynamics, we introduce AdaDim, a method that adaptively balances between dimensionality and mutual information by identifying situations where both families of SSL approaches are better utilized. We show that AdaDim results in performance exceeding common SSL baselines by as much as 3% accuracy on common SSL benchmark tasks.