Self-Supervised Learning with Swin Transformers

In plain words

Most self-supervised image methods are built for convolutional networks. This paper asks what happens when the backbone is a transformer instead. It uses Swin Transformer, a general-purpose vision transformer that also works for detection and segmentation, so the learned features can be tested on those tasks as well as on classification.

The training method is a mix of two existing ideas. There are two copies of the network. One copy learns from gradients, and the other follows it slowly, as an average of its past versions. Each image is shown twice in different random versions. The network must match each version to the other one and separate it from a memory of earlier images. The method combines two existing ideas. The contribution is a tuned set of settings for transformers, and a check of how the features transfer to other tasks.

TL;DR

• MoBY is a combination of MoCo v2 and BYOL, with a transformer backbone and a tuned recipe §2.
• Two encoders (online and target) each with a projector; the online side adds a predictor. The target is a momentum average of the online weights, and the momentum value rises toward 1 during training §2 fig. 1.
• The loss is a contrastive (InfoNCE) loss against a queue of 4096 target keys, temperature 0.2 §2 eq. 1 alg. 1.
• Asymmetric drop path: 0.2 on the online encoder, none on the target tab. 5.
• Author-reported: 75.0% ImageNet-1K linear top-1 for Swin-T and 72.8% for DeiT-S at 300 epochs tab. 1. Transfer to COCO and ADE20K is about on par with supervised pretraining tab. 3 tab. 4.
• Code is MIT-licensed. It shows several settings the paper does not state (LR schedule, normalisation, augmentation details). Read Paper vs code first.

Problem

• Self-supervised learning had shown strong transfer mainly with convolutional networks. Transformer-based self-supervised work had reported only ImageNet linear evaluation, not detection or segmentation §1.
• ViT and DeiT are not designed for dense prediction, so they are hard to use for those downstream tasks §1. The paper uses Swin Transformer, which is described as general-purpose, to allow such tests §1 §2.
• Recipes from MoCo v3 and DINO reached similar linear accuracy, but the paper says they rely on heavier tricks, and MoCo v3 uses a batch size of 4096 §3.1.

Method

The MoBY approach has an online encoder and a target encoder. Each consists of a backbone and a projector head, a 2-layer MLP. The online side has an extra predictor head, also a 2-layer MLP §2 fig. 1.

1. Views. Each image gives two random views. The two encoders use asymmetric augmentations, inherited from BYOL §2. The specific augmentations are not listed in the paper build.py.
2. Online branch. The online backbone, projector and predictor encode each view; the predictor output is the query §2 alg. 1.
3. Target branch. The target backbone and projector encode the other view with no gradient; the output is the key §2 alg. 1.
4. Contrastive loss. For each online query, the positive is the target key of the other view of the same image. The negatives are the entries of a queue of target keys §2 eq. 1. The queue is separate for each view alg. 1.
5. Symmetric sum. The total loss adds the two directions, query 1 against key 2 and query 2 against key 1 alg. 1.
6. Update. The online encoder is updated by gradients. Then the target is updated as a momentum average of the online weights, and the new keys are enqueued §2 alg. 1.
7. Momentum schedule. The momentum starts at 0.99 and rises toward 1 over training §2 tab. 6.

Backbone. The default backbone is Swin-T, which the paper says has similar complexity to ResNet-50 and DeiT-S §2. The architecture and hyperparameters are said to be in the Swin paper §2.

Key equation for one online view q, with positive key k₊, queue keys kᵢ, temperature τ and queue size K eq. 1:

\mathcal{L}_q = -\log \frac{\exp(q \cdot k_{+}/\tau)}{\sum_{i=0}^{K} \exp(q \cdot k_i/\tau)} eq. 1

Pseudocode. The paper's PyTorch-style listing is alg. 1. The version below follows it, with the normalisation and LR details from the official code, which the paper does not state.

# MoBY training step  [alg. 1] [§2]
# view1, view2: two augmented views of each image, (B, 3, H, W)
# queue1, queue2: negatives from the target branch, (C, K)
#   [code: models/moby.py#L81-L84]
# Returns the symmetric contrastive loss

