Modulate Your Spectrum in Self-Supervised Learning

TL;DR

Spectral Transformation (ST) is introduced, a framework to modulate the spectrum of embedding and to seek for functions beyond whitening that can avoid dimensional collapse, and its empirical investigations unveil other ST instances capable of preventing collapse.

Abstract

Whitening loss offers a theoretical guarantee against feature collapse in self-supervised learning (SSL) with joint embedding architectures. Typically, it involves a hard whitening approach, transforming the embedding and applying loss to the whitened output. In this work, we introduce Spectral Transformation (ST), a framework to modulate the spectrum of embedding and to seek for functions beyond whitening that can avoid dimensional collapse. We show that whitening is a special instance of ST by definition, and our empirical investigations unveil other ST instances capable of preventing collapse. Additionally, we propose a novel ST instance named IterNorm with trace loss (INTL). Theoretical analysis confirms INTL's efficacy in preventing collapse and modulating the spectrum of embedding toward equal-eigenvalues during optimization. Our experiments on ImageNet classification and COCO object detection demonstrate INTL's potential in learning superior representations. The code is available at https://github.com/winci-ai/INTL.

Figures (9)

• Figure 1: The framework using spectral transformation (ST) to modulate the spectrum of embedding in joint embedding architecture for SSL.
• Figure 2: Investigate ST using power functions. We choose several $p$ from $0$ to $1.5$. We show (a) the visualization of the toy model output; (b) top-1 and 5-nearest neighbors (5-nn) accuracy on CIFAR-10; (c) condition indicator of embedding $\mathbf{Z}_{}$ and transformed output $\widehat{\mathbf{Z}}_{}$ on CIFAR-10. We use the inverse of the condition number (IoC) in logarithmic scale with base 10 ( $lg{IoC} = lg{c^{-1}} = lg\frac{\lambda_d}{\lambda_1}$) as the condition indicator. The results on CIFAR-10 are obtained through training with ResNet-18 for 200 epochs and averaged over five runs, with standard deviation shown as error bars. We show the details of experimental setup in Appendix\ref{['sec:analytical_experiments']}. Similar phenomena can be observed when using other datasets (e.g., ImageNet) and other networks (e.g., ResNet-50).
• Figure 3: Investigate the effectiveness of IterNorm with and without trace loss. We train the models on CIFAR-10 with ResNet-18 for 100 epochs. We apply IterNorm with various iteration numbers $T$, and show the results with (solid lines) and without (dashed lines) trace loss respectively. (a) The spectrum of the embedding $\mathbf{Z}_{}$; (b) The spectrum of the transformed output $\widehat{\mathbf{Z}}_{}$; (c) The top-1 accuracy. (d) indicates that IterNorm (without trace loss) suffers from numeric divergence when using a large iteration number, e.g. $T=9$. It is noteworthy that when $T\geq11$, the loss values are all NAN, making the model unable to be trained. Similar phenomena can be observed when using other datasets (e.g., ImageNet) and other networks (e.g., ResNet-50).
• Figure 4: Algorithm of INTL, PyTorch-style Pseudocode.
• Figure 5: Visualization of our synthetic 2D dataset. We show (a) the distribution of our 2D dataset; (b) the initial output of the toy Siamese network.

Theorems & Definitions (9)

• Definition 1
• Theorem 1
• Proposition 1
• Proposition 2
• Theorem 2
• proof
• proof
• proof
• proof