Source-linked AI summary
Visual Transformers: Token-based Image Representation and Processing for Computer Vision
Bichen Wu, Chenfeng Xu, Xiaoliang Dai, Alvin Wan, Peizhao Zhang, Zhicheng Yan, Masayoshi Tomizuka, Joseph Gonzalez, Kurt Keutzer, Peter Vajda
TL;DR
Computer vision’s pixel-convolution paradigm treats pixels and concepts uniformly and struggles with spatially distant relationships. The paper introduces semantic visual tokens and token-space transformers, reporting higher accuracy with lower computation across classification and segmentation.
Problem
Convolutions uniformly process pixels, model high-level concepts across all images, and struggle to relate spatially distant concepts.
Method
Visual Transformers use spatial attention to convert feature maps into compact semantic tokens, apply transformers to model token interactions, and project tokens back when needed.
Results
VTs improve accuracy with lower computational cost across image classification and semantic segmentation, including 4.6 to 7 points higher ImageNet ResNet accuracy and 0.35 points higher segmentation mIOU.
Takeaways & Limitations
Visual Transformers provide a token-space alternative to pixel-convolution processing that attends to important regions and relates spatially distant semantic concepts.
Takeaways & Limitations
The filter-based tokenizer can waste computation by modeling sparse high-level concepts with a fixed set of learned weights, and comparisons also reflect differing training recipes.
Abstract
from arXiv · showhide
Computer vision has achieved remarkable success by (a) representing images as uniformly-arranged pixel arrays and (b) convolving highly-localized features. However, convolutions treat all image pixels equally regardless of importance; explicitly model all concepts across all images, regardless of content; and struggle to relate spatially-distant concepts. In this work, we challenge this paradigm by (a) representing images as semantic visual tokens and (b) running transformers to densely model token relationships. Critically, our Visual Transformer operates in a semantic token space, judiciously attending to different image parts based on context. This is in sharp contrast to pixel-space transformers that require orders-of-magnitude more compute. Using an advanced training recipe, our VTs significantly outperform their convolutional counterparts, raising ResNet accuracy on ImageNet top-1 by 4.6 to 7 points while using fewer FLOPs and parameters. For semantic segmentation on LIP and COCO-stuff, VT-based feature pyramid networks (FPN) achieve 0.35 points higher mIoU while reducing the FPN module's FLOPs by 6.5x.
1 Facebook Reality Labs, 2 Facebook AI, 3 UC Berkeley
The authors are affiliated with Facebook and UC Berkeley.
- The Facebook-affiliated email addresses include wbc, xiaoliangdai, stzpz, zyan3, and vajdap.
- The listed affiliations span Facebook and UC Berkeley.
- The UC Berkeley-affiliated email addresses include xuchenfeng, alvinwan, tomizuka, jegonzal, and keutzer.
1. Introduction
The paper argues that pixel-based convolutions waste computation and struggle with long-range semantic relationships. Visual Transformers address these issues by converting feature maps into compact semantic tokens and modeling token interactions with transformers, improving efficiency and accuracy across classification and segmentation.
- Challenges: Convolutions process all image patches uniformly, despite foreground regions being more important than backgrounds for many vision tasks.This creates spatial inefficiency in computation and representation.
- Challenges: Convolutions apply high-level filters across all images, wasting computation on concepts absent from the current image.The paper contrasts ubiquitous low-level patterns with sparse high-level concepts such as dog features.
- Challenges: Convolutional filters struggle to relate spatially distant semantic concepts without added model and computational complexity.Prior remedies include larger kernels, deeper models, dilated convolutions, global pooling, and non-local attention.
- Visual Transformer: Visual Transformers convert feature maps into a compact set of semantic visual tokens and use self-attention to model their relationships.Tokens support image-level prediction directly and can be projected back for semantic segmentation.
- Visual Transformer: VTs attend to important regions, encode image-relevant concepts, and relate distant concepts in token-space.
- Results: VT-based ResNets improve ImageNet top-1 accuracy by 4.6 to 7 points while reducing the replaced stage's FLOPs by 6.9x.For semantic segmentation, VT-based FPN improves mIOU by 0.35 points while reducing regular FPN module FLOPs by 6.4x.
2. Relationship to previous work
The paper distinguishes Visual Transformers from prior vision uses of transformers, graph convolutions, attention, and efficiency-focused convolutional models. Its approach combines convolutional low-level feature extraction with token-space modeling of higher-level concepts.
- Transformers in vision models: Vision Transformer divides images into 16 × 16 patches and feeds them to a standard transformer, incurring an extremely high computational price.The paper reports that ViT required up to 7 GPU years and 300M JFT images to outperform competing convolutional variants.
- Transformers in vision models: VTs use convolutions for low-level features and transformers for relating high-level concepts, while spatial attention focuses on important regions.The paper contrasts this with treating every image patch equally.
- Transformers in vision models: DETR and VT use transformers for orthogonal applications, so their insights could be combined in one model.
- Graph convolutions in vision models: Prior graph-based methods extract latent variables, model node interactions, and project nodes back to feature maps, analogous to VT processing.The paper identifies concept redundancy as a susceptibility of these approaches.
- Attention in vision models: Vision attention methods modulate feature maps, make convolution spatially adaptive, or use non-local operators equivalent to self-attention.
- Efficient vision models: Efficiency-focused work improves computation through efficient operators or neural architecture search over convolutional networks.
3. Visual Transformer
Visual Transformers replace late-stage pixel-convolution processing with compact semantic tokens, transformer-based token interaction, and feature-map refinement. The design combines convolutional processing for low-level patterns with token-based modeling for higher-order concepts and addresses limitations of fixed filters through recurrent tokenization.
- Architecture: VT models use convolutions early for densely distributed low-level patterns, then visual transformers later for sparse higher-order semantic concepts.Image-level tasks use visual tokens, while pixel-level tasks use the augmented feature map.
- Architecture: A VT module groups pixels into semantic concepts, applies a transformer to model token relationships, and projects tokens back to pixel-space.This produces an augmented feature map for downstream pixel-level processing.
- Transformer: As few as 16 visual tokens can replace hundreds of analogous latent nodes while using input-dependent token-to-token weights through self-attention.The transformer computes relationships from key-query products rather than fixed inference weights.
- Filter-based Tokenizer: Filter-based tokenizers use point-wise convolutions to assign pixels to semantic groups and spatially pool them into weighted visual tokens.The attention weights form semantic groups and produce weighted averages over the feature map.
- Filter-based Tokenizer: Fixed filter-based tokenizers may waste computation because sparse high-level concepts appear in only a few images but are modeled simultaneously.The method identifies this as a limitation of convolutional filters used for token extraction.
- Recurrent Tokenizer: Recurrent tokenizers condition current token extraction on previous-layer tokens, incrementally refining visual concepts across VT modules.In practice, recurrent tokenizers begin with the second VT because they require tokens from a previous module.
4. Using Visual Transformers in vision models
This section applies Visual Transformers as modular building blocks for classification and semantic segmentation, replacing selected convolutional processing with compact token-based modeling.
- VT modules are specified by feature-map channel size, visual-token channel size, and number of visual tokens.
- Image classification: For classification, VT-ResNets replace the last ResNet convolutional stage with the same number of VT modules.The replacements are made in ResNet-{18, 34, 50, 101}.
- Semantic segmentation: VT-based segmentation addresses convolutional complexity and long-range interaction challenges by modeling a small number of visual tokens independently of image resolution.Token-space modeling bypasses long-range interactions between pixels.
- Semantic segmentation: VT-FPN extracts visual tokens from feature maps at different resolutions, merges them with one transformer, and projects them back to the original feature maps.
- Semantic segmentation: VT-based FPNs replace expensive high-resolution spatial convolutions while preserving multiscale, detail-preserving feature maps for segmentation.The baseline FPN fuses ResNet feature maps in a top-down manner to produce rich semantics.
5. Experiments
Experiments show that Visual Transformers improve classification and segmentation while reducing computation, with controlled ablations supporting semantic tokenization, recurrent tokenizers, token interaction modeling, compact token sets, and feature-map projection.
- Image classification: VT-ResNets achieved up to 2.2 points higher top-1 validation accuracy while using 244M and 384M fewer FLOPs than ResNet18 and ResNet34, respectively.Under the default training recipe, VT-ResNets also showed larger training-accuracy gains, indicating heavier overfitting and motivating stronger regularization.
- Tokenizer ablations: Filter-based and clustering-based tokenizers significantly outperformed pooling-based tokenizers, supporting semantic grouping rather than spatial downsampling.The two stronger tokenizer types performed similarly overall, with differences varying between ResNet18 and ResNet34.
- Tokenizer ablations: Recurrent tokenizers produced higher accuracy than alternatives because they are content-aware.The experiments used a filter-based tokenizer initially and recurrent tokenizers in subsequent VT modules.
- Token relationship modeling: Transformer-based token interaction outperformed graph convolutions and omitting token interactions, supporting dense modeling of relationships between semantic concepts.The authors attribute this advantage to transformers allowing tokens to encode image-dependent semantic concepts rather than fixed concepts.
- Token efficiency and projection: Increasing the number of visual tokens yielded negligible or no accuracy improvement, while projecting tokens back to feature maps improved performance.The projection result is consistent with feature maps retaining important spatial information needed for vision understanding.
- Advanced training and segmentation: With the advanced training recipe, VT-ResNets achieved 4.6 to 7 points higher accuracy than baseline ResNets, while segmentation VT-FPNs used 6.5x fewer FLOPs with slightly higher mIoU.Replacing the original ResNet stage reduced its FLOPs by up to 6.9x; the advanced recipe used longer training, stronger augmentation and regularization, and distillation.
6. Conclusion
Visual Transformers replace pixel-array convolutions with semantic visual tokens and token-space transformers, improving accuracy across image classification and segmentation tasks while reducing computation in segmentation.
- Visual Transformers represent high-level image concepts as visual tokens and use transformers to relate those concepts in token-space.
- 4.6 to 7 points: VT improves ResNet accuracy on ImageNet.The comparison is reported for ImageNet accuracy using an advanced training recipe.
- 0.35 points higher mIoU: VT-based FPNs improve semantic segmentation on LIP and COCO-stuff while using 6.5x fewer FLOPs than convolutional FPN modules.
A. Stage-wise model description of VT-ResNet
VT-ResNet configurations are specified stage by stage using three hyper-parameters that control feature-map channels, token channels, and the number of visual tokens.
- Three hyper-parameters control a VT module: output feature-map channel size C, visual-token channel size CT, and token count L.
- The VT-ResNet model configurations are described in Table 11.
B. More visualization results
Additional Figure 7 visualizations show the spatial attention produced by the tokenizer on images from the LIP dataset.
- Figure 7 provides more visualizations of spatial attention on LIP dataset images.
C. Clustering-based tokenizer
The clustering-based tokenizer forms semantic groups from image-specific feature content rather than applying fixed filters, producing tokens from clustered pixel representations.
- Clustering-based tokenizer: The clustering-based tokenizer treats each pixel as a sample and applies k-means to find L centroids representing semantic concepts.
- Clustering-based tokenizer: The centroids are stacked into WK, which replaces WA to form L semantic groups of channels.
- Clustering-based tokenizer: T = SOFTMAXHW (XWK)T X computes the visual tokens from the feature map and the centroid-based grouping weights.
- Clustering-based tokenizer: The implementation normalizes pixels, initializes centroids with a spatially downsampled feature map, and runs Lloyd’s algorithm.
- Clustering-based tokenizer: The tokenizer models only concepts present in the current image but is not designed to choose the most discriminative concepts.
- More visualization results: Figure 7 visualizes spatial attention, with red denoting higher values and blue denoting lower values; tokens focus on different semantic image areas without supervision.