Token Pooling in Vision Transformers

TL;DR

This paper introduces Token Pooling, a novel token downsampling method for vision transformers that reduces computational costs by exploiting redundancies, achieving similar accuracy with significantly fewer computations.

Contribution

The paper proposes Token Pooling, a new token downsampling technique that approximates tokens through clustering, improving the cost-accuracy trade-off in vision transformers.

Findings

Token Pooling reduces computation by 42% on DeiT with no accuracy loss.

Softmax-attention acts as a low-pass filter, enabling effective token redundancy pruning.

Token Pooling outperforms prior downsampling methods in efficiency and accuracy.

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.