[论文解读] MSG-Transformer: Exchanging Local Spatial Information by Manipulating Messenger Tokens

Jiemin Fang, Lingxi Xie+4 more|arXiv (Cornell University)|May 31, 2021

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

一句话总结

本文提出MSG-Transformer，一种视觉Transformer模型，通过轻量级消息传递（MSG）令牌增强局部空间信息交换，这些令牌汇总并转发非重叠局部窗口之间的特征。通过用MSG令牌操作（尤其是混洗操作）替代直接的跨窗口注意力机制，该方法在ImageNet（84.0% Top-1准确率）和MS-COCO（52.8 mAP）上实现了最先进性能，且推理速度更快，尤其在CPU上表现更优，原因在于计算开销减少和实现更简洁。

ABSTRACT

Transformers have offered a new methodology of designing neural networks for visual recognition. Compared to convolutional networks, Transformers enjoy the ability of referring to global features at each stage, yet the attention module brings higher computational overhead that obstructs the application of Transformers to process high-resolution visual data. This paper aims to alleviate the conflict between efficiency and flexibility, for which we propose a specialized token for each region that serves as a messenger (MSG). Hence, by manipulating these MSG tokens, one can flexibly exchange visual information across regions and the computational complexity is reduced. We then integrate the MSG token into a multi-scale architecture named MSG-Transformer. In standard image classification and object detection, MSG-Transformer achieves competitive performance and the inference on both GPU and CPU is accelerated. Code is available at https://github.com/hustvl/MSG-Transformer.

研究动机与目标

• 为解决标准Transformer在视觉任务中计算成本过高的问题，尤其是由于其二次方复杂度的注意力机制。
• 克服现有局部注意力方法（如Swin Transformer和HaloNet）存在的内存浪费、实现复杂或通信受限等问题。
• 设计一种模块化、灵活的局部特征交换框架，将信息传播与图像块级别的注意力计算解耦。
• 通过最小化冗余特征传递和数据复制，实现更低的FLOPs和更快的推理速度，尤其在CPU上表现更优。
• 通过分层的非重叠窗口设计与集中式信息路由机制（通过MSG令牌实现），实现高效的多尺度特征学习。

提出的方法

• 为每个局部窗口引入一个轻量级消息传递（MSG）令牌，用于汇总并传播空间特征，替代直接的跨窗口注意力机制。
• 采用非重叠局部窗口设计，以减少计算和内存开销，避免HaloNet中冗余特征复制或Swin Transformer中窗口移位带来的问题。
• 实现一种混洗操作，将MSG令牌的特征重新分配至不同空间位置，以极低计算成本实现跨窗口信息交换。
• 将MSG令牌整合到多尺度架构（MSG-Transformer）中，其中每个阶段处理不同分辨率的特征。
• 设计网络结构，使图像块令牌仅在其本地窗口内进行注意力计算，而MSG令牌负责跨窗口通信，从而减轻图像块级别注意力的负担。
• 采用分层设计，包含多个阶段，其中MSG令牌在分辨率逐步提升的特征图之间更新并共享，实现多层级特征学习。

实验结果

研究问题

• RQ1 中心化、轻量级的令牌机制（MSG）是否能有效替代直接的跨窗口注意力机制，从而降低视觉Transformer中的计算成本？
• RQ2 与Swin Transformer相比，所提出的MSG-Transformer在准确率、FLOPs和推理速度（尤其在CPU上）方面表现如何？
• RQ3 不同混洗区域大小对MSG-Transformer性能和效率的影响是什么？
• RQ4 MSG令牌是否能灵活控制信息交换模式，而无需增加架构复杂度？
• RQ5 MSG-Transformer在图像分类和目标检测等多样化视觉任务中是否保持强大性能？

主要发现

• MSG-Transformer在ImageNet图像分类任务中达到84.0% Top-1准确率，优于近期的Swin Transformer。
• 在MS-COCO目标检测任务中，MSG-Transformer达到52.8 mAP，超过Swin Transformer的性能表现。
• 该模型在推理速度方面具有显著优势，尤其在CPU上，得益于特征传递和内存开销的减少。
• 混洗操作以可忽略的计算成本实现了有效的跨窗口信息交换，且增大混洗区域尺寸可进一步提升准确率。
• 更深更窄的网络结构（64维，[2,4,12,4]个块）在MSG-Transformer中实现了比Swin Transformer更优的准确率-计算量权衡。
• 注意力图可视化结果表明，深层中的MSG令牌聚焦于与物体相关的区域，表明尽管采用局部注意力机制，仍具备有效的空间建模能力。