The Mechanism of Prediction Head in Non-contrastive Self-supervised Learning

Abstract

Recently the surprising discovery of Bootstrap Your Own Latent (BYOL) method by Grill et al.shows the negative term in contrastive loss can be removed if we add the so-called prediction head to the network architecture, which breaks the symmetry between the positive pairs. This initiated the research of non-contrastive self-supervised learning. It is mysterious why even when trivial collapsed global optimal solutions exist, neural networks trained by (stochastic) gradient descent can still learn competitive representations and avoid collapsed solutions. This phenomenon is one of the most typical examples of implicit bias in deep learning optimization, and its underlying mechanism remains little understood to this day.In this work, we present our empirical and theoretical discoveries about the mechanism of prediction head in non-contrastive self-supervised learning methods. Empirically, we find that when the prediction head is initialized as an identity matrix with only its off-diagonal entries being trained, the network can learn competitive representations even though the trivial optima still exist in the training objective. Moreover, we observe a consistent rise and fall trajectory of off-diagonal entries during training. Our evidence suggests that understanding the identity-initialized prediction head is a good starting point for understanding the mechanism of the trainable prediction head.Theoretically, we present a framework to understand the behavior of the trainable, but identity-initialized prediction head. Under a simple setting, we characterized the substitution effect and acceleration effect of the prediction head during the training process. The substitution effect happens when learning the stronger features in some neurons can substitute for learning these features in other neurons through updating the prediction head. And the acceleration effect happens when the substituted features can accelerate the learning of other weaker features to prevent them from being ignored. These two effects together enable the neural networks to learn all the features rather than focus only on learning the stronger features, which is likely the cause of the dimensional collapse phenomenon. To the best of our knowledge, this is also the first end-to-end optimization guarantee for non-contrastive methods using nonlinear neural networks with a trainable prediction head and normalization.

ContrasInver: Voxel-wise Contrastive Semi-supervised Learning for Seismic Inversion

“…In this regard LID is characterized by contrastive learning, and this feature consistency approach is important in that it also plays a semi-supervised role in ConstraInver, hence it is called contrastive semi-supervised learning. We experimentally demonstrate that the performance of projected feature consistency is significantly higher than that of unprojected and label consistency, which we hypothesize is due to the filtering effect of the projection head on meaningful information [66], [67]. We include at least one sample containing real wells in each batch to ensure that the learning process does not deviate from the right track.…”

Augmentation Invariant Manifold Learning

“…First, the new loss function keeps the local similarity between the negative pair of data (X i , X i − ), while the existing self-supervised learning methods either ignore the negative pair (in non-contrastive methods) or push the negative pair far from each other despite their Euclidean distance (in contrastive methods). Second, the regularization part in (5) can help avoid dimensional collapse observed in some self-supervised learning methods (Hua et al, 2021;Wen and Li, 2022).…”

“…Since the introduction of self-supervised learning, recent works have tried to theoretically understand why the learned data representation from augmented data can help improve the downstream analysis (Arora et al, 2019;Tsai et al, 2020;Wei et al, 2020;Tian et al, 2020b;Tosh et al, 2021;Wen and Li, 2021;HaoChen et al, 2021;Wang, 2022;Wen and Li, 2022;Balestriero and LeCun, 2022). Most of these theoretical works focus on the setting where the special conditional independence structure for the augmented data is assumed.…”

Towards a Unified Theoretical Understanding of Non-contrastive Learning via Rank Differential Mechanism

“…From an empirical side, Chen & He (2021) think that the predictor helps approximate the expectation over augmentations, and Zhang et al (2022a) take a center-residual decomposition of representations for analyzing the collapse. From a theoretical perspective, Tian et al (2021) analyze the dynamics of predictor weights under simple linear networks, and Wen & Li (2022) obtain optimization guarantees for two-layer nonlinear networks. These theoretical discussions often need strong assumptions on the data distribution (e.g., standard normal (Tian et al, 2021)) and augmentations (e.g., random masking (Wen & Li, 2022)).…”

“…From a theoretical perspective, Tian et al (2021) analyze the dynamics of predictor weights under simple linear networks, and Wen & Li (2022) obtain optimization guarantees for two-layer nonlinear networks. These theoretical discussions often need strong assumptions on the data distribution (e.g., standard normal (Tian et al, 2021)) and augmentations (e.g., random masking (Wen & Li, 2022)). Besides, their analyses are often problem-specific, which is hardly extendable to other non-contrastive variants without a predictor, e.g., DINO.…”