A Cookbook of Self-Supervised Learning — part 6/8

Cleaned-text code points [77936, 93347). Bibliography omitted here; source.md includes it.

K-NN classifiers have the great advantage of not relying on many hyperparameters, being fast and light to deploy, without requiring any domain adaptation.

##### Linear

In the context of SSL evaluation, training a linear classifier on top of pre-trained feature representations, a.k.a. linear-probing evaluation, was introduced by Zhang et al. (2016) ; Zhang et al. (2017) . It is the most popular protocol for several reasons: it achieves high-accuracy, its performance heavily rely on the quality of the representation since their discriminative power is low, it imitates how the features can be used in practice, and last but not least, it is not very computationally expensive.

Most of the time, it is done simply by appending a linear layer at the end of the frozen backbone, and optimizing its parameters for a few epochs (around 100 100 ). Sometimes, as introduced by Bao et al. (2021b) , we can benefit from the fact that the linear evaluation is lightweight and evaluate multiple linear heads at the same time, to test many hyper-parameters at the same time (learning rate, averaging features or using a class-token for ViT-like architectures, number of features, etc.). A linear probe can also be trained online, by simply cutting gradient from the representations. Though only an approximation, an online linear probe is extremely cheap as it reuses the computations for the SSL pretraining, and gives a good indication of downstream performance, as shown in Figure 13 .

##### MLP

Instead of a simple linear probing, a multi-layer perceptron (with two or three layers) could also be used to extract which information is learned in a SSL model. Non linear evaluation is rarely present in work around SSL, however it is needed when the learned features are not linearly separable, and when it is too difficult to extract information present in features with a linear model. In fact, comparing results with a linear and a non linear probe, can give us some ideas about how well structured a representation is. Bordes et al. (2023a) present some results that compare different evaluation regime using a linear or a non linear probe. In Figure 13 , one can observe that it’s possible to get some gain in accuracy when using a multilater layer perceptron instead of a linear probe. However, the main issue with adding capacity into the probe is one related to overfitting: the best MLP head might not be the ones you get after 100 epochs, as showed in Figure 13 .

##### Full Fine-tuning

The Masked Auto-encoders (MAE) paper ( He et al., 2022 ) re-introduced fine-tuning as the main evaluation metrics. The main arguments are that linear-probing is uncorrelated with fine-tuning and transfer learning performances, and that small MLP heads do not evaluate the strength of the method to create strong but non-linear features. The majority of works that followed ( Bao et al., 2021b ; Zhou et al., 2022a ; Dong et al., 2021 ) focused on this type of evaluation (and sometimes do not report linear/MLP results). It has been shown that contrasting methods show inferior performance than masked image modeling with regards to fine-tuning, because they are less “optimization friendly” ( Wei et al., 2022 ) - which explains the overall interest over MIM. It is by far the most computationally expensive of the evaluation methods, since it needs to re-train the whole network. The most common benchmark on ImageNet runs the optimization over 100 100 epochs for ViT smaller than base, and for 50 50 epochs for larger models ( He et al., 2022 ) . Other works ( Bao et al., 2021b ; Peng et al., 2022 ; Wang et al., 2022b ) first fine-tune on ImageNet-21k for 60 60 epochs, and further fine-tune on ImageNet-1k, which represents between 1 / 5 1/5 to 2 2 times the cost of the pre-training phase.

#### 3.7.2 Evaluation without labels

As we just discussed, most evaluations rely on the use of labels and training an auxiliary model. This can make evaluations expensive and sensitive to hyperparameters or their optimizations. To help alleviate these issues multiple methods have been proposed to evaluate or help tune hyperparameters of methods without relying on labels. Using a pretext-task such as rotation prediction can facilitate performance evaluation without labels, as demonstrated in Reed et al. (2021) for data augmentation policy selection. However, a drawback of this approach is the requirement for training the classifier for the pretext-task and the assumption that rotations were not part of the pretraining augmentations, or the model would be invariant to it. The eigenspectrum of representations is used in conjunction with the loss value to evaluate performance in Li et al. (2022a) . While a correlation with performance is shown, it requires training a performance classifier with the rank and loss value, making it hard to use for unsupervised evaluation. In Agrawal et al. (2022) α \alpha -ReQ is introduced to evaluate methods by looking at the eigenspectrum decay of representations before the projector.

