Source-linked AI summary

CAT: Cross Attention in Vision Transformer

Hezheng Lin, Xing Cheng, Xiangyu Wu, Fan Yang, Dong Shen, Zhongyuan Wang, Qing Song, Wei Yuan

arXiv:2106.05786v1cs.CVcs.AI

TL;DR

Vision Transformers face rapidly increasing computation when image patches are used as tokens, limiting efficient processing and backbone use for downstream tasks. CAT addresses this with alternating inner-patch and single-channel feature-map attention, forming a hierarchical backbone that combines local and global information. It achieves state-of-the-art performance across reported vision tasks, including 82.8% top-1 accuracy for its base model on ImageNet-1K and improvements on COCO and ADE20K tasks.

  • Problem

    Image-token Transformers require vast computation that grows rapidly with input size, while their same-shaped feature maps limit use as backbones for downstream vision tasks.

  • Method

    CAT alternates inner-patch attention for local information with attention on single-channel feature maps for global information, producing a hierarchical backbone with multi-scale features.

  • Results

    CAT achieves state-of-the-art performance across reported vision tasks, including 82.8% top-1 accuracy on ImageNet-1K and improved performance for object detection and semantic segmentation.

  • Takeaways & Limitations

    CAT provides a backbone that combines CNN-like local and hierarchical features with Transformer-like global information without substantially increasing computation.

Abstract

from arXiv · show

