[论文解读] 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 Sensors参考文献 44被引用 20

一句话总结

PSViT 通过引入可学习的标记池化和注意力共享，减轻了空间冗余和层间冗余，从而提升了特征表示能力，并实现了更优的速度-精度权衡。采用紧凑的搜索空间与 AutoML（SPOS），PSViT 在 ImageNet 上相较 DeiT 最高实现 6.6% 的 top-1 准确率提升，并在下游检测与分割任务中超越 ResNets。

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.

研究动机与目标

• 为解决因各层固定标记数量导致的视觉变换器中的空间冗余问题。
• 减少相邻层间因注意力模式相似而产生的注意力图冗余。
• 通过学习最优标记数量与注意力共享配置，实现灵活的计算分配。
• 提升特征表示能力，并在视觉变换器架构中实现更优的速度-精度权衡。
• 利用 AutoML 开发一种系统化、数据驱动的 ViT 架构设计方法。

提出的方法

• 引入标记池化机制，以在不同网络阶段动态减少标记数量，适应特征复杂度。
• 在相邻的 Transformer 层之间应用注意力共享，复用相似的注意力图，降低计算量。
• 构建一个紧凑的搜索空间，将标记池化位置与注意力共享选择作为超参数。
• 使用 SPOS（Single Path One-Shot）进行神经架构搜索，联合优化标记数量与注意力共享模式。
• 将标记池化与注意力共享视为可学习的超参数，通过可微架构搜索端到端训练。
• 采用 1D 与 2D 标记池化变体，其中 2D 池化在图像任务中能更好地保留空间结构。

实验结果

研究问题

• RQ1 在不同层间动态进行标记池化能否提升视觉变换器的特征表示能力？
• RQ2 相邻 ViT 层中的注意力图在多大程度上表现出冗余？这种冗余是否可被有效利用？
• RQ3 针对标记池化与注意力共享的紧凑搜索空间能否支持高效的 AutoML 架构搜索？
• RQ4 联合优化标记数量与注意力共享是否能带来更优的准确率与效率权衡？
• RQ5 PSViT 在目标检测与实例分割等下游视觉任务中的迁移性能如何？

主要发现

• PSViT 在 ImageNet 上相较 DeiT 最高实现 6.6% 的 top-1 准确率提升，展现出显著的性能增益。
• PSViT-2D-Tiny 在 COCO 目标检测任务中达到 40.8% mAP，在实例分割任务中达到 37.7% mAP，优于 ResNet18，甚至超过 ResNet101。
• 2D 标记池化变体优于 1D 版本，表明保留空间结构能提升特征质量。
• 最后一层 Transformer 避免注意力共享，因其独立的自注意力机制对最终分类最为有利。
• 网络的前几层最受益于注意力共享，因其计算成本高且特征平滑。
• 搜索空间设计支持高效的 AutoML 优化，PSViT-2D-Tiny 在极少人工干预下即实现优异性能。