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

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

Despite their recent successes, there is an important limitation of SSL methods: poor performance on unbalanced datasets. Since real world data is imbalanced, such a limitation is an important factor that made the use of SSL methods on vast amount of uncurated data challenging. Assran et al. (2022a) explains such a limitation by the use of an hidden uniform prior that is common to many SSL methods. By distributing the data uniformly in the representation space, SSL methods learn to find the most discriminative features in a given mini batch. When data is uniformly distributed across classes labels, the most discriminative features that the model will learn will be class specific. However, when using imbalanced data, the most discriminative features inside the mini batch might not be the class anymore but more low level information which decrease the performances on downstream classification tasks. To alleviate this issue, Assran et al. (2022a) introduce the use of an additional regularization term on the SSL method MSN ( Assran et al., 2022c ) to change the distribution of the SSL clustering.

### 3.4 Teacher-Student Architecture Specific Tricks

#### 3.4.1 Role of the Moving Average Teacher

While the original BYOL method is based on exponential moving average (EMA) updates of the weights for the target (teacher) network, it was later confirmed that EMA is not necessary (i.e., the online and target networks can be identical). This is also confirmed with SimSiam ( Chen and He, 2021 ) , as long as the predictor is updated more often or has larger learning rate compared to the backbone. In the case of DQN, the target network with EMA is shown to remove bias Fan et al. (2020) and Piché et al. (2021) showed that the EMA could be removed from the target network by using the correct regularizer. For BYOL, a stop gradient of the online network, meaning the decay rate is 0 for the target network, collapses as shown in Table 5 of Grill et al. (2020) . Pham et al. (2022) shows the idea of exponential moving averages provide training stability that can even be used in non student-teacher frameworks such as SimCLR. Specifically, they show applying EMA updates to the projector of SimCLR can boost performance. Wang et al. (2022c) shows that training could also benefit from other kinds of asymmetries in the teacher-student setting (e.g., stronger augmentation on the student side).

#### 3.4.2 Role of the Predictor in Self-Labeling SSL

The predictor network plays a central role in BYOL’s success by predicting the representation of the teacher network from the student networks’ representation. Shi et al. (2020) shows removing the predictor leads to a performance drop from 68% to 21% top-1 accuracy on ImageNet (compared to the original two-layer MLP predictor in BYOL). In Figure 1 of Shi et al. (2020) , they demonstrate even a linear predictor leads to good performance and can recover from poor initialization in 10-20 epochs of training. For SimSiam, Table 1 of Chen and He (2021) shows removing the predictor in SimSiam also leads to collapse with a top-1 accuracy of < 1% on ImageNet. Tian et al. (2021) , whose implementation can be found 2 2 2 https://github.com/facebookresearch/luckmatters/tree/main/ssl , proves that in the presence of the predictor, the training dynamics of BYOL and SimSiam contains nontrivial stable fixed points, and thus avoid being trapped into trivial solutions during training, even if these trivial solutions are global optimal. It further proposed a contrastive method, DirectPred, that directly sets the predictor via eigenvalue decomposition during training and leads to comparable performance in ImageNet. Its follow-up work (DirectSet Wang et al. (2021b) ) further removes the overhead of eigenvalue decomposition.

### 3.5 Role of Standard Hyper-Parameters

A common issue in SSL research is that each method has different configurations of hyper-parameters. Hence comparisons directly between different SSL methods or models is often challenging. In this section, we present and describe the impact of each hyper-parameters to help SSL practitioners identify which are most important depending of their setup.

#### 3.5.1 Role of Mini-Batch Size

It was originally thought that contrastive methods such as SimCLR or MoCo require large batch sizes or memory banks to work. This turns out to be misleading as both methods can be made to work at small batch sizes. A square root scaling of the learning rate was discussed in the appendix of Chen et al. (2020b) which already gave a significant increase in performance of up to 5 points in top-1 accuracy on ImageNet for a 100 epochs training. Similarly, Bordes et al. (2023a) investigated the impact of the learning rate with small batch sizes and found how one can train SimCLR on ImageNet using a single gpu without an important drop in performances. Furthermore, some works such as DCL ( Yeh et al., 2021 ) show that you can reach top performance with a batch size of 256 or more for SimCLR, and a queue size of only 256 or more for MoCo, by simply removing the positive pair from the denominator of the softmax and with more careful hyperparameter tuning. Similarly, it was shown by Zhang et al. (2022a) that by decomposing the dictionary in MoCo and by using different temperatures for the positive and negative pairs it is possible to increase the robustness to the dictionary dimension.

#### 3.5.2 Role of Learning Rate (Schedulers) and Optimizers

