PSViT: Better Vision Transformer via Token Pooling and Attention Sharing

TL;DR

This paper introduces PSViT, a novel vision transformer architecture that reduces redundancy through token pooling and attention sharing, leading to improved accuracy and efficiency in image recognition.

Contribution

The paper proposes a new ViT model with token pooling and attention sharing, automatically learned as hyper-parameters, to enhance feature representation and speed-accuracy trade-off.

Findings

Achieves up to 6.6% accuracy improvement on ImageNet.

Effectively reduces redundancy in tokens and attention maps.

Enhances feature representation and computational efficiency.

Abstract

In this paper, we observe two levels of redundancies when applying vision transformers (ViT) for image recognition. First, fixing the number of tokens through the whole network produces redundant features at the spatial level. Second, the attention maps among different transformer layers are redundant. Based on the observations above, we propose a PSViT: a ViT with token Pooling and attention Sharing to reduce the redundancy, effectively enhancing the feature representation ability, and achieving a better speed-accuracy trade-off. Specifically, in our PSViT, token pooling can be defined as the operation that decreases the number of tokens at the spatial level. Besides, attention sharing will be built between the neighboring transformer layers for reusing the attention maps having a strong correlation among adjacent layers. Then, a compact set of the possible combinations for different token pooling and attention sharing mechanisms are constructed. Based on the proposed compact set, the number of tokens in each layer and the choices of layers sharing attention can be treated as hyper-parameters that are learned from data automatically. Experimental results show that the proposed scheme can achieve up to 6.6% accuracy improvement in ImageNet classification compared with the DeiT.