Towards Demystifying Representation Learning with Non-contrastive Self-supervision

Abstract

Non-contrastive methods of self-supervised learning (such as BYOL and SimSiam) learn representations by minimizing the distance between two views of the same image. These approaches have achieved remarkable performance in practice, but it is not well understood 1) why these methods do not collapse to the trivial solutions and 2) how the representation is learned. Tian et al. ( 2021) made an initial attempt on the first question and proposed DirectPred that sets the predictor directly. In our work, we analyze a generalized version of DirectPred, called DirectSet(α). We show that in a simple linear network, DirectSet(α) provably learns a desirable projection matrix and also reduces the sample complexity on downstream tasks. Our analysis suggests that weight decay acts as an implicit threshold that discard the features with high variance under augmentation, and keep the features with low variance. Inspired by our theory, we simplify DirectPred by removing the expensive eigen-decomposition step. On CIFAR-10, CIFAR-100, STL-10 and ImageNet, DirectCopy, our simpler and more computationally efficient algorithm, rivals or even outperforms DirectPred.

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

“…In this section, we will clarify the differences between our results and some similar studies. Especially the theoretical papers by Tian et al [82] and Wang et al [87]. Pokle et al [67] compared the landscapes between contrastive and non-contrastive learning and points out the existence of non-collapsed bad minima for non-contrastive learning without a prediction head.…”

“…In the paper [82], experiments over the STL-10 dataset showed that the linear prediction head tends to converge to a symmetric matrix during training. And the follow-up paper [87] established a theory under the symmetric prediction head (which is not trained but manually set at each iteration). However, similar to the reason why eigenspace alignment cannot fully explain the effects of the prediction head, the symmetric prediction head given in [87] might not explain the trainable prediction head as well.…”

“…And the follow-up paper [87] established a theory under the symmetric prediction head (which is not trained but manually set at each iteration). However, similar to the reason why eigenspace alignment cannot fully explain the effects of the prediction head, the symmetric prediction head given in [87] might not explain the trainable prediction head as well. Under their linear network setting, where W is the weight matrix of the base encoder, they manually set the prediction head W p at iteration t to be…”

“…There are already some theoretical papers [82,87,67] that try to address similar questions. While none of these papers studied the training process of the prediction head, our results provide a completely different perspective: We explain why training the prediction head can encourage the network to learn diversified features and avoid dimensional collapses, even when the trivial collapsed optima still exist in the training objective, which is not covered by the prior works.…”

“…This proves the trainable prediction head does not need to behave like a symmetric matrix during most of the training process. Therefore the theories based on symmetric prediction head [82,87] cannot fully explain the behaviors of the trainable prediction head. See Figure 3 and Figure 4.…”