Implicit variance regularization in non-contrastive SSL

Summary

This article builds on the theory of non-constrastive (nc) self-supervised learning (SSL), with methods such as BYOL or SimSiam. Contrary to the existing litterature, they study directly the cosine similarity loss used rather than a Euclidian loss on the eigenspace, using NTK dynamics. They show that collapse is avoided in this case with implicit variance reduction. With their analysis, they find an undesirable anisotropy which they fix with a new family of losses that outperforms the state-of-the-art, especially without needing an Exponiential Moving Average (EMA) target network.

Strengths

The theory provides strong results in the linear framework, explaining the eigenvalues dynamics of both euclidian and cosine losses with and without projectors and stop gradients. The use of the NTK dynamics and the analysis of the cosine similarity is novel to me. The findings that the eigenvalues affect learning as a learning rate provide a good justification for the new IsoLosses, which gives state-of-the-art results. In particular, the removal of the EMA and the increased dimensionality of the losses ensures their consistency.

Weaknesses

One of the main claims of the article is the analysis of the cosine similarity loss. However in that case, if the dynamics are quite different as claimed by the authors, I find the use of the result of Tian et al. that the predictor eigenspace aligns with the one of the correlation matrix of the representations surprising since it was established with a Euclidian loss. I would have hoped atleast a comment on this, if not an analysis of the alignment as done in Tian et al. I was awaiting an analysis of the eigenvalues similar to Figure 2 for a real network, and not only a comparison with other losses. As such, the analysis of the IsoLoss feels a bit disconnected from the theoretical analysis. Why do the dimensionality increase? Is it because small eigenvalues do not collapse due to their low learning rates ? Do the eigenvalues evolve during training similarly to the linear case, for BYOL/SimSiam and IsoLoss ? (and their euclidian counterparts) Two important notions in the theory of NC-SSL are the EMA and the weight decay. Since an important improvement of the method is the removal of the need of the EMA of the target network, I would find natural to mention the effect of the EMA in the dynamic. Similarly, the results of Wang et al. indicate the importance of the role of weight decay to collapse unwanted eigenvalues. I am surprised that the role of weight decay is not mentioned anywhere, especially despite its use in the implementation. The formulation of the Isoloss (in particular the cosine similarity one) is very hard to interpret. However, it is logical due to the way it was found. This is an interesting theoretical article despite some flaws and questions raised, and I am ready to raise my score if my questions are answered.

Questions

If would define $\hat z = U z$ a bit more clearly maybe before Lemma 1. Equation 3: There is a dimensional issue with the diagonal matrix on the right with the product "$zD$" if I am not mistaken. Shouldn't $D$ be on the left? Definition 1: I find the different notations often a bit unclear, but most importantly is the sudden usage of the time $t$ index which is not consistent. It is present at line 114 but not at line 111, and appears only on part of the variables. Similarly for lemma 2, where $\hat z_t$ loses the time index again. The use of bold variables is also often inconsistent. Table 2: I would have expected the inclusion of DirectCopy and of SimSiam in the comparison.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

The authors address adequately the limits of the theory they established, although they do not talk about the effect of EMA and weight decay in their theory or their hypothesis of DirectPred which was established with an Euclidian loss.

Summary

The learning dynamics of non-contrastive self-supervised learning is an important problem to understand how these methods avoid collapse without using negative samples. In this paper, the authors provide a rigorous analysis of this problem on a simple linear network with Gaussian inputs, especially the difference between Euclidean and cosine loss functions. It is demonstrated that both losses have implicit regularization effects on the variance of representation, and the role of predictor and stop-grad operations are thoroughly investigated. Based on these insights, the authors propose an isotropic loss to equalize the convergence rate and lead to a better performance in various settings.

Strengths

The theoretical analysis in this paper is novel and clearly improves the understanding of how loss function, stop grad and projector affect the performance of non-contrastive self-supervise learning. Based on the mathematical understanding, the authors propose a new loss function that could beat the baseline method on different datasets, which illustrates the validity and power of the proposed theory in practice.

