Vision Transformer with Super Token Sampling

TL;DR

This work targets the redundancy and inefficiency of capturing global context in shallow Vision Transformer layers. It introduces STViT, a hierarchical backbone that uses Super Token Attention (STA) to compress visual tokens into a small set of semantically meaningful super tokens, perform global reasoning on them, and upsample back to the original tokens. STA decomposes global attention into a sparse association map and a low-dimensional self-attention over $m$ super tokens, dramatically reducing computation while preserving long-range dependencies. Across ImageNet classification, COCO detection/segmentation, and ADE20K semantic segmentation, STViT achieves strong accuracy with competitive FLOPs, including $86.4\%$ top-1 on ImageNet-1K for STViT-L without extra data. The approach offers a practical path to efficient global context modeling in early network stages for diverse vision tasks.

Abstract

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.

Figures (4)

• Figure 1: FLOPs vs. Accuracy on $224^2$ ImageNet-1K images.
• Figure 2: Visualization of early-stage attention maps for different vision transformers. For global attention in DeiT touvron2021training and local attention in Swin liu2021swin, only a few neighboring tokens (filled with red color) work for an anchor token (green box), resulting in local representations with high redundancy. Compared with such ViTs, our method can learn global representations even for shallow layers.
• Figure 3: The architecture of Super Token Vision Transformer (STViT).
• Figure 4: Visualization of super tokens from initial grid to learned ones. Only the surrounding super tokens in the red box are used to compute the associations for every token in the green box.