def training_step(view1, view2, momentum, queue1, queue2):
    # Online branch: encoder, projector, predictor  [§2] [alg. 1]
    queries1 = predictor(projector(encoder_online(view1)))  # (B, C)
    queries2 = predictor(projector(encoder_online(view2)))  # (B, C)
    # Normalise queries to unit length; not in eq. 1
    #   [code: models/moby.py#L142-L147]
    queries1 = normalize(queries1)
    queries2 = normalize(queries2)

    # Target branch: momentum-averaged weights, no gradients  [§2] [alg. 1]
    with no_grad:
        keys1 = projector_target(encoder_target(view1))  # (B, C)
        keys2 = projector_target(encoder_target(view2))  # (B, C)
        # Normalise keys to unit length  [code: models/moby.py#L155-L159]
        keys1 = normalize(keys1)
        keys2 = normalize(keys2)

    # Symmetric loss: each query against the other view's key  [alg. 1] [eq. 1]
    loss = contrastive_loss(queries1, keys2, queue2)
    loss = loss + contrastive_loss(queries2, keys1, queue1)

    # Gradient step on the online branch only  [§2] [alg. 1]
    update_online_branch(loss)

    # Momentum update of the target, after the online step  [§2] [alg. 1]
    encoder_target = momentum * encoder_target + (1 - momentum) * encoder_online
    projector_target = momentum * projector_target + (1 - momentum) * projector

    # Push the new keys into each queue and drop the oldest  [alg. 1]
    enqueue(queue1, keys1)
    enqueue(queue2, keys2)
    return loss


def contrastive_loss(queries, keys, queue):
    # Positive logit per sample: (B, 1)  [eq. 1]
    positive = sum(queries * keys, dim=1, keepdim=True)
    # Negative logits against the queue: (B, K)  [eq. 1]
    negatives = queries @ queue  # queue is detached
    # Temperature scaling of all logits  [alg. 1] [code: models/moby.py#L131]
    logits = concat([positive, negatives], dim=1) / temperature
    # Cross-entropy with the positive at index 0  [alg. 1]
    return cross_entropy(logits, zeros(B))


def momentum_at(step, total_steps, start=0.99):
    # Cosine ramp from the start value up to 1  [code: models/moby.py#L93]
    # The paper says only "gradually increased to 1"  [§2]
    cosine = (1 + cos(pi * step / total_steps)) / 2
    return 1 - (1 - start) * cosine

Hyperparameters

Pretraining. Values from the paper unless marked "code".

Linear evaluation on ImageNet-1K.

Transfer.

Results

Author-reported. Linear evaluation on ImageNet-1K uses top-1 with center crop on the validation set §3.1. Supervised rows are from ImageNet-1K classification, as reported in the paper tab. 1.

Parameters and throughput from tab. 1: DeiT-S 22M, 4.6 GFLOPs; Swin-T 29M, 4.5 GFLOPs; 940.4 and 755.2 img/s respectively.

Ablations, Swin-T, ImageNet-1K linear top-1 (%) §3.3 tab. 5 tab. 6 tab. 2:

The BN variant gives +1.1 points but changes the Swin architecture, so the paper does not use it in the standard settings tab. 2. The paper's 100-epoch numbers are not comparable to its 300-epoch headline numbers.

Paper vs code

Official code: github.com/SwinTransformer/Transformer-SSL, commit 4510de1f21ee6f9810a74494c254081dd8f2c383, MIT (Microsoft Corporation) LICENSE. The repository is linked from the paper's abstract and conclusion [abstract] §4. We cloned it and read moby_main.py, models/moby.py, optimizer.py, lr_scheduler.py, data/build.py, moby_linear.py, config.py and the configs.

