[论文解读] Shunted Self-Attention via Multi-Scale Token Aggregation

Sucheng Ren, Daquan Zhou+3 more|arXiv (Cornell University)|Nov 30, 2021

Advanced Neural Network Applications参考文献 34被引用 10

一句话总结

本文提出了一种新型多尺度标记聚合机制——分流自注意力（Shunted Self-Attention, SSA），通过选择性地合并标记，使视觉变换器（Vision Transformers）能够在单个注意力层内同时建模细粒度与粗粒度特征。SSA 在降低计算成本的同时，提升了图像分类、目标检测和语义分割任务的性能，在模型参数量和浮点运算量（FLOPs）仅为当前最先进模型 Focal Transformer 一半的情况下，ImageNet 上达到 84.0% 的 Top-1 准确率，并在 COCO 上超越其 1.3 mAP，在 ADE20K 上超越其 2.9 mIOU，且在相似效率约束下表现更优。

ABSTRACT

Recent Vision Transformer~(ViT) models have demonstrated encouraging results across various computer vision tasks, thanks to their competence in modeling long-range dependencies of image patches or tokens via self-attention. These models, however, usually designate the similar receptive fields of each token feature within each layer. Such a constraint inevitably limits the ability of each self-attention layer in capturing multi-scale features, thereby leading to performance degradation in handling images with multiple objects of different scales. To address this issue, we propose a novel and generic strategy, termed shunted self-attention~(SSA), that allows ViTs to model the attentions at hybrid scales per attention layer. The key idea of SSA is to inject heterogeneous receptive field sizes into tokens: before computing the self-attention matrix, it selectively merges tokens to represent larger object features while keeping certain tokens to preserve fine-grained features. This novel merging scheme enables the self-attention to learn relationships between objects with different sizes and simultaneously reduces the token numbers and the computational cost. Extensive experiments across various tasks demonstrate the superiority of SSA. Specifically, the SSA-based transformer achieves 84.0\% Top-1 accuracy and outperforms the state-of-the-art Focal Transformer on ImageNet with only half of the model size and computation cost, and surpasses Focal Transformer by 1.3 mAP on COCO and 2.9 mIOU on ADE20K under similar parameter and computation cost. Code has been released at https://github.com/OliverRensu/Shunted-Transformer.

研究动机与目标

• 为解决视觉变换器在注意力层中因所有标记具有统一且静态的感受野而导致难以捕捉多尺度特征的问题。
• 在保留图像中不同尺度物体的细粒度细节的同时，降低计算成本与显存消耗。
• 使单个自注意力层能够同时关注大对象与小对象，从而提升下游视觉任务的性能。
• 设计一种通用、即插即用的机制，无需对现有 ViT 架构进行大规模结构调整即可增强其性能。

提出的方法

• 引入一种分流自注意力机制，将注意力计算分为两条分支：一条用于聚合后的粗粒度标记，另一条用于未聚合的细粒度标记。
• 应用可学习的标记聚合函数，基于特征相似性与空间上下文信息，选择性地合并空间相邻的标记，从而为大对象形成更大的感受野。
• 为粗粒度与细粒度特征分别维护独立的注意力图，使模型能够在同一层内学习不同尺度之间的关系。
• 集成一种面向细节的前馈网络，通过引入邻近标记的上下文信息，增强局部特征表示，提升特征精炼能力。
• 采用多尺度的补丁嵌入头，以保留空间分辨率并减少初始特征提取过程中的信息损失。
• 采用混合计算策略：聚合后的标记可减少序列长度与 FLOPs，而未聚合的标记则保留高分辨率细节。

实验结果

研究问题

• RQ1 单个自注意力层是否能在不依赖分层特征图的情况下，有效建模多尺度特征？
• RQ2 选择性标记合并是否能在降低计算成本的同时，提升目标检测与语义分割等密集预测任务的性能？
• RQ3 统一的注意力机制是否能超越采用独立分支或下采样策略处理不同尺度的模型？
• RQ4 所提出的标记聚合函数与标准卷积或线性下采样相比，在准确率与效率方面表现如何？
• RQ5 在前馈层中引入局部细节信息在多大程度上能改善基于 ViT 模型的特征表示？

主要发现

• 分流视觉变换器（Shunted Transformer）在 ImageNet 上达到 84.0% 的 Top-1 准确率，优于当前最先进模型 Focal Transformer，且参数量减少 50%，FLOPs 降低 50%。
• 在 COCO 目标检测任务中，该模型在相似参数量与计算预算下，mAP 超过 Focal Transformer 1.3 个百分点。
• 在 ADE20K 语义分割任务中，该模型在相近效率约束下，mIOU 比 Focal Transformer 提升 2.9 个百分点。
• 结合 SegFormer 框架时，所提出的主干网络在 ADE20K 上达到 48.3% 的 mIoU，优于 MiT-B2 的 46.5%，且仅使用 25.1M 参数（MiT-B2 为 27.5M）。
• 与标准前馈网络相比，面向细节的前馈层在 ViT 中将 Top-1 准确率提升 0.2%，在分流视觉变换器中提升 0.3%，证明其在保留局部细节方面的有效性。
• 所提出的补丁嵌入头相比非重叠补丁嵌入方式，将 ImageNet 准确率提升 1.4%，相比重叠补丁嵌入方式提升 0.3%，且 FLOPs 增加极少。