PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient Self-supervised Vision Transformers for Representation Learning

Chunyuan Li, Jianwei Yang, Pengchuan Zhang, Mei Gao, Bin Xiao, Xiyang Dai, Lu Yuan, Jianfeng Gao

arXiv:2106.09785v2cs.CVcs.AIcs.LG

TL;DR

Self-supervised vision Transformers can be accurate but computationally expensive, and efficient multi-stage designs lose fine-grained region correspondences. EsViT combines sparse multi-stage Transformers with non-contrastive region matching, achieving strong ImageNet and transfer results while reducing resource demands.

  • Problem

    Transformer-based self-supervised vision systems require substantial computation, while multi-stage efficiency can lose automatically learned fine-grained region correspondences.

  • Method

    EsViT combines a multi-stage Transformer with sparse self-attention and a non-contrastive region-matching pre-training task alongside view-level prediction.

  • Results

    EsViT outperforms prior systems on ImageNet linear evaluation and transfers better than its supervised counterpart on 17 of 18 downstream datasets.

  • Takeaways & Limitations

    Combining architectural efficiency with region matching yields a self-supervised vision Transformer with strong accuracy, smaller model size, and substantially lower compute demands.

  • Takeaways & Limitations

    Applying region matching to monolithic Transformers can be prohibitively expensive because its computation scales as O(T^2).

Abstract

from arXiv · show

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

1 INTRODUCTION

Transformer-based self-supervised vision methods aim to learn general-purpose representations, but existing approaches can require substantial computation. EsViT combines efficient multi-stage Transformers with region matching to recover fine-grained correspondences and achieve strong accuracy and transfer performance.

  • Motivation: Transformer-based self-supervised vision methods inherit strong representation-learning results but require substantial computation for long image-patch sequences.Existing approaches can consume extensive resources, while computer-vision SSL seeks features reusable for classification, detection, and segmentation.
  • Contribution: EsViT presents a multi-stage architecture with sparse self-attention and a region-based pre-training task for efficient self-supervised representation learning.The paper combines both techniques rather than relying only on architectural efficiency.
  • Findings: Self-supervised monolithic Transformers can automatically discover semantic correspondences between local image regions.This property is reported as an intriguing finding of the paper.
  • Findings: Multi-stage architectures greatly reduce compute complexity but cause the loss of automatically learned local-region correspondences.Region matching is introduced to address this lost property.
  • Results: 81.3% top-1 accuracy is achieved by EsViT on ImageNet linear evaluation, with 3.5× parameter efficiency and at least 10× higher throughput than the cited prior SoTA.The cited prior comparison is MoCo-v3 at 81.0% using ViT-BN-L/7.
  • Results: EsViT outperforms its supervised Swin Transformer counterpart on 17 of 18 downstream linear-classification datasets.The comparison concerns transfer of learned representations to downstream tasks.

2 METHODS

EsViT uses multi-stage sparse-attention Transformers to reduce computational cost while adding non-contrastive region matching to preserve local correspondences. The method combines view-level and region-level prediction, with region matching designed to recover dependencies lost through patch merging.

  • Multi-stage ViT: Multi-stage ViTs use patch merging and sparse self-attention modules across stages to reduce Transformer computation.The first stage embeds non-overlapping image patches, while later stages merge neighboring patch features.
  • Correspondence learning: Self-supervised monolithic ViTs automatically learn correspondences between regions of differently augmented views without an explicit region-level objective.The paper evaluates this property using a simulated ImageNet correspondence benchmark.
  • View-level task: View-level prediction matches student and teacher representations across different augmented views, using pooled view features and an MLP head.The view set may contain global and local crops, with the teacher receiving only global views.
  • Region-level task: Patch merging dilutes region correspondences in multi-stage Transformers because view-level matching leaves region-to-region correspondence unspecified.The paper motivates a region-level task to leverage local features and model their co-occurrences and structures.
  • Region-level task: Non-contrastive region matching directly matches corresponding local features across views to learn invariant contextualized region representations.The best-matching teacher feature is selected by cosine similarity, and the loss encourages region dependencies to be captured.
  • Training objective: EsViT minimizes the sum of region-level and view-level losses while updating the teacher as an exponential moving average of the student.The combined objective is L = LR + LV.
  • Computational overhead: For multi-stage Transformers, region matching adds around 1.2× memory and 1.05× computational cost, whereas it quickly runs out of memory for monolithic Transformers as batch size increases.With 224×224 images, the top-layer sequence length is 49 for the multi-stage model versus 196 for ViT/DeiT, requiring 3 times less region-matching computation.

