Source-linked AI summary

CrossViT: Cross-Attention Multi-Scale Vision Transformer for Image Classification

Chun-Fu Chen, Quanfu Fan, Rameswar Panda

arXiv:2103.14899v2cs.CV

TL;DR

Vision transformers provide promising image-classification performance, but the paper studies how to learn multi-scale representations within transformer models. CrossViT uses separate small- and large-patch branches with linear-complexity cross-attention fusion, and experiments show performance better than or comparable to concurrent vision transformers and efficient CNNs.

  • Problem

    The paper asks how to learn multi-scale feature representations in transformer models for image classification.

  • Method

    CrossViT uses two branches for small- and large-patch tokens and fuses them with cross-attention using one query token per branch for linear computation and memory complexity.

  • Results

    CrossViT performs better than or on par with several concurrent vision transformers and achieves comparable results to EfficientNet across accuracy, throughput, and model parameters.

  • Takeaways & Limitations

    The results support using dual-branch multi-scale representations and efficient cross-attention fusion in vision transformers for image classification.

Abstract

from arXiv · show

The recently developed vision transformer (ViT) has achieved promising results on image classification compared to convolutional neural networks. Inspired by this, in this paper, we study how to learn multi-scale feature representations in transformer models for image classification. To this end, we propose a dual-branch transformer to combine image patches (i.e., tokens in a transformer) of different sizes to produce stronger image features. Our approach processes small-patch and large-patch tokens with two separate branches of different computational complexity and these tokens are then fused purely by attention multiple times to complement each other. Furthermore, to reduce computation, we develop a simple yet effective token fusion module based on cross attention, which uses a single token for each branch as a query to exchange information with other branches. Our proposed cross-attention only requires linear time for both computational and memory complexity instead of quadratic time otherwise. Extensive experiments demonstrate that our approach performs better than or on par with several concurrent works on vision transformer, in addition to efficient CNN models. For example, on the ImageNet1K dataset, with some architectural changes, our approach outperforms the recent DeiT by a large margin of 2\% with a small to moderate increase in FLOPs and model parameters. Our source codes and models are available at \url{https://github.com/IBM/CrossViT}.

1. Introduction

The paper investigates multi-scale feature learning in vision transformers and introduces a dual-branch architecture with efficient cross-attention fusion. It reports stronger or comparable performance than concurrent vision transformers and efficient CNNs.

  • Motivation: Vision transformers emerged as convolution-free alternatives to CNNs, but ViT requires very large training datasets while DeiT improves data efficiency.ViT uses embedded image patches as transformer inputs; DeiT uses augmentation and regularization to train high-performance models with fewer data.
  • Approach: Multi-scale representations are studied for vision transformers by combining small- and large-patch tokens in separate branches with different computational complexities.The branches are fused repeatedly so their features can complement each other.
  • Approach: Cross-attention fusion uses a non-patch token from each branch to exchange information, reducing attention-map generation from quadratic to linear time.The proposed token fusion scheme is linear in both computation and memory.
  • Contributions: The paper proposes a dual-branch vision transformer and a linear-complexity cross-attention token-fusion scheme for image classification.These are identified as the work's main contributions.
  • Results: The approach performs better than or on par with concurrent ViT-based methods and achieves comparable accuracy, throughput, and parameter counts to EfficientNet.Figure 1 compares the proposed approach with DeiT and ViT using ImageNet1K models trained from scratch, with circle size indicating model size.

2. Related Works

The related-work discussion situates CrossViT among attention-augmented CNNs, pure vision transformers, and multi-scale CNNs. Unlike these approaches, CrossViT uses a dual-path transformer to extract multi-scale features.

  • CNNs with Attention: Attention-augmented CNNs use channel, spatial, efficient channel, or global self-attention to enhance representations and improve image-classification efficiency.Examples include SENet, CBAM, ECANet, BoTNet, and related hybrid architectures.
  • Vision Transformer: Pure vision transformers include ViT, distilled and pyramid variants, Perceiver, and T2T-ViT, which target performance, data efficiency, scalability, or local structure.These approaches differ from CrossViT's dual-path architecture for multi-scale feature extraction.
  • Multi-Scale CNNs: Multi-scale representations have long been used in image pyramids, scale-space methods, coarse-to-fine approaches, and CNN architectures for recognition and efficiency.Examples include Big-Little Net, OctNet, and bLVNet-TAM.

3. Method

CrossViT is a dual-branch vision transformer that learns multi-scale image features by processing different patch sizes and repeatedly fusing the branches with attention. Its cross-attention fusion uses CLS tokens to exchange information efficiently between branches.

  • Vision Transformer: ViT converts images into patch tokens, adds a CLS token and positional embeddings, then uses the final CLS representation for classification.The CLS token summarizes patch-token information through repeated transformer encoders.
  • Motivation: Fine-grained patches improve accuracy but increase computation: patch size 16 outperforms patch size 32 by 6% while requiring 4× more FLOPs.CrossViT addresses this accuracy-complexity trade-off with branches operating at different patch scales.
  • Architecture: CrossViT uses large and small branches with different patch sizes, encoder counts, embedding widths, and computational complexities.The large branch uses coarse patches with more encoders and wider embeddings, while the small branch uses fine patches with fewer encoders and smaller embeddings.
  • Feature Fusion: The model repeatedly fuses branch features using all-attention, class-token, pairwise, or cross-attention strategies.These alternatives respectively combine all tokens, CLS tokens only, spatially aligned patch tokens, or CLS tokens with the other branch’s patch tokens.
  • Cross-Attention Fusion: Cross-attention uses one branch’s CLS token to query the other branch’s patch tokens, then projects the exchanged information back to the originating branch.The module aligns dimensions with projections and omits a feed-forward network after cross-attention.
  • Cross-Attention Fusion: Cross-attention has linear computation and memory complexity because only CLS tokens serve as queries, and it achieves the best accuracy among the four fusion strategies.All-attention instead requires quadratic computation because it processes all tokens jointly.

4. Experiments

Experiments evaluate CrossViT on ImageNet1K and downstream classification tasks against DeiT, recent vision transformers, and CNNs. Results show improved or competitive accuracy, with efficiency depending on model configuration and ablations favoring lightweight branches and cross-attention fusion.

  • Comparisons with DeiT: CrossViT improves DeiT-Ti, DeiT-S, and DeiT-B accuracy by 1.2%, 1.2%, and 0.4% points, respectively.Architectural changes add another 0.3–0.5% point with a small increase in FLOPs and parameters.
  • Comparisons with CNNs: CrossViT-18† reduces FLOPs and parameters almost by half compared to DeiT-B while remaining 1.0% more accurate.CrossViT-9† and CrossViT-15† improve accuracy by approximately 2.5–5% but incur 30–50% more FLOPs and parameters than baselines.
  • Comparisons with SOTA Transformers: CrossViT-18† reaches 82.8% accuracy versus ViT-B’s 77.9%, while requiring 50% less FLOPs and parameters.CrossViT-15† outperforms small models from the compared transformer approaches with comparable FLOPs and parameters.
  • Comparisons with CNNs: CrossViT-15 outperforms the compared ResNet-family CNNs in accuracy while being smaller and more efficient, except for slightly faster ResNet-101.At higher resolution, CrossViT-15† and CrossViT-18† are competitive with EfficientNet on accuracy, throughput, and parameters.
  • Transfer Learning: CrossViT is on par with DeiT models across five downstream classification tasks, including natural-image and medical-image datasets.The transfer experiments evaluate CIFAR10, CIFAR100, Pet, CropDisease, and ChestXRay8.
  • Ablation Studies: Cross-attention achieves the best fusion accuracy with a minor increase in FLOPs and parameters, while more complex S-branches add cost without improving accuracy.More frequent cross-attention or multi-scale encoder stacking likewise increases computation without performance improvement; removing CLS tokens lowers accuracy from 81.0% to 80.0%.
  • Cooperation with Concurrent Works: CrossViT-18+T2T achieves 83.0% top-1 accuracy on ImageNet1K, improving CrossViT-18 by 0.5%.This result demonstrates compatibility between the proposed cross-attention and another ViT variant.

5. Conclusion

CrossViT combines multi-scale patch tokens through a dual-branch vision transformer and cross-attention fusion for image classification. Experiments show performance better than or comparable to concurrent vision transformers and efficient CNN models.

  • CrossViT uses separate branches for different patch scales to learn multi-scale features for image classification.The architecture is designed to improve recognition accuracy by combining image patch tokens of different scales.
  • Cross-attention efficiently exchanges information between the two branches in linear time.The fusion method is intended to combine patch tokens of different scales without quadratic computational complexity.
  • CrossViT performs better than or on par with several concurrent vision transformers and efficient CNN models.

A. More Comparisons and Analysis

Additional analyses show that cross-attention fusion consistently improves baseline vision transformers and produces complementary representations across the two branches. The section also presents feature visualizations and supporting experimental tables.

  • Comparisons with Baselines: Cross-attention fusion consistently improves baseline vision transformers across different primary branches and patch embeddings.The result suggests that multi-scale fusion is effective for different vision transformer configurations.
  • Feature Analysis: The small branch generates more low-level features, while the large branch produces more abstract features.The branches use different numbers of transformer encoders, and their features complement each other in the ensemble.
  • Experimental Details: The section includes comparisons with baselines on ImageNet1K and documents the associated model details and training settings.Table 8 compares various baselines, while Table 9 lists training settings.
  • Feature Analysis: Feature visualizations show selected channels from patch tokens in both branches of CrossViT-S after the last multi-scale transformer encoder.The figure displays 36 randomly selected channels.
Loading 2103.14899v2…