[Paper Review] PSViT: Better Vision Transformer via Token Pooling and Attention Sharing

Boyu Chen, Peixia Li+7 more|arXiv (Cornell University)|Aug 7, 2021

CCD and CMOS Imaging SensorsEngineering44 references20 citations

TL;DR

PSViT improves Vision Transformers by introducing learnable token pooling and attention sharing to reduce spatial and inter-layer redundancy, enabling better feature representation and a superior speed-accuracy trade-off. Using a compact search space and AutoML (SPOS), PSViT achieves up to 6.6% top-1 accuracy gain over DeiT on ImageNet and outperforms ResNets on downstream detection and segmentation tasks.

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.

Motivation & Objective

• To address spatial redundancy in Vision Transformers caused by fixed token counts across layers.
• To reduce inter-layer attention map redundancy due to similar attention patterns in adjacent layers.
• To enable flexible computation allocation by learning optimal token numbers and attention sharing configurations.
• To improve feature representation and achieve better speed-accuracy trade-offs in vision transformer architectures.
• To develop a principled, data-driven approach to ViT architecture design using AutoML.

Proposed method

• Introduces token pooling as a mechanism to dynamically reduce the number of tokens at different network stages, adapting to feature complexity.
• Applies attention sharing between adjacent transformer layers to reuse similar attention maps, reducing computation.
• Constructs a compact search space combining token pooling locations and attention sharing choices as hyper-parameters.
• Uses SPOS (Single Path One-Shot) for Neural Architecture Search to jointly optimize token numbers and attention sharing patterns.
• Treats token pooling and attention sharing as learnable hyper-parameters, trained end-to-end via differentiable architecture search.
• Employs 1D and 2D token pooling variants, with 2D pooling preserving spatial structure better for image tasks.

Experimental results

Research questions

• RQ1 Can dynamic token pooling across layers improve feature representation in Vision Transformers?
• RQ2 To what extent do attention maps in adjacent ViT layers exhibit redundancy, and can it be exploited?
• RQ3 Can a compact search space for token pooling and attention sharing enable effective AutoML-based architecture search?
• RQ4 Does joint optimization of token count and attention sharing lead to better accuracy and efficiency trade-offs?
• RQ5 How does PSViT transfer to downstream vision tasks like object detection and instance segmentation?

Key findings

• PSViT achieves up to 6.6% higher top-1 accuracy on ImageNet compared to DeiT, demonstrating significant performance gains.
• PSViT-2D-Tiny achieves 40.8% mAP on COCO object detection and 37.7% mAP on instance segmentation, surpassing ResNet18 and even ResNet101.
• The 2D token pooling variant outperforms the 1D version, indicating that preserving spatial structure improves feature quality.
• The last transformer layer avoids attention sharing, as it benefits most from independent self-attention for final classification.
• The first stages of the network benefit most from attention sharing due to high computational cost and feature smoothness.
• The search space design enables effective AutoML optimization, with PSViT-2D-Tiny achieving strong performance with minimal human intervention.