1. Learning rate schedule. Paper: "a fixed learning rate of 0.001" §2. Code: cosine decay after a 5-epoch warmup, to a minimum of 5e-6 moby_swin_tiny.yaml config.py. The code's scheduler is timm cosine, called per step moby_main.py.
2. Batch scaling. Paper: batch size 512 §3.1. Code: base LR is multiplied by the total batch size divided by 512 at startup, so the LR scales with batch size moby_main.py. The same scaling is applied to the linear evaluation moby_linear.py.
3. Momentum schedule. Paper: "gradually increased to 1" §2; no formula. Code: cosine ramp from the start value to 1 over all steps moby.py.
4. Normalisation. Paper: eq. 1 uses plain dot products §2 eq. 1. Code: predictor outputs and target projections are L2-normalised moby.py moby.py moby.py moby.py.
5. Online drop path. Paper: the headline value is not named; the ablation covers 0.05 to 0.2 tab. 5. Code: 0.2 in the Swin config moby_swin_tiny.yaml, 0.1 as the library default config.py. The DeiT config also uses 0.2 moby_deit_small.yaml.
6. Drop path per block. Paper: not stated. Code: a linear ramp from 0 to the rate across all Swin blocks swin_transformer.py.
7. Augmentations. Paper: "asymmetric data augmentations" from BYOL §2; no parameters given. Code: both views use random resized crop (scale from 0.08 to 1), horizontal flip, colour jitter with probability 0.8, grayscale with probability 0.2 and a Gaussian blur. View 1 always blurs; view 2 blurs with probability 0.1 and solarises with probability 0.2 build.py. Blur sigma is drawn from 0.1 to 2.0 build.py.
8. Headline Swin-T number. Paper: 75.0% tab. 1. Repository README: 75.3% for 300-epoch MoBY Swin-T README.md README.md. We found no explanation of the gap in the repository.
9. Queue and gather. Paper: queue of K keys §2. Code: keys are gathered from all GPUs before enqueuing, and the global batch must divide K moby.py moby.py.
10. Queue initialisation. Paper: not stated. Code: random normal vectors, normalised per column moby.py.
11. Projector and predictor width. Paper: "2-layer MLP" §2. Code: hidden 4096 with BatchNorm and ReLU, output 256, and the queue has 256 rows moby.py moby.py.
12. BatchNorm in heads. Projector and predictor use BatchNorm, converted to synchronised BatchNorm across GPUs moby.py. Paper: not stated.
13. Linear evaluation. Paper: linear classifier on a frozen backbone, learning rate from the grid §3.1. Code additionally uses SGD with momentum 0.9 moby_linear.py, label smoothing 0.1 moby_linear.py, and drop path 0.2 moby_linear.py.
14. Transfer code. Paper: detection and segmentation follow Swin's implementations §3.2. Code: no detection or segmentation code is in this repository. Not verified.

No discrepancies were found in the InfoNCE form, the queue size of 4096, the temperature of 0.2, or the target drop path of 0.0.

Pitfalls

• Equation 1 and the denominator. As written, the sum runs from 0 to K and does not show the positive inside it. The code and the algorithm use one positive plus K negatives, so the positive is in the denominator alg. 1 moby.py.
• Normalisation is essential. Temperature 0.2 is calibrated for unit-length vectors. The code normalises both sides moby.py. Implement it even though eq. 1 does not show it.
• Learning rate depends on batch size. The paper's "fixed" 0.001 matches only the 512 batch. Other batch sizes get a linearly scaled LR moby_main.py. The paper does not say this.
• Asymmetry is the main recipe choice. Adding drop path to the target costs 1.9 points at 100 epochs §3.3. The target must use drop path 0 build.py.
• Momentum schedule shape. Paper says only "gradually increased". The code's cosine ramp reaches 1 at the final step moby.py. Other shapes are untested here.
• Random queue at the start. Early steps use random vectors as negatives moby.py.
• Compute. Each step runs the backbone four times: two online views and two target views moby.py. The target is a full second backbone, so parameter memory roughly doubles.
• BatchNorm across views. The projector and predictor see each view as a separate batch moby.py. Batch statistics then depend on how views are grouped.
• Selection bias. The paper reports the best of three drop path values for detection §3.2 and tunes ADE20K LR over three values §3.2. It does not say whether tuning used a held-out split. The linear LR is also picked from a grid §3.1, and the paper does not say which split was used.
• Not a pure SSL comparison. The paper notes that its ResNet-based comparisons use weaker supervised augmentation §3.2. The supervised Swin baselines use Swin's supervised settings §3.2.
• BN variant. The moby_swin_tiny_bn config changes the Swin normalisation moby_swin_tiny_bn.yaml. Use it only if you want the non-standard result from Table 2.

Lineage

• Builds on: MoCo v2, arXiv:2003.04297 (no card) for the queue and contrastive loss; BYOL, arXiv:2006.07733 (no card) for the asymmetric encoders and momentum schedule §2. Swin Transformer [/papers/2103.14030] is the backbone §2. DeiT, arXiv:2012.12877 (no card), is the second backbone §3.1.
• Compared against: MoCo v3 [/papers/2104.02057] and DINO [/papers/2104.14294], both with ViT or DeiT backbones §3.1 tab. 1.
• Superseded by: not known from content in this wiki.