Here we overview typical standard settings for learning rate schedulers and optimizers across methods. To determine the learning rate, methods often scale a base learning rate based on the batch size according to the heuristic by Goyal et al. (2017) : learning rate = batch size 256 ∗ base learning rate \frac{\text{batch size}}{256}*\text{base learning rate} . For ImageNet pretraining, VICREg, Barlow Twins, BYOL, and SimCLR use a base learning rate of 0.2 − 0.3 0.2-0.3 with the LARS optimizer ( You et al., 2017 ) . Additionally for some methods such as Barlow twins, a much smaller learning rate (0.0048) is used to update the bias terms and batch norm parameters. Other methods such as MAE, DINO, and iBot use the AdamW optimizer ( Loshchilov and Hutter, 2017 ) with a smaller base learning rate of 1 ​ e − 5 − 5 ​ e − 4 1e-5-5e-4 . For a discussion of weight decay see Section 3.5.3 . The most common training schedule involves a warmup period, usually 10 epochs, where the learning rate is linearly increased to its base value. After the warmup period, most methods use cosine decay.

#### 3.5.3 Role of Weight-Decay

Weight-decay is an important component of backprogagation for many SSL methods. Table 15 in BYOL ( Grill et al., 2020 ) indicates that no weight decay may lead to unstable results. A recent blogpost 3 3 3 https://generallyintelligent.ai/blog/2020-08-24-understanding-self-supervised-contrastive-learning/ also mentions using weight decay leads to stable learning in BYOL. In Figure 4 of Tian et al. (2020b) the effect of weight decay is explained in terms of its effect on memory of the initial conditions. The hypothesis is that weight decay allows the online network and predictor to better model invariance to augmentations regardless of the initial condition. For further reading, Zhang et al. (2022b) provides a good review of SimSIAM collapse understanding and Shi et al. (2020) does the same for BYOL.

#### 3.5.4 Vision Transformers Considerations

Training Vision Transformers (ViT) ( Dosovitskiy et al., ) requires special care. They are more prone to collapse and instability, and are more sensitive to the setting of hyper-parameters ( Touvron et al., 2021a ) .

Batch size. ( Chen et al., 2021b ) found that large batch (e.g., 4096) training for joint-embedding ViT SSL methods can be unstable. This instability does not reflect as a large drop in the final accuracy, but appears as drops in kNN probe accuracy during training when the L ∞ − n ​ o ​ r ​ m L_{\infty}-norm of the gradient spikes. Using a random (versus a learned) patch projection layer to embed pixel patches into input tokens for ViT stabilizes training for MoCo-V3, SimCLR, and BYOL and also improves the final accuracy. A learning rate warm-up period of 10k iterations ( Goyal et al., 2017 ; Dosovitskiy et al., ) also improves training stability. On the other hand, Caron et al. (2021) noted a drop in final k-NN accuracy when training with very small batch sizes (128). So, a batch size of 1024 or 2048 seems to be the sweet spot for SSL pre-training of ViTs.

While the ViT architecture does not have any BatchNorm layers, training a MoCo-V3 model with BN layers in the projector heads improved the linear probing accuracy of the ViT ( Chen et al., 2021b ) . Note that for joint embedding methods, batching can be done either together for all samples and crops in one batch, or separately for each batch of crops. SimCLR adopts the former, while BYOL and MoCo-V3 adopt the latter.

Patch size. ( Caron et al., 2021 ) found that training with smaller patch sizes ( 5 × 5 5\times 5 , or 8 × 8 8\times 8 instead of 16 × 16 16\times 16 ) leads to improved linear probing accuracy on DINO ViT pre-training. Note that while increasing patch sizes leads to a reduction in running time, it also increases memory usage (which makes it hard to train on patches smaller than 8 × 8 8\times 8 ).

Stochastic depth ( Huang et al., 2016 ) originated from NLP and was subsequently used in vision models ( Touvron et al., 2021b ) to train deeper models. It randomly drops blocks of the ViT as a regularization. The per-layer drop-rate may depend linearly on the layer depth or uniformly as suggested in recent works ( Touvron et al., 2021b ) . It has huge importance when training larger models (ViT-L, ViT-H, etc.). For instance Touvron et al. (2022) use 0.5 0.5 drop path rate for ViT-H models. Conversely, when training smaller models like ViT-B, such regularization usually hurts the performance ( Steiner et al., 2021 ) .

LayerDecay ( Clark et al., 2020 ) decreases the learning rate geometrically along the layers. Put differently, the last layer is not affected, while the first has very small learning rate. In SSL vision models, LayerDecay increases performance when fine-tuning on downstream tasks ( Bao et al., 2021b ; Zhou et al., 2022a ; He et al., 2022 ) . Depending on the model size, the parameter is set between 0.65 0.65 to 0.85 0.85 – larger models usually need higher values because there are more layers. The underlying principle is that SSL builds strong model backbones, therefore we only need to fine-tune the shallowest layers.