Since Transformer has found widespread use in NLP, the potential of Transformer in CV has been realized and has inspired many new approaches. However, the computation required for replacing word tokens with image patches for Transformer after the tokenization of the image is vast(e.g., ViT), which bottlenecks model training and inference. In this paper, we propose a new attention mechanism in Transformer termed Cross Attention, which alternates attention inner the image patch instead of the whole image to capture local information and apply attention between image patches which are divided from single-channel feature maps capture global information. Both operations have less computation than standard self-attention in Transformer. By alternately applying attention inner patch and between patches, we implement cross attention to maintain the performance with lower computational cost and build a hierarchical network called Cross Attention Transformer(CAT) for other vision tasks. Our base model achieves state-of-the-arts on ImageNet-1K, and improves the performance of other methods on COCO and ADE20K, illustrating that our network has the potential to serve as general backbones. The code and models are available at \url{https://github.com/linhezheng19/CAT}.

1 Introduction

Transformer offers global information extraction for vision, but image-token computation grows rapidly with input size and typical features lack hierarchical scales. CAT alternates local inner-patch attention with efficient global attention to combine CNN-like locality and Transformer-like global information while generating multi-scale features.

  • Motivation: Transformer-based image methods face rapidly increasing computational complexity as input resolution grows and typically produce same-shaped feature maps.These properties make them unsuitable as backbones for some downstream vision tasks.
  • Cross Attention: CAT alternates attention among pixels within each image patch and attention on single-channel feature maps to capture local and global information.The two operations reduce computation relative to standard self-attention.
  • CAT Backbone: CAT builds a hierarchical backbone that generates feature maps at different scales for downstream vision tasks.The proposed hierarchy is described as having characteristics of CNN hierarchical networks.
  • Results: 82.8% top-1 accuracy is achieved by the CAT base model on ImageNet-1K, comparable with state-of-the-art CNN- and Transformer-based networks.CAT also improves the performance of object detection and semantic segmentation methods when used as a backbone.
  • Cross Attention: CAT combines CNN-inspired local feature extraction with Transformer-inspired global information extraction while introducing global attention with no or only a small computation increase.The paper presents this combination as a more efficient way to join the complementary features of CNNs and Transformers.

2 Related work

Related work spans CNN backbones, global- and local-attention Transformers, and hierarchical networks for downstream vision tasks. Existing Transformer approaches remain constrained by costly global interactions, limited local modeling, or difficulty producing CNN-like multi-scale features.

  • CNN-based networks: CNN-based networks became standard vision backbones through shared weights, locality, and other architectural properties supporting feature extraction.The related work discusses applications across object detection and semantic segmentation.
  • Global attention in Transformer-based network: Global-attention Transformers such as ViT and DeiT tokenize image patches and apply multi-head self-attention across patches.Other methods add convolutional projections or combine multiple branches.
  • Global attention in Transformer-based network: Global-attention methods place all patches together for self-attention, causing computational complexity to increase dramatically with input size.This makes them difficult to apply to vision tasks requiring large-resolution inputs.
  • Local attention in Transformer-based network: Local-attention methods such as TNT and Swin capture within-patch relationships while also attempting to model global information.The paper notes computational costs or other problems when combining local and global interactions.
  • Hierarchy networks and downstream tasks: Typical Transformers preserve consistent input and output shapes, making CNN-like hierarchical structures difficult despite their importance for downstream tasks.Pyramidal features are used in object detection and semantic segmentation, while later methods reduce resolution to address hierarchy.

3 Method

CAT combines local inner-patch attention with global cross-patch attention in a hierarchical Transformer architecture. Its stages use patch embedding, cross-attention blocks, and configurable layer counts to support vision tasks.

  • Cross Attention based Transformer: CAT reduces feature-map resolution through patch embedding, then extracts features at multiple scales with successive CAT layers.The initial resolution becomes H1 = H/P and W1 = W/P, with P = 4 in the experiments.
  • Cross Attention based Transformer: CAT stages contain different numbers of layers and use patch projection for downsampling, producing CAT-T, CAT-S, and CAT-B variants with 1×, 2×, and 3× computation.The variants differ according to the number of CABs in stage 3 and the dimension of the patch projection layer.
  • Cross Attention Block: A Cross Attention Block stacks IPSA and CPSA blocks with layer normalization, MLP, and shortcut connections.The architecture uses these blocks to integrate information within patches and between patches.
  • Inner-Patch Self-Attention Block: Inner-Patch Self-Attention applies self-attention within each patch to model local pixel relationships with lower resolution-dependent computation.For H = W = 56, C = 96, and N = 7, IPSA requires approximately 0.15 G FLOPs versus approximately 2.0 G for MSA.
  • Cross-Patch Self-Attention Block: Cross-Patch Self-Attention separates single-channel feature maps into patches and performs attention across patches to capture global spatial information.The design alternates local within-patch processing with communication between patches in each feature map.

4 Experiment

CAT is evaluated across ImageNet-1K classification, COCO detection, ADE20K segmentation, and ablations, showing competitive or improved performance across tasks. The experiments also examine architectural choices including patch embedding, multi-head attention, shifted windows, position encoding, and dropout.

  • Image Classification: 81.8% and 82.8% top-1 accuracy are achieved by CAT-S and CAT-B, respectively, at 224 × 224 resolution.
  • Object Detection: 3.4%, 3.7%, and 4.8% improvements are reported over FCOS, RetinaNet, and Cascade R-CNN with CAT-S under multi-scale COCO detection.
  • Object Detection: 4.2% mask mAP improvement is reported for instance segmentation with CAT-S, and all tested methods outperform their original versions.
  • Semantic Segmentation: 44.9% mIoU is achieved by Semantic FPN with CAT-B at 160k iterations, improving 4.2% over ResNet101.
  • Ablation Study: Absolute position encoding improves performance on three benchmarks, while dropout rate 0.2 performs best and indicates slight overfitting in CPSA.

5 Conclusion

The paper concludes that Cross Attention combines CNN-like local extraction with Transformer-like global information while keeping computation relatively low. CAT provides hierarchical, multi-scale features and achieves state-of-the-art performance across several vision-task datasets.

  • CAT alternates attention within feature-map patches and attention on single-channel feature maps to capture local and global information.
  • CAT generates features at different scales and adapts to different input sizes for other vision tasks.
  • CAT achieves state-of-the-art performance on ImageNet-1K, COCO 2017, and ADE20K.
Loading 2106.05786v1…