[Paper Review] Visual Transformers: Token-based Image Representation and Processing for Computer Vision

BoRui Wu, Chenfeng Xu+8 more|arXiv (Cornell University)|Jun 5, 2020

Advanced Neural Network Applications49 references377 citations

TL;DR

The paper introduces Visual Transformers (VTs) that replace part of convolutional networks with token-based transformers operating on a compact set of semantic visual tokens, achieving higher accuracy with fewer FLOPs while enabling efficient classification and segmentation.

ABSTRACT

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.

Motivation & Objective

• Revisit the pixel-convolution paradigm as a root cause of redundancy in vision models.
• Propose a tokenizer to convert feature maps into a small set of semantic visual tokens.
• Relate visual tokens via transformers to model inter-token relationships.
• Project token interactions back to pixel-space for task-specific predictions.
• Demonstrate improved accuracy and reduced compute on classification and segmentation benchmarks.

Proposed method

• Process input with convolutions to learn low-level features, then tokenize the feature map into a small set of visual tokens (L << HW).
• Model interactions among visual tokens using a transformer with content-dependent attention weights.
• Project the transformer outputs back to the pixel space to augment the feature map for downstream tasks.
• Use a tokenizer that can be filter-based or recurrent (token dependencies on prior tokens) to obtain tokens.
• Train VT-ResNets by replacing the last stage of ResNets with VT modules, using 16 tokens and a 1024-channel token representation.
• For segmentation, replace FPN with VT-FPN to compute cross-resolution token interactions and project back to pixel-space.

Experimental results

Research questions

• RQ1 Can a token-based transformer operating on a compact semantic token space outperform traditional pixel-space convolutions in image classification and segmentation?
• RQ2 Do visual tokens enable better modeling of long-range semantic relationships with reduced computational cost?
• RQ3 Is a recurrent, content-aware tokenizer superior to fixed filter-based or pooling-based tokenizers?
• RQ4 Does projecting token interactions back to the feature map preserve pixel-level information and improve task performance?

Key findings

• Replacing the last stage of ResNet with VT modules reduces last-stage FLOPs by up to 6.9x and improves top-1 ImageNet accuracy by 4.6 to 7 points.
• VT-based ResNets achieve higher validation accuracy with substantially fewer FLOPs and similar or fewer parameters compared to baselines (e.g., VT-R18: 72.1% vs 69.9% val accuracy; VT-R34: 75.0% vs 73.3%).
• VT-based FPN for segmentation achieves 0.35 points higher mIoU on COCO-stuff/LIP with 6.4x to 6.5x fewer FLOPs than conventional FPN modules.
• Using 16 visual tokens suffices; increasing token count yields negligible gains.
• Recurrent tokenizers and transformer-based token interaction outperform pooling, clustering, or graph-convolution approaches for token relations.
• Projecting tokens back to feature maps is beneficial for performance, validating the need to preserve spatial information.