Another simple way to evaluate SSL methods, called RankMe, was introduced by Garrido et al. (2022a) . The idea is to use the effective rank of representations, defined as the entropy of the singular value distribution of the embeddings. It can be computed as: RankMe ( 𝒁 ) = exp ( − ∑ k = 1 min ⁡ ( N , K ) p k log p k ) , p k = σ k ( 𝒁 ) ‖ σ ⁡ ( 𝒁 ) ‖ 1 + ϵ \text{RankMe}({\bm{Z}})=\exp\left(-\sum_{k=1}^{\min(N,K)}p_{k}\log p_{k}\right),\;p_{k}=\frac{\sigma_{k}({\bm{Z}})}{\|\sigma({\bm{Z}})\|_{1}}+\epsilon (17) It is shown to be a necessary condition for good performance, though you can achieve full rank representations with degenerate results (e.g. a random matrix with entries sampled i.i.d. from a Gaussian distribution). While this cannot be used to evaluate different methods, it works well for hyperparameter selection, as shown in Table 3 .

#### 3.7.3 Going beyond classification

While classification is a commonly used performance metric for evaluating self-supervised learning models, it is important to consider other types of vision tasks as well. Tasks such as object detection and semantic segmentation have gained popularity as they require models to learn more complex representations of visual information. Recent works Caron et al. (2021) ; Zhou et al. (2022a) ; Bardes et al. (2022) have demonstrated the effectiveness of self-supervised learning for these tasks. However, a limitation is that there is currently no standardized protocol for evaluating self-supervised models on these tasks. Various evaluation methods exist, such as finetuning the encoder on a downstream task or using the encoder as a feature extractor. Further research is needed to establish a standardized evaluation protocol for these tasks in the context of self-supervised learning.

#### 3.7.4 Visual Evaluation

Another way to evaluate what information is contained or not in a representation is to use a decoder over the representation that is able to map back this information to pixel space. Some methods like ( He et al., 2022 ) are built with a specific decoder which make such visual analysis easy, however most SSL methods aren’t shipped with a decoder. To alleviate this issue and to allow researchers to visualize what can be learned by any type of SSL method, Bordes et al. (2022b) suggest training a conditional generative diffusion model using a SSL representation as conditioning. By analyzing which information remains constant across different generated samples using a specific conditioning and what information does not remain constant (because of the stochasticity in the generative model), one can get some hints about what information is contained in the representation. If a representation encodes every information about each pixel, the conditional generative model would exploit every bit of this information to perform a perfect reconstruction which will lead to no variance across different samples. If the representation encodes only the class information, the conditional generative model will only be able to use that to reconstruct the image belonging to this class, which means that when generating different samples, the object class will remain constant but the background/context/color would change across samples. In Figure 14 , we show how RCDM was used by Bordes et al. (2022b) to compare the representations learned at the projector level versus the representations learned at the backbone level. In this Figure, we observe that the representations at the projector level are much more invariant since the color/background information does not remain constant across different samples while this is not the case at the backbone level.

### 3.8 Speeding up Training

#### 3.8.1 Distributed Training

Training self-supervised models often requires large batch sizes ( Chen et al., 2020b ; He et al., 2020b ) , or can be considerably speed up by increasing the batch size, which is ultimately limited by the memory capacity of the device the model is trained on. Distributed training divides batches across several devices that run in parallel, which increases the overall size of the batch. This is mainly done with DDP: Distributed Data Parallel or FSDP: Fully Sharded Data Parallel, available in libraries like FairScale ( FairScale, 2021 ) or Apex ( NVidia, 2021 ) . However some self-supervised methods rely on the statistics of the current batch for the computation of their loss value ( Chen et al., 2020b ; Zbontar et al., 2021 ; Bardes et al., 2021 ) , which has to be taken into account when distributing the training across multiple devices. In this section, we present the elements that need to be taken into account in order to correctly distribute the training of common self-supervised learning methods. We call effective batch size, the size of the full batch distributed on the devices, and per device batch size, the size of each sub-batch on a single device.

Synchronized batch normalization. Batch normalization is one a the most common technique for stabilizing neural network training, as well as improving the performance of the network. It is present in most convolutional backbones used in self-supervised learning, in particular in ResNet. Batch norm uses the statistics from the current batch, which need to be aggregated for distributed training. This can be done easily in PyTorch by wrapping your distributed model the following way: model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model) This will replace all the BatchNorm modules in the network by a custom BatchNorm class that aggregates the statistics automatically.