3 RELATED WORKS

Prior region-level methods and self-supervised vision Transformer approaches motivate EsViT’s non-contrastive region matching and efficiency focus.

  • LR is a non-contrastive region-matching task that determines correspondences from the two most similar grid features without negative samples or a queue.DenseCL uses contrastive region matching, whereas LR uses soft targets without negatives.
  • LR targets automatic correspondence learning in self-supervised monolithic Transformers, unlike most region-level tasks designed primarily for dense visual prediction.
  • Most region-level tasks improve dense prediction but sacrifice image classification, while LR consistently improves classification performance.
  • Transformer-based self-supervised vision research remains early, with prior methods exploring large compute budgets, large models, or smaller and medium-size Transformers.

4 EXPERIMENTAL RESULTS

Experiments evaluate EsViT against prior self-supervised systems on ImageNet, transfer tasks, architectures, pre-training datasets, and region-matching objectives. EsViT combines strong accuracy with substantially improved efficiency, while LR consistently improves representations and convergence.

  • ImageNet comparisons: 81.3% top-1 accuracy is achieved by EsViT (Swin-B/W=14) on ImageNet linear probing, with 95.5% top-5 and 79.3% k-NN accuracy.This configuration is reported as the best-performing system in the comparison.
  • ImageNet comparisons: At similar model size and compute complexity, EsViT (Swin-B) exceeds DINO (ViT-B/16) by 2.2% linear probe and 2.8% k-NN accuracy.EsViT (Swin-B) also slightly exceeds DINO (ViT-B/8) while providing 4× higher throughput.
  • ImageNet comparisons: EsViT is 3.5× parameter-efficient and has at least 10× higher throughput than previous SoTA MoCo-v3.The comparison reports MoCo-v3 at 81.0% with ViT-BN-L/7.
  • Transfer learning: EsViT outperforms its supervised counterpart on 17 out of 18 downstream linear-classification datasets.The evaluation reports averaged scores for each method across the 18 datasets.
  • Region matching: LR consistently improves performance across model scales and window sizes, with k-NN gains of around 1–2%.Figure 4 also shows that LR helps model convergence and can be used as a drop-in addition to the view-level task.
  • Region matching: LR improves both classification and dense visual prediction performance, whereas DenseCL improves dense prediction but hurts classification performance.The paper notes that non-contrastive methods generally show lower dense-prediction performance, consistent with observations for BYOL.
  • Qualitative and correspondence studies: LR improves correspondence accuracy from 66% to 91% in the ImageNet validation evaluation.The evaluation uses augmented views of the same image with ground-truth correspondences.

5 CONCLUSIONS

EsViT combines sparse multi-stage Transformers with non-contrastive region matching to achieve strong self-supervised vision performance with less compute and smaller models.

  • 5 CONCLUSIONS: EsViT combines a multi-stage Transformer with sparse self-attentions and a non-contrastive region-matching pre-training task.The multi-stage design improves efficiency, while region matching addresses lost correspondence learning.
  • 5 CONCLUSIONS: The paper identifies automatic correspondence learning between local regions as a property of self-supervised monolithic Vision Transformers.
  • 5 CONCLUSIONS: The study finds that multi-stage architectures can lose this correspondence-learning property despite reducing computational complexity.

ETHICS STATEMENT