Weaknesses

The major weakness of this paper is the strong assumptions that require the network to be linear and the input to be iid Gaussian. For real-world datasets and practical network structures, it is not clear if those insights are still valid.

Questions

1. In Table 1, it is predicted that $L_{noPred}$ converges slower than $L$, however, in Figure 2 there is an eigenvalue in $L$ that converge much slower, there seems to be a discrepancy between theory and experiments. 2. In theorem 1 and 2, the dynamics are taken to be the expectation of original dynamics over the distribution of augmentation. It seems that after taking expectations, it is essentially equivalent to inputting the original images without augmentations, which may lose an essential part of contrastive learning. 3. In Theorem 2, it is assumed that the eigenvalues are of comparable magnitude, is this assumption required only on the initial value, or on the whole trajectory? In the latter case, it seems not appropriate to draw conclusions on the behavior of eigenvalues.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

4 excellent

Limitations

No certain limits.

Summary

This work analyzes the learning dynamics of non-contrastive SSL approaches such as BYOL and Simsiam. Based on the proposed theory, the authors analyze the how the stop-grad and predictor module affect the learning dynamics. Importantly, the authors design a theoretically inspired loss and gain improvement on classification tasks.

Strengths

1. This work provides solid analysis and discussion on the learning dynamics of non-contrastive SSL algorithm, which contributes to our understanding of the algorithm. Especially, they consider the widely used cosine similarity loss, while existing works only discuss Euclidean loss. They also propose an isotropic loss which is inspired by the theory. 2. The simulation experiments helps understanding the proposed theory.

Weaknesses

1. Lack of basic introduction to the used techniques, e.g., the neural tangent kernal (NTK). 2. Lack of awareness of dividing the article into paragraphs. The preliminary, theoretical analysis and the simulation experiment are all mixed in chapter 2, which might be confusing for readers. 3. The accuracy for the baseline model (BYOL) is too low. In the well-known SSL repository, solo-learn [1], the top-1 accuracy of BYOL (R-18, 1000ep) is 92.58 on CIFAR-10, and 70.46 on CIFAR-100. However, in this paper the baseline is only 89.4 and 61.1. I am afraid that this experiment cannot justify the effectiveness of the proposed isotropic loss. 4. This paper mainly provides theory to understand how non-contrastive SSL algorithms work (4 pages), and the design of the isotropic loss is inspired by the theory (only half a page). I think the title of this paper is not accurate enough.

Questions

1. What does "Gaussian i.i.d inputs" in line 133 refer to specifically? To what extent does this setting fit the real-world scenario? Others: Line 68: Replace z^{(1/2)} with z^{(1)}, z^{(2)} to avoid ambiguity Line 146: typo: "lammata" $\ell_{iso}$ first appears in Fig 2. and Tab 1., which is before where it is defined. If the author solves my concerns, then I would be willing to raise the rating. [1] https://github.com/vturrisi/solo-learn

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

See weaknesses

Good Rebuttal

The authors did a good job at the rebuttal. Most of my previous concerns are resolved. They reproduce the experiments to match SOTA performance, and adjust their claims accordingly. This could help deliver the right message to readers (as far as I know, there is still no linear predictor on par with nonlinear ones, and the authors' new results again confirm this observation). Second, the authors also show that their advantage over DirectPred in the new setting, which I thought, would be enough for an empirical justification. I will further go through other reviewers' comments. I would be happy to raise the score when the button is available. One thing that came to my mind is that the theory from a recent work [1] might also be related, because it could also explain the effectiveness of IsoLoss. As far as I see, in their framework, the IsoLoss serves as a high-pass filter (the filter function is 2-\lambda, roughly speaking), which often outperforms the low-pass one like DirectPred (Table 1 & 2 [1]). A discussion of their relationship would be helpful. [1] Zhuo et al. Towards a Unified Theoretical Understanding of Non-contrastive Learning via Rank Differential Mechanism. ICLR. 2023.