[论文解读] 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 Applications被引用 6

一句话总结

本文提出PPT，一种新颖的框架，通过自适应地结合视觉Transformer中的token剪枝与池化，以在不引入可学习参数的情况下减少计算冗余。通过根据层深度和输入特定的token分布动态选择剪枝或池化，PPT在DeiT-S上实现了超过37%的FLOPs减少和45%的吞吐量提升，且在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/

研究动机与目标

• 解决视觉Transformer（ViTs）因所有token间注意力计算冗余而导致的计算效率低下问题。
• 通过同时解决无注意力和重复性冗余两类问题，克服现有方法仅针对单一类型冗余（如无注意力或重复）的局限性。
• 设计一种轻量级、无参数的机制，实现跨层和输入实例的token剪枝与池化自适应集成。
• 在无需微调或架构重构的情况下，实现ViTs的高效推理与更快训练。
• 与最先进token压缩方法相比，实现更优的准确率-FLOPs权衡。

提出的方法

• 提出一种自适应的token压缩模块，根据每层中token重要性得分的分布，动态选择剪枝或池化。
• 基于token得分的方差设计启发式决策准则：高方差（重要性差异显著）时选择剪枝，低方差（重要性相似）时选择池化。
• 将模块集成到标准transformer块中，不引入任何可学习参数，确保与预训练模型的兼容性。
• 在不同层中差异化应用该自适应策略：浅层偏好池化（相似性高），深层偏好剪枝（重要性差异显著）。
• 根据输入实例的局部token得分分布动态调整策略，实现实例感知的压缩。
• 设置可学习阈值τ以控制剪枝与池化的偏好，通过微调优化，或在现成场景中经验性设定。

实验结果

研究问题

• RQ1 在视觉Transformer中，联合应用token剪枝与池化是否能比单独使用任一方法获得更优的效率-准确率权衡？
• RQ2 是否存在与层深度相关的token冗余模式，使得在不同深度采用不同压缩策略（剪枝vs.池化）具有合理性？
• RQ3 一种实例感知的、无训练参数的策略选择机制，能否有效平衡多样化输入中的无注意力与重复性冗余？
• RQ4 决策阈值τ的选择如何影响自适应压缩框架的性能与鲁棒性？
• RQ5 所提方法是否可在无需微调的情况下，在ImageNet等标准基准上实现最小准确率损失的部署？

主要发现

• PPT在ImageNet上对DeiT-S实现FLOPs减少超过37%，吞吐量提升超过45%，且准确率无下降。
• 自适应选择策略优于固定规则策略（如浅层池化、深层剪枝）和随机选择，尤其在低FLOP水平下表现更优。
• 在所提框架中，仅使用token池化优于仅使用token剪枝，但通过自适应选择结合两者可实现最佳整体性能。
• 通过微调确定的最优决策阈值τ为6×10⁻⁵，而在现成设置中为7×10⁻⁵，能有效平衡剪枝与池化的偏好。
• 该方法在不同ViT架构上均保持强泛化能力，并在ImageNet上实现了最先进水平的效率-准确率权衡。
• 该框架在推理加速方面有效，且可在无额外参数的情况下实现更快训练，准确率损失极小。