PaperlayerView exact revision on arXiv

Source-linked AI summary

Vision Transformer with Super Token Sampling

Huaibo Huang, Xiaoqiang Zhou, Jie Cao, Ran He, Tieniu Tan

arXiv:2211.11167v2cs.CV

TL;DR

Shallow vision-transformer layers can incur redundant computation on local features, while local attention and early convolutions limit long-range dependency modeling. The paper introduces super token attention and the STViT backbone to model global context efficiently, and reports results across classification, detection, and segmentation tasks.

  • Problem

    Shallow vision-transformer layers can redundantly compare tokens for local features, while local attention and early-stage convolutions sacrifice long-range dependency modeling.

    Source: §1 · p. 1

  • Method

    STA sparsely associates visual tokens with super tokens, applies self-attention among super tokens, and maps the result back to the original token space within a hierarchical vision-transformer backbone.

    Source: §1 · p. 2

  • Results

    STViT performs strongly across vision tasks; STViT-L reaches 86.4% top-1 accuracy on ImageNet-1K without extra training data, while STViT-B surpasses the Swin Transformer counterpart on COCO box AP, COCO mask AP, and ADE20K mIOU.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Super tokens provide a compact representation that suppresses local redundancy while allowing global dependencies to be modeled through attention in the super-token space.

    Source: §5 · p. 8

    STViT uses operations that are not computationally efficient on GPUs, so its models are not among the fastest at similar FLOPs; it was also not trained on large-scale datasets such as ImageNet-21K due to computational constraints.

    Source: E. Limitations and Broader Impacts · p. 13

Abstract

from arXiv · show

Vision transformer has achieved impressive performance for many vision tasks. However, it may suffer from high redundancy in capturing local features for shallow layers. Local self-attention or early-stage convolutions are thus utilized, which sacrifice the capacity to capture long-range dependency. A challenge then arises: can we access efficient and effective global context modeling at the early stages of a neural network? To address this issue, we draw inspiration from the design of superpixels, which reduces the number of image primitives in subsequent processing, and introduce super tokens into vision transformer. Super tokens attempt to provide a semantically meaningful tessellation of visual content, thus reducing the token number in self-attention as well as preserving global modeling. Specifically, we propose a simple yet strong super token attention (STA) mechanism with three steps: the first samples super tokens from visual tokens via sparse association learning, the second performs self-attention on super tokens, and the last maps them back to the original token space. STA decomposes vanilla global attention into multiplications of a sparse association map and a low-dimensional attention, leading to high efficiency in capturing global dependencies. Based on STA, we develop a hierarchical vision transformer. Extensive experiments demonstrate its strong performance on various vision tasks. In particular, without any extra training data or label, it achieves 86.4% top-1 accuracy on ImageNet-1K with less than 100M parameters. It also achieves 53.9 box AP and 46.8 mask AP on the COCO detection task, and 51.9 mIOU on the ADE20K semantic segmentation task. Code is released at https://github.com/hhb072/STViT.

1. Introduction

Shallow vision-transformer layers can spend computation on redundant local features, while local attention and early convolutions sacrifice long-range dependency modeling. The paper introduces super token attention (STA) and the STViT backbone to combine efficient early-stage global representations with local-feature modeling.

1. Introduction

  • Global attention in shallow layers focuses on a few adjacent tokens, while local attention and early-stage convolutions reduce redundancy but sacrifice global dependency modeling.This motivates seeking efficient global representations at early network stages.Source: §1 · pp. 1–2
  • STA samples super tokens through sparse token-to-super-token associations, applies self-attention in the super-token space, and maps the representations back to the original tokens.The lower-dimensional attention is intended to capture long-range dependency with less computation than token-space self-attention.Source: §1 · p. 2
  • STViT combines convolutional layers with stacks of super token transformer blocks containing CPE, STA, and ConvFFN modules.CPE and ConvFFN use depth-wise convolutions to enhance local-feature representation at low computation cost.Source: §1 · p. 2
  • STViT-L achieves 86.4% top-1 accuracy on ImageNet-1K without extra training data; STViT-B achieves 53.9 box AP and 46.8 mask AP on COCO, and 51.9 mIOU on ADE20K, surpassing the Swin Transformer counterpart by +2.1, +2.1 and +2.4, respectively.The experiments cover image classification, object detection, instance segmentation, and semantic segmentation.Source: §1 · p. 2

2. Related Works