Aggregate batches for exact loss computation. Batch norm is not the only operation that operate on batches, multiple self-supervised loss functions do as well, such as SimCLR ( Chen et al., 2020b ) that uses the examples in the current batch as negative example for its contrastive loss, or VICReg ( Bardes et al., 2021 ) that computes the covariance matrix of its embeddings. In these cases the batches from each device need to be aggregated into the full batch manually. This can be done using the all_gather operation from PyTorch, however this operation does not allow back-propagation through it. We therefore implement a custom gather operation that does, the code is provided below:

We use an all_reduce operation on the gradient, which sums them, because DDP will divide them later by the number of devices. One can use the operation by simply calling: FullGatherLayer.apply(x) on the input x . Practically, for the methods above, this needs to be done on the embeddings just before the computation of the loss.

Additional tricks. We advise to always use the effective batch size as argument to the training script, as well as for comparing runs. The DataLoader class takes the per device batch size as, argument, which can be obtained by dividing the effective batch size by the number of devices which is world_size in PyTorch. We also advise to use an adaptive learning rate scaled with the effective batch_size, for example using effective_lr = base_lr * effective_batch_size / 256 where is the base_lr is the argument of the training script. This reduce the learning rate search range when changing batch_size. When using small batch size, it is recommended by Chen et al. (2020b) to use effective_lr = base_lr * 𝙾𝙿𝙴𝙽 ( effective_batch_size ) \sqrt{(}\text{effective\_batch\_size}) / 256 .

#### 3.8.2 Even Faster Training with FFCV and Other Speedups

Since most join-embedding SSL methods requires different set of handcrafted data augmentation, data processing can become a real bottleneck when training SSL models. Some approaches 4 4 4 https://github.com/vturrisi/solo-learn have used DALI as an alternative data loader to pytorch vision while some other have relied on FFCV-SSL 5 5 5 https://github.com/facebookresearch/FFCV-SSL which is based on the FFCV library ( Leclerc et al., 2022 ) . FFCV-SSL ( Bordes et al., 2023a ) shows that one can train SimCLR on ImageNet in less than 2 days on a single GPU or in just a few hours using 8 GPUs (Figure 15 ).

#### 3.8.3 Speeding Up Training of Vision Transformers

Training ViT can be made more efficient for two reasons. First, it is made easy for ViTs not to process all patches. This is especially helpful when using masked prediction pre-training objectives such as MAE ( He et al., 2022 ) or Masked Siamese Networks ( Assran et al., 2022b ) . For instance, with ViT and such objectives, Data2vec 2.0 ( Baevski et al., 2022 ) achieves 84 % 84\% top-1 accuracy after pre-training for only 3 3 hours on 32 32 GPUs.

The second reason is linked to the architecture. Since transformers ( Vaswani et al., 2017 ) are employed in almost all domains of computer science, many works aim to reduce the compute and memory requirements of the attention mechanism. One approach is with low-rank and/or sparse approximation mechanisms ( Kitaev et al., 2020 ; Choromanski et al., 2020 ; Wang et al., 2020a ; Chen et al., 2021a ; Zaheer et al., 2020 ) . For instance, Li et al. (2022b) use sparse self-attention to improve efficiency in the context of SSL vision models. Another approach, is to resort to IO-aware optimizations ( Ivanov et al., 2021 ) , the most known one perhaps being FlashAttention ( Dao et al., 2022 ) .

These speed-ups are available in open-source libraries: Fairseq ( Ott et al., 2019 ) , FairScale ( FairScale, 2021 ) , XFormers ( Lefaudeux et al., 2022 ) , Apex ( NVidia, 2021 ) , etc.

Another simple way to speed up the training of vision transformers is to use Pytorch bfloat16 which allow faster training while keeping the same precision range as float32 (this is useful to avoid the usual numerical instability issues one can encounter when training vision transformers in float16).

## 4 Extending Self-Supervised Learning Beyond Images and Classification

### 4.1 Strategies for Other Data Domains

Pre-training large models with self-supervision objectives is popular not only for vision systems, but also for audio, text, and tabular data as well. The performance of existing SSL methods varies across these domains – yielding state-of-the-art language models but limited success on tabular data – which may either reflect better suitability of self-supervision or alternatively the wildly differing amount of attention which has been paid to the various domains in the SSL literature.