The paper frames efficient self-supervised vision learning as important because existing systems require extensive computational resources, and presents EsViT as a more affordable alternative.

  • ETHICS STATEMENT: Existing state-of-the-art self-supervised vision systems can require extensive computation and may be environmentally unfriendly.The paper describes these resource demands as barriers to participation in AI research.
  • ETHICS STATEMENT: EsViT is proposed to provide more affordable and efficient solutions for experimenting with self-supervised vision.
  • ETHICS STATEMENT: EsViT reports 3.5× parameter efficiency and 10× higher throughput than previous state-of-the-art systems.

REPRODUCIBILITY STATEMENT

The paper supports reproducibility through released code, checkpoints, detailed specifications, pseudocode, and an explicit training algorithm description.

  • REPRODUCIBILITY STATEMENT: The authors release pretrained checkpoints and a codebase for reproducible research.
  • REPRODUCIBILITY STATEMENT: The paper provides PyTorch-style pseudocode, supplementary example code, and detailed hyperparameter specifications.
  • REPRODUCIBILITY STATEMENT: Algorithm 1 describes EsViT with view-level and region-level losses using two augmented views.
  • REPRODUCIBILITY STATEMENT: The algorithm specifies student and teacher outputs as view-level probabilities, region-level features, and region-level probability vectors.
  • REPRODUCIBILITY STATEMENT: Training updates the student with AdamW, the teacher with exponential moving averages, and view and region centers with momentum updates.

A.2 NETWORK ARCHITECTURE CONFIGURATIONS AND IMPLEMENTATION DETAILS

The appendix describes multi-stage Transformer architectures, sparse attention mechanisms, implementation settings, and region-level learning comparisons used to study efficient self-supervised vision.

  • A.2 NETWORK ARCHITECTURE CONFIGURATIONS AND IMPLEMENTATION DETAILS: Multi-stage Transformers use patch merging and sparse self-attention modules across stages to reduce computational complexity.Early stages process more patches, making sparse attention particularly useful there.
  • A.2 NETWORK ARCHITECTURE CONFIGURATIONS AND IMPLEMENTATION DETAILS: Swin Transformer alternates shifted local windows so neighboring features participate in different self-attention groupings.
  • A.2 NETWORK ARCHITECTURE CONFIGURATIONS AND IMPLEMENTATION DETAILS: Swin-T and Swin-S variants with W = 14 enable full self-attention in stages 3 and 4 to study longer sequences.
  • A.2 NETWORK ARCHITECTURE CONFIGURATIONS AND IMPLEMENTATION DETAILS: The experiments use relative position bias rather than absolute position embeddings to accommodate augmented views with varied resolutions.
  • A.2 NETWORK ARCHITECTURE CONFIGURATIONS AND IMPLEMENTATION DETAILS: Most compared region-level tasks improve object detection but reduce ImageNet classification accuracy, while DetCo improves both tasks with a sophisticated objective.
  • A.2 NETWORK ARCHITECTURE CONFIGURATIONS AND IMPLEMENTATION DETAILS: Region matching adds acceptable computational cost to multi-stage architectures, unlike its large overhead when applied directly to monolithic Transformers.

C.3 EXPERIMENTAL SETTINGS OF PRE-TRAINING AND EVALUATION ON IMAGENET

The experiments define pre-training and linear-probe evaluation settings for ImageNet, including optimizer schedules, architecture configurations, computational comparisons, and benchmark protocols.

  • Experimental settings: Pre-training uses ImageNet-1K without labels, AdamW, batch size 512, 300 epochs, linear warmup, and cosine learning-rate decay.The base learning rate follows the linear scaling rule.
  • Comparisons: The study compares computational cost and linear-probe performance across architecture and SSL configurations, including ResNet-50 and Swin-T-scale models.Tables report memory usage, iteration time, and performance comparisons.
  • Architecture configurations: Swin-B is designed to match the model size and computational complexity of ViT-B/DeiT-B, while Swin-T and Swin-S target ResNet-50- and ResNet-101-scale complexity.These architectures form the basis of the experimental systems.
  • Evaluation: Linear probing freezes the backbone and trains a supervised linear classifier using concatenated view-level features from the last four Transformer layers.The results are reported on ImageNet-1K and related benchmark settings.
  • Comparisons: The evaluation includes linear-probe results for ResNet-50 networks trained with different self-supervised learning methods.This comparison is presented in Table 12.