Related work covers vision-transformer backbones and methods for reducing attention cost, alongside superpixel algorithms that group image content. STA differs from the described attention alternatives by using sparse soft associations to super tokens and performing self-attention in their space.

Vision Transformers

  • Vision-transformer research includes hierarchical architectures, self-attention variants, and positional encodings designed for vision tasks.ViT adapts the Transformer architecture to vision, following its use in NLP.Source: §2 · p. 3
  • Efficient attention methods reduce computation through local regions, downsampled or sampled tokens, hashing-based buckets, merged tokens, or orthogonal-space attention; STA instead uses sparse soft token-to-super-token mappings.STA performs self-attention in the super-token space.Source: §2 · p. 3

Superpixel Algorithms

  • Superpixel methods include graph-based partitioning, clustering, deep clustering, and differentiable segmentation; STA's sampling algorithm is motivated mainly by SSN but uses sparse sampling in a more attention-like manner.SSN is described as an end-to-end differentiable superpixel segmentation framework.Source: §2 · p. 3

3. Method

STViT combines a hierarchical convolutional backbone with super token attention (STA), which models global dependencies in a compact token space while retaining local modeling. Its sampling, attention, and upsampling steps provide an efficient route between visual tokens and super tokens.

3.1. Overall Architecture

  • STViT uses four hierarchical stages of Super Token Transformer blocks, with convolutional layers to support local feature modeling.Each block combines Convolutional Position Embedding, STA, and a convolutional feed-forward network.Source: §3.1 · p. 3

3.2. Super Token Attention

  • STA aggregates visual tokens into super tokens, applies global self-attention there, then maps the result back to the visual-token space.The paper describes STA as Super Token Sampling, Multi-Head Self-Attention, and Token Upsampling.Source: §3.2 · p. 3

3.2.1 Super Token Sampling

  • Super Token Sampling initializes super tokens by averaging tokens in regular grid regions, computes sparse token-to-super-token associations, and updates super tokens as weighted sums.Each token's association uses only its nine surrounding super tokens; Figure 4 depicts the initial grid and learned super tokens.Source: §3.2.1 · p. 4, Figure 4 · p. 4

3.2.2 Self-Attention for Super Tokens

  • Standard self-attention operates on the sampled super tokens, whose compact representations can emphasize global contextual dependencies.The stated super-token attention map has size m×m.Source: §3.2.2 · pp. 4–5

3.2.3 Token Upsampling

  • Token Upsampling uses the association map to map super tokens back to visual tokens and add them to the original tokens.The paper notes that sampling loses most local details, motivating this return to the visual-token space.Source: §3.2.3 · p. 5

3.2.4 Analysis of Complexity and Redundancy

  • STA decomposes global attention into sparse associations and a smaller super-token attention matrix, with lower computation cost than global attention when m<N.The paper frames this decomposition as a way to reduce redundant comparisons while retaining effective global representations.Source: §3.2.4 · p. 5

3.3. Implementation Details

  • The first two stages use sampling grids of 8×8 and 4×4, yielding 49 super tokens at 224×224 resolution; later stages use tokens directly as super tokens.The sampling iteration count is set to 1 to reduce sampling cost.Source: §3.3 · p. 5

4. Experiments

STViT is evaluated on image classification, detection, instance segmentation, and semantic segmentation, with ablations examining its components. The reported results compare favorably with prior backbones across these tasks.

4.1. Image Classification

  • 86.4% ImageNet-1K top-1 accuracy for STViT-L at 384 × 384 surpasses Swin-B by 2.2% and CSwin-B by 0.9%, with 22% fewer FLOPs and double speed than CaiT-s48.At 224 × 224, STViT-S reaches 83.6% top-1 accuracy with 4.4G FLOPs, surpassing Swin-T, CSwin-T, and MPViT-S by 2.3%, 0.9%, and 0.6%.Source: Table 2 · p. 6, §4.1 · p. 7

4.2. Object Detection and Instance Segmentation

  • Under Cascade Mask R-CNN, STViT-B beats Swin-B by +2.0 box AP and +1.5 mask AP, and DAT-B by +0.9 box AP and +1.0 mask AP.The comparison includes backbones with similar FLOPs and models with more FLOPs.Source: Table 4 · p. 7, §4.2 · p. 7