LayerScale ( Touvron et al., 2021b ) is a per-channel multiplication of the vector produced by each residual block of the transformer. It increases the stability of the optimization and permits deeper ViT (larger than ViT-B).

[cls] token. When it is not explicitly needed by the method, using the average of the patch tokens instead of the class token saves memory without much change on the accuracies of the network ( Zhai et al., 2022a ) .

### 3.6 Techniques for High Performance Masked Image Modeling

While there are several approaches to masked pretraining, the state-of-the-art systems that employ them tend to pair MIM with other techniques. For example, the ConvNextV2 architecture, which was state of the art on ImageNet (for models trained with only public data) when released, employs MAE pretraining ( Woo et al., 2023 ) . Interestingly, the authors point out that simply pretraining a ConvNextV2 with the MAE framework is subpar. They propose adding a novel normalization layer, called “global response normalization,” that proves vital to reaching state-of-the-art results ( Woo et al., 2023 ) .

In other works that claim state-of-the-art performance on image classification and semantic segmentation, MIM pretraining is paired with distillation. While some MIM routines involve reconstructing the masked portion of the input in pixel space, another option is to use a teacher network to generate target representations of the unmasked image. Zhou et al. (2022a) propose iBOT, which uses ViTs for both the teacher and the student in distillation-based MIM and outperforms prior methods on ImageNet classification. Subsequently, Liu et al. (2022b) propose dBOT, an updated distillation-based MIM approach which also achieves state-of-the-art results on image classification and semantic segmentation. A major finding in their work is that the choice of the teacher model does not have to be chosen carefully if the distillation is done in stages. This is where the teacher is updated periodically to match the student’s weights and the student is reinitialized. Oquab et al. (2023) employ similar distillations to train smaller models from a ViT-g teacher with much better performance than training from scratch. This line of work highlights that pairing distillation with MIM is extremely effective.

For object detectors that utilize MIM to outperform prior work, techniques that allow MIM to work with recent and high performing pyramid ViTs like Swin are critical. Since pyramid ViTs collapse patches, random masking can leave some local windows with no information. Li et al. (2022d) propose an approach to masking that accounts for the hierarchical structure of these models called “uniform masking.” They constrain the masking to hide equal amounts of information in each local window ensuring that each has some information intact. This technique helps self-supervised models (on ImageNet1K) outperform supervised models (even on ImageNet22K) on object detection benchmarks Li et al. (2022d) .

### 3.7 Evaluating Your SSL Models

#### 3.7.1 Evaluation with labels

Self-supervised pre-training is mainly evaluated on image classification, since it has been at the core of computer vision for decades. The three main common protocols are referred to as k k -nearest neighbors (KNN), linear and full fine-tuning evaluations (ranked by order of complexity). They are offline evaluations, meaning that they are done independently of the self-supervised training procedure, conversely to online evaluation, which are performed during training. While online evaluation can provide a useful signal of downstream performance, because it’s optimized alongside the varying self-supervised learning objective, it can be misleading. In addition, to these procedures which require labels for the downstream task, more recently, RankMe ( Garrido et al., 2022a ) has appeared as a viable alternative to costly evaluations, and is used as an oracle to final accuracy without having to do any training.

##### KNN

is one of the best known algorithms of machine learning and has been extensively used throughout the fields. With regards to image classification, a KNN classifier determines the label of a data point from the labels of its neighbors.

Formally speaking, the model is first used to extract frozen features 𝒳 = x 1 , … , x n \mathcal{X}=x_{1},...,x_{n} (often l 2 l_{2} -normalized), from all the images in the training dataset. To classify a new image, we extract its feature representation x ′ x^{\prime} , and retrieve its k k nearest-neighbors. They are the k k vectors of the training set 𝒳 \mathcal{X} that have highest cosine-similarity with x ′ x^{\prime} . Then, the vanilla approach applies a majority voting scheme: every neighbor counts + 1 +1 in its corresponding label, and we choose the label with most votes at the end. More sophisticated approaches use a weigthed voting scheme. Instead of counting + 1 +1 in its corresponding label, every neighbor counts a weight w = f ⁡ ( x T ​ x ′ ) w=f(x^{T}x^{\prime}) , for instance DINO implemenation employs w = e x T ​ x ′ / T w=e^{x^{T}x^{\prime}/T} ( Caron et al., 2021 ) . This allows to account for imbalanced training set, not i.i.d. features, and usually gives more accurate results, at the cost of introducing an additional hyperparameter T T .