[Paper Review] Efficient Self-supervised Vision Transformers for Representation Learning

Chunyuan Li, Jianwei Yang+6 more|arXiv (Cornell University)|Jun 17, 2021

Domain Adaptation and Few-Shot Learning72 references83 citations

TL;DR

EsViT combines a multi-stage ViT architecture with a non-contrastive region-matching pre-training task to achieve high ImageNet linear probe accuracy with much higher throughput and parameter efficiency than prior SSL methods.

ABSTRACT

This paper investigates two techniques for developing efficient self-supervised vision transformers (EsViT) for visual representation learning. First, we show through a comprehensive empirical study that multi-stage architectures with sparse self-attentions can significantly reduce modeling complexity but with a cost of losing the ability to capture fine-grained correspondences between image regions. Second, we propose a new pre-training task of region matching which allows the model to capture fine-grained region dependencies and as a result significantly improves the quality of the learned vision representations. Our results show that combining the two techniques, EsViT achieves 81.3% top-1 on the ImageNet linear probe evaluation, outperforming prior arts with around an order magnitude of higher throughput. When transferring to downstream linear classification tasks, EsViT outperforms its supervised counterpart on 17 out of 18 datasets. The code and models are publicly available: https://github.com/microsoft/esvit

Motivation & Objective

• Investigate efficiency limitations of monolithic versus multi-stage Vision Transformers (ViTs) in self-supervised learning.
• Propose a region-based pre-training task to recover fine-grained region correspondences lost in multi-stage architectures.
• Demonstrate that combining multi-stage architectures with region-matching improves downstream transfer and linear probe performance.
• Provide evidence that EsViT achieves high accuracy with substantially better throughput and parameter efficiency than prior SSL methods.

Proposed method

• Study the trade-offs of multi-stage ViT architectures with sparse self-attention for SSL.
• Introduce a non-contrastive region-matching loss L_R operating on local features across augmented views.
• Combine L_R with a view-level non-contrastive loss L_V in a joint objective L = L_R + L_V.
• Use a teacher-student setup with EMA parameter updates for SSL training.
• Evaluate region matching and attention behavior through qualitative visualizations and ablations.

Experimental results

Research questions

• RQ1 How do multi-stage ViTs compare to monolithic ViTs for self-supervised representation learning in terms of efficiency and accuracy?
• RQ2 Can a region-matching pre-training task recover fine-grained region correspondences lost in multi-stage architectures and improve representations?
• RQ3 What is the impact of combining L_R with L_V on ImageNet linear probe performance and transfer learning?
• RQ4 What are the design choices and limitations of the region-matching objective across architectures?
• RQ5 Do EsViT representations generalize better than supervised counterparts on downstream tasks?

Key findings

• EsViT achieves 81.3% top-1 accuracy on ImageNet linear probing, outperforming prior SSL methods with substantially higher throughput and fewer parameters.
• EsViT shows better transfer than its supervised counterpart on 17 of 18 downstream classification datasets.
• A region-matching loss L_R mitigates loss of fine-grained region correspondences in multi-stage ViTs and improves representation quality.
• The non-contrastive L_R combined with L_V yields improved k-NN and linear probe performance and preserves efficiency gains from multi-stage architectures.
• EsViT with region matching demonstrates notably higher throughput (up to ~10x) and parameter efficiency (up to ~3.5x) relative to prior SOTA SSL methods on ImageNet.
• Qualitative analyses show EsViT learns meaningful region correspondences and attention patterns, especially when L_R is used.