4.3. Semantic Segmentation

  • STViT-B surpasses Swin Transformer by +3.1 mIOU at similar model size and by +2.6 mIOU at larger model size on ADE20K semantic segmentation.It also outperforms Uniformer by +1.2 mIOU and +1.2 Multi-Scale(MS) mIOU.Source: §4.3 · p. 7, Table 5 · p. 8

4.4. Ablation Study

  • Adding the convolution stem, 1 × 1 convolution projection, CPE, and ConvFFN progressively improves a hierarchical four-stage network with STA.The ablation sequence in Table 6 starts with the STA network and adds the remaining STViT modules.Source: Table 6 · p. 8, §4.4 · p. 8
  • A four-stage STA backbone without convolutional modules outperforms DeiT-S by 2.4% and Swin-T by 1.0%, despite using no position embedding.DeiT uses absolute PE, while Swin uses relative PE.Source: §4.4 · p. 8
  • Super token sampling helps in the first two stages but hurts performance in the third; larger grid sizes yield better performance.A single sampling iteration performs best; zero iterations use average-pooled features, while two iterations perform slightly worse with 0.03G more FLOPs.Source: §4.4 · p. 8
  • All tested positional encodings improve performance: APE by 0.2%, RPE by 0.3%, and CPE by 0.4%; ConvFFN and its shortcut also help.CPE is selected for flexibility across arbitrary resolutions.Source: §4.4 · p. 8

5. Conclusion

STViT adapts superpixels to the token space, using super tokens to reduce local redundancy while retaining global representations. Its STA factors global attention into sparse association and low-dimensional attention operations, and experiments report strong performance across vision tasks.

5. Conclusion

  • STViT uses super tokens to aggregate similar tokens, reducing local redundancy while enabling efficient global dependency modeling through sparse association and low-dimensional attention.Experiments cover image classification, object detection, instance segmentation, and semantic segmentation.Source: §5 · p. 8

A. Algorithm of Super Token Attention

STA samples super tokens, applies self-attention to them, then maps them back to the token space. Sparse associations restrict each token's association computation to its surrounding super tokens.

A. Algorithm of Super Token Attention

  • STA generates initial super tokens by average pooling, computes sparse associations with surrounding super tokens, applies multi-head self-attention, and maps the attended super tokens back to the original token space.Associations use only each token's 3 × 3 surrounding super tokens; the described configuration sets the iteration count to one.Source: A. Algorithm of Super Token Attention · p. 11

B. Architecture Details

The architecture details specify convolutional operations for token embedding, between-stage downsampling, and local modeling. Table 8 is identified as the architecture reference for ImageNet classification at 224 × 224.

B. Architecture Details

  • The architecture uses four 3×3 convolutions in its stem, stride-2 3×3 convolutions between stages, and depth-wise 3×3 convolutions in CPE and ConvFFN.Table 8 describes architectures for ImageNet classification at 224 × 224.Source: B. Architecture Details · p. 11, Table 8 · p. 12

C. Experimental Settings

ImageNet-1K classification uses 300-epoch scratch training at 224 × 224, while COCO and ADE20K follow task-specific detection and segmentation setups.

ImageNet Image Classification

  • ImageNet-1K models train from scratch for 300 epochs at 224 × 224, using DeiT's training strategy and specified augmentation and regularization.For 384 × 384 input, models are fine-tuned for 30 epochs.Source: C. Experimental Settings · p. 12

COCO Object Detection and Instance Segmentation

  • COCO evaluation uses Mask R-CNN or Cascade Mask R-CNN with either 12-epoch 1× or 36-epoch 3× +MS training.The 3× +MS setting uses multi-scale resizing.Source: C. Experimental Settings · p. 12

ADE20K Semantic Segmentation

Isotropic comparison

  • Table 9 compares isotropic STViT and Swin without Conv Stem, Projection, CPE, or ConvFFN; Swin† uses global self-attention in stage three.Source: Table 9 · p. 13

D. Isotropic comparison

Under isotropic settings, STViT outperforms Swin in comparisons with similar configurations, which the authors say indicates gains from STA itself.

D. Isotropic comparison

  • STViT outperforms Swin under similar isotropic settings, which the authors interpret as evidence that STA itself brings performance gains.Source: D. Isotropic comparison · p. 13

E. Limitations and Broader Impacts

STViT's GPU-inefficient operations limit its speed relative to models with similar FLOPs, and the paper reports no direct negative social impact.

Limitations

Broader Impacts

Related papers

Loading 2211.11167v2…