Token Pooling in Vision Transformers
arXiv:2110.03860
Abstract
Despite the recent success in many applications, the high computational requirements of vision transformers limit their use in resource-constrained settings. While many existing methods improve the quadratic complexity of attention, in most vision transformers, self-attention is not the major computation bottleneck, e.g., more than 80% of the computation is spent on fully-connected layers. To improve the computational complexity of all layers, we propose a novel token downsampling method, called Token Pooling, efficiently exploiting redundancies in the images and intermediate token representations. We show that, under mild assumptions, softmax-attention acts as a high-dimensional low-pass (smoothing) filter. Thus, its output contains redundancy that can be pruned to achieve a better trade-off between the computational cost and accuracy. Our new technique accurately approximates a set of tokens by minimizing the reconstruction error caused by downsampling. We solve this optimization problem via cost-efficient clustering. We rigorously analyze and compare to prior downsampling methods. Our experiments show that Token Pooling significantly improves the cost-accuracy trade-off over the state-of-the-art downsampling. Token Pooling is a simple and effective operator that can benefit many architectures. Applied to DeiT, it achieves the same ImageNet top-1 accuracy using 42% fewer computations.
References in corpus (21)
- MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications
- Longformer: The Long-Document Transformer
- Neural Architecture Search: A Survey
- ENet: A Deep Neural Network Architecture for Real-Time Semantic Segmentation
- Linformer: Self-Attention with Linear Complexity
- ConViT: Improving Vision Transformers with Soft Convolutional Inductive Biases
- Generating Long Sequences with Sparse Transformers
- The State of Sparsity in Deep Neural Networks
- Visual Transformers: Token-based Image Representation and Processing for Computer Vision
- Reformer: The Efficient Transformer
- DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification
- Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention
- Big Bird: Transformers for Longer Sequences
- Image Transformer
- Stand-Alone Self-Attention in Vision Models
- Synthesizer: Rethinking Self-Attention in Transformer Models
- Random Feature Attention
- Sparse Sinkhorn Attention
- PoWER-BERT: Accelerating BERT Inference via Progressive Word-vector Elimination
- Centroid Transformers: Learning to Abstract with Attention
- Scalable Vision Transformers with Hierarchical Pooling