AdaDim: Dimensionality Adaptation for SSL Representational Dynamics

TL;DR

The paper addresses dimensional collapse in self-supervised learning by examining how the representation space dimensionality $H(R)$ and the projection-space mutual information $I(R;Z)$ co-evolve during training. It reveals that top-performing SSL models achieve a balance between $H(R)$ and $I(R;Z)$ rather than maximizing one term alone, and it introduces AdaDim, an adaptive training objective that modulates between dimension-contrastive and sample-contrastive signals while gradually regularizing $I(R;Z)$. The method computes an adaptive interpolation parameter $oldsymbol{ extalpha}$ from the current effective rank of $Z$, and applies a mutual-information regularizer with $oldsymbol{ extbeta}=oldsymbol{ extgamma}oldsymbol{ extalpha}$, resulting in a loss L_AdaDim that blends $(1-oldsymbol{ extalpha})L_{VICReg} + oldsymbol{ extalpha}L_{NCE}$ and a reg term. Empirically, AdaDim improves performance across diverse datasets without expensive training techniques, and the work provides a principled framework to tailor SSL objectives to representational dynamics, with broader implications for domain-specific SSL deployment.

Abstract

A key factor in effective Self-Supervised learning (SSL) is preventing dimensional collapse, where higher-dimensional representation spaces ($R$) span a lower-dimensional subspace. Therefore, SSL optimization strategies involve guiding a model to produce $R$ with a higher dimensionality ($H(R)$) through objectives that encourage decorrelation of features or sample uniformity in $R$. A higher $H(R)$ indicates that $R$ has greater feature diversity which is useful for generalization to downstream tasks. Alongside dimensionality optimization, SSL algorithms also utilize a projection head that maps $R$ into an embedding space $Z$. Recent work has characterized the projection head as a filter of noisy or irrelevant features from the SSL objective by reducing the mutual information $I(R;Z)$. Therefore, the current literature's view is that a good SSL representation space should have a high $H(R)$ and a low $I(R;Z)$. However, this view of SSL is lacking in terms of an understanding of the underlying training dynamics that influences the relationship between both terms. Our analysis shows that the best performing SSL models do not have the highest $H(R)$ nor the lowest $I(R;Z)$, but effectively arrive at a balance between both. To take advantage of this analysis, we introduce AdaDim, a training strategy that leverages SSL training dynamics by adaptively balancing between increasing $H(R)$ through feature decorrelation and sample uniformity as well as gradual regularization of $I(R;Z)$ as training progresses. We show performance improvements of up to 3% over common SSL baselines despite our method not utilizing expensive techniques such as queues, clustering, predictor networks, or student-teacher architectures.

Figures (23)

• Figure 1: a) This figure shows how performance varies for 20 different pre-trained ResNet-50 models as a function of $H(R)$ and $I(R;Z)$. b) The first three figures show how $H(R)$ and $I(R;Z)$ vary across training of a ResNet-18 encoder with SimCLR chen2020simple for 1000 epochs on three different datasets. c) This toy graphic shows how the representation space ($R$) and embedding space ($Z$) of a 3D dataset changes when following SSL training dynamics. We also demonstrate how these changes effect $H(R)$ and $I(R;Z)$.
• Figure 2: a) and b) show how $I(R;Z)$ in a gaussian setting changes as the number of features of $R$ is increased. c) and d) show how $I(R;Z)$ varies as the sample cluster variance increases.
• Figure 3: This is an analysis of 4 different SSL models $R$ and $Z$ space trained for 2000 epochs on Cifar-100 with ResNet-50. This analysis includes a) the number of eigenvalues above a threshold of $\tau = .01$, b) the cumulative explained variance ratio for top 30% of eigenvalues, c) the uniformity of each space, and d) $I(R;Z)$.
• Figure 4: In Figures a), b), and c), the $H(R)$ and $I(R;Z)$ across 15 ResNet-50 models trained with randomized hyperparameters with 3 different SSL strategies are shown. In Figure d), we show the same plot across 11 different SSL methods trained on ResNet-18 for 1000 epochs.
• Figure 5: This table shows the pearson correlation coefficient between the performance of a set of SSL models trained with different hyperparameters on a specific dataset and the effective rank ($H(R)$), $I(R;Z)$, and the ratio between them.