C.5 LINEAR PROBE ON A SUITE OF SMALL DATASETS

The paper evaluates transferability with linear classification on 18 small datasets, using standardized dataset descriptions and automated logistic-regression hyperparameter selection.

  • Datasets: 18 datasets are used for linear-probe evaluation, with dataset-specific class counts, train/test sizes, metrics, and public sources documented.UCF101 is converted from video by extracting the middle frame, using the first TensorFlow split.
  • Evaluation protocol: Each downstream classifier is logistic regression trained with scikit-learn’s L-BFGS implementation for up to 1,000 iterations.The reported metric is determined separately for each dataset.
  • Evaluation protocol: The L2 regularization strength λ is selected by validation-set sweeps over 10^-6 to 10^6 using 96 logarithmically spaced values.This procedure follows the cited evaluation protocol.
  • Results: At comparable ResNet-50 or Swin-T scale, Table 11 reports transfer results for EsViT, supervised models, and CLIP-based comparisons.EsViT with Swin-T and the supervised Swin-T baseline are both pre-trained on ImageNet-1K.
  • Pre-training settings: Larger and less curated pre-training datasets are accompanied by dataset statistics and schedules chosen to match the number of augmented views.These settings are summarized in Table 13.

C.7 RESULTS ON CORRESPONDENCE LEARNING

The correspondence experiments evaluate whether self-supervised Transformer representations preserve region correspondences across augmented and practically varied image views.

  • Evaluation: Correspondence quality is evaluated on 50K ImageNet validation images using mild augmentations with known ground-truth correspondences and top-10 predictions.The evaluation considers correspondence accuracy and distance error.
  • Architecture and task effects: Multi-stage Transformers lose the correspondence property learned by monolithic DINO, while region matching alleviates this issue.The comparison is visualized across the monolithic, multi-stage, and region-matching configurations.
  • Quantitative results: 0.95 accuracy and 2.49 distance error are achieved by DINO with a monolithic Transformer, providing a reference for automatic correspondence learning.The figure caption describes this as strong evidence of the property.
  • Augmented correspondences: For left-to-right flipped views, view-level pre-training alone is insufficient for good multi-stage correspondences, whereas region matching improves them significantly.Similar observations are reported for rows 3 and 4.
  • Practical correspondences: Region matching works particularly well for one-to-many mappings, cartoon-to-toy views, and seasonal scene changes beyond standard pre-training augmentations.These cases are shown in rows 2, 5, and 6 of Figure 9.

C.8 MORE VISUALIZATION RESULTS OF ATTENTION MAPS

Attention visualizations show that monolithic DINO, multi-stage Transformers, and EsViT with region matching learn different attention patterns and foreground-selection behaviors.

  • Foreground attention: Monolithic DINO identifies main foreground objects, but this property is lost when changing to a multi-stage architecture.The multi-stage example has 24 heads versus 6 for the monolithic architecture.
  • Foreground attention: Region matching reduces redundant attention patterns in EsViT’s multi-stage architecture.This comparison is made between the multi-stage model and EsViT with the region-level task.
  • Query-dependent attention: DINO attends mainly to the foreground object even when the query lies in a background region.The figure caption and accompanying analysis identify this behavior.
  • Query-dependent attention: Region-level pre-training improves background-region behavior by training the model to match individual regions.The paper attributes DINO’s foreground focus to view-level invariance on ImageNet.
  • Attention diversity: Region matching diversifies attention across heads by increasing higher-entropy heads and decreasing lower-entropy heads.The paper describes this as modulation of the entropy distribution among heads.
  • Visualization layout: Figures 10 and 11 rank top-layer heads by softmax entropy and show top-60%-mass and full attention maps for specified query locations.The figures compare attention patterns for LR and LV.
Loading 2106.09785v2…