[Paper Review] Self-Supervised Learning with Swin Transformers

Zhenda Xie, Yutong Lin+5 more|arXiv (Cornell University)|May 10, 2021

Domain Adaptation and Few-Shot LearningComputer Science18 references112 citations

TL;DR

MoBY combines MoCo v2 and BYOL with Swin Transformer backbones to achieve strong ImageNet linear evaluation and comparable downstream task transfer, using lighter tricks and tuned hyper-parameters.

ABSTRACT

We are witnessing a modeling shift from CNN to Transformers in computer vision. In this work, we present a self-supervised learning approach called MoBY, with Vision Transformers as its backbone architecture. The approach basically has no new inventions, which is combined from MoCo v2 and BYOL and tuned to achieve reasonably high accuracy on ImageNet-1K linear evaluation: 72.8% and 75.0% top-1 accuracy using DeiT-S and Swin-T, respectively, by 300-epoch training. The performance is slightly better than recent works of MoCo v3 and DINO which adopt DeiT as the backbone, but with much lighter tricks. More importantly, the general-purpose Swin Transformer backbone enables us to also evaluate the learnt representations on downstream tasks such as object detection and semantic segmentation, in contrast to a few recent approaches built on ViT/DeiT which only report linear evaluation results on ImageNet-1K due to ViT/DeiT not tamed for these dense prediction tasks. We hope our results can facilitate more comprehensive evaluation of self-supervised learning methods designed for Transformer architectures. Our code and models are available at https://github.com/SwinTransformer/Transformer-SSL, which will be continually enriched.

Motivation & Objective

• Motivate studying self-supervised learning (SSL) with Transformer backbones for both classification and dense prediction tasks.
• Present MoBY, a lightweight SSL method combining MoCo v2 and BYOL with Swin Transformers.
• Evaluate linear ImageNet-1K performance and transfer to COCO object detection/segmentation and ADE20K semantic segmentation.
• Provide baselines and ablations to guide future SSL research using Transformer architectures.

Proposed method

• MoBY blends MoCo v2 style momentum encoder, a key queue, and contrastive loss with BYOL-style asymmetric encoders, augmentations, and momentum schedule.
• Two encoders (online and target) with a 2-layer MLP projector; online adds a 2-layer predictor head; target is a moving-average of online with increasing momentum.
• AdamW optimizer is used with a fixed learning rate (0.001) and weight decay (0.05); queue size K defaults to 4096.
• Asymmetric drop path regularization is applied; targets are updated via momentum; temperature tau and other hyperparameters are tuned.
• Swin Transformer (Swin-T) is used as backbone by default to enable downstream task evaluation; compared against DeiT-S in linear evaluation.
• A 300-epoch training regime for linear evaluation on ImageNet-1K with standard augmentations and a 5-epoch warm-up.

Experimental results

Research questions

• RQ1 How well does MoBY perform for ImageNet-1K linear evaluation with Swin Transformer backbones compared to other SSL methods (MoCo v3, DINO) and DeiT backbones?
• RQ2 Do learned representations from MoBY transfer effectively to downstream tasks such as COCO object detection/instance segmentation and ADE20K semantic segmentation when using Swin backbones?
• RQ3 What is the impact of architectural choices (Swin-T vs DeiT-S) and specific regularization tricks (asymmetric drop path, queue size, temperature) on SSL performance?
• RQ4 Can lighter SSL tricks with Transformer backbones close the gap to supervised pre-training on downstream tasks?

Key findings

• MoBY achieves 72.8% top-1 accuracy on ImageNet-1K linear evaluation with DeiT-S at 300 epochs and 75.0% with Swin-T at 300 epochs.
• Swin-T outperforms DeiT-S by 2.2 percentage points in linear evaluation under the same 300-epoch regime.
• MoBY with Swin-T yields competitive downstream performance on COCO object detection/instance segmentation and ADE20K semantic segmentation, comparable to supervised pre-training, with some metrics showing parity or slight differences.
• Applying MoCo v3/DINO-inspired tricks (e.g., BN before MLP) can yield gains (e.g., +1.1% top-1 with 100 epochs), indicating potential for further improvements.
• MoBY’s linear evaluation and transfer results suggest substantial room for improvement in SSL with Transformer architectures compared to ResNet-based SSL methods.