[Paper Review] PPT: Token Pruning and Pooling for Efficient Vision Transformers

WU Xin-jian, Fanhu Zeng+2 more|arXiv (Cornell University)|Oct 3, 2023

Advanced Neural Network Applications6 citations

TL;DR

This paper proposes PPT, a novel framework that adaptively combines token pruning and pooling in Vision Transformers to reduce computational redundancy without trainable parameters. By dynamically selecting pruning or pooling based on layer depth and input-specific token distributions, PPT reduces over 37% FLOPs and improves throughput by 45% for DeiT-S with no accuracy drop on ImageNet.

ABSTRACT

Vision Transformers (ViTs) have emerged as powerful models in the field of computer vision, delivering superior performance across various vision tasks. However, the high computational complexity poses a significant barrier to their practical applications in real-world scenarios. Motivated by the fact that not all tokens contribute equally to the final predictions and fewer tokens bring less computational cost, reducing redundant tokens has become a prevailing paradigm for accelerating vision transformers. However, we argue that it is not optimal to either only reduce inattentive redundancy by token pruning, or only reduce duplicative redundancy by token merging. To this end, in this paper we propose a novel acceleration framework, namely token Pruning & Pooling Transformers (PPT), to adaptively tackle these two types of redundancy in different layers. By heuristically integrating both token pruning and token pooling techniques in ViTs without additional trainable parameters, PPT effectively reduces the model complexity while maintaining its predictive accuracy. For example, PPT reduces over 37% FLOPs and improves the throughput by over 45% for DeiT-S without any accuracy drop on the ImageNet dataset. The code is available at https://github.com/xjwu1024/PPT and https://github.com/mindspore-lab/models/

Motivation & Objective

• Address the computational inefficiency of Vision Transformers (ViTs) caused by redundant attention across all tokens.
• Overcome the limitation of existing methods that target only one type of redundancy—either inattentive or duplicative—by jointly addressing both.
• Develop a lightweight, parameter-free mechanism to integrate token pruning and pooling adaptively across layers and input instances.
• Enable efficient inference and faster training for ViTs without fine-tuning or architectural overhaul.
• Achieve a superior accuracy-FLOPs trade-off compared to state-of-the-art token compression methods.

Proposed method

• Propose an adaptive token compression module that selects between pruning and pooling based on the distribution of token importance scores in each layer.
• Use a heuristic decision criterion based on the variance of token scores to determine whether to apply pruning (for high-variance, distinct tokens) or pooling (for low-variance, similar tokens).
• Integrate the module into standard transformer blocks without adding any trainable parameters, ensuring compatibility with pre-trained models.
• Apply the adaptive strategy differently across layers: favor pooling in shallow layers (high similarity) and pruning in deeper layers (distinct importance).
• Dynamically adjust the policy per input instance based on the local token score distribution, enabling instance-aware compression.
• Set a learnable threshold τ to control the preference between pruning and pooling, optimized via fine-tuning or set empirically in off-the-shelf scenarios.

Experimental results

Research questions

• RQ1 Can jointly applying token pruning and pooling lead to better efficiency-accuracy trade-offs than using either method alone in Vision Transformers?
• RQ2 Is there a layer-dependent pattern in token redundancy that justifies different compression strategies (pruning vs. pooling) at different depths?
• RQ3 Can an instance-aware, non-trainable policy selection mechanism effectively balance inattentive and duplicative redundancy across diverse inputs?
• RQ4 How does the choice of decision threshold τ affect the performance and robustness of the adaptive compression framework?
• RQ5 Can the proposed method be deployed with minimal accuracy degradation on standard benchmarks like ImageNet without fine-tuning?

Key findings

• PPT reduces FLOPs by over 37% and improves throughput by more than 45% for DeiT-S on ImageNet without any accuracy drop.
• The adaptive selection strategy outperforms fixed rule-based policies (e.g., pooling in early layers, pruning in later ones) and random selection, especially at lower FLOP levels.
• Token pooling alone performs better than token pruning alone in the proposed framework, but combining both via adaptive selection yields the best overall performance.
• The optimal decision threshold τ is found to be 6×10⁻⁵ with fine-tuning and 7×10⁻⁵ in the off-the-shelf setting, balancing pruning and pooling preferences.
• The method maintains strong generalization across different ViT architectures and achieves state-of-the-art efficiency-accuracy trade-offs on ImageNet.
• The framework is effective for inference acceleration and enables faster training with minimal accuracy degradation, even without additional parameters.