[论文解读] Lightweight Vision Transformer with Cross Feature Attention

Youpeng Zhao, Huadong Tang+3 more|arXiv (Cornell University)|Jul 15, 2022

Advanced Neural Network Applications被引用 8

一句话总结

本文提出了一种轻量级自注意力机制——交叉特征注意力（XFA），通过在特征维度而非 token 序列上操作，将计算复杂度从二次方降低为线性。该方法构建了混合 CNN-ViT 模型 XFormer，在 ImageNet-1K（78.5% top-1 准确率）、COCO 目标检测（33.2 mAP）和 Cityscapes 语义分割（78.5 mIoU）任务上达到最先进性能，且参数量和 FLOPs 更少，优于 MobileNetV2 和 DeiT。

ABSTRACT

Recent advances in vision transformers (ViTs) have achieved great performance in visual recognition tasks. Convolutional neural networks (CNNs) exploit spatial inductive bias to learn visual representations, but these networks are spatially local. ViTs can learn global representations with their self-attention mechanism, but they are usually heavy-weight and unsuitable for mobile devices. In this paper, we propose cross feature attention (XFA) to bring down computation cost for transformers, and combine efficient mobile CNNs to form a novel efficient light-weight CNN-ViT hybrid model, XFormer, which can serve as a general-purpose backbone to learn both global and local representation. Experimental results show that XFormer outperforms numerous CNN and ViT-based models across different tasks and datasets. On ImageNet1K dataset, XFormer achieves top-1 accuracy of 78.5% with 5.5 million parameters, which is 2.2% and 6.3% more accurate than EfficientNet-B0 (CNN-based) and DeiT (ViT-based) for similar number of parameters. Our model also performs well when transferring to object detection and semantic segmentation tasks. On MS COCO dataset, XFormer exceeds MobileNetV2 by 10.5 AP (22.7 -> 33.2 AP) in YOLOv3 framework with only 6.3M parameters and 3.8G FLOPs. On Cityscapes dataset, with only a simple all-MLP decoder, XFormer achieves mIoU of 78.5 and FPS of 15.3, surpassing state-of-the-art lightweight segmentation networks.

研究动机与目标

• 解决视觉变换器（ViTs）在移动端和轻量化场景下的高计算成本问题。
• 克服标准自注意力机制随序列长度呈二次方增长的局限性，从而提升在边缘设备上的部署可行性。
• 设计一种混合架构，结合 CNN 的局部归纳偏置与高效 ViTs 的全局建模能力。
• 开发一种通用、高效的主干网络，在分类、检测和分割等任务中均表现优异。

提出的方法

• 提出一种新型注意力机制——交叉特征注意力（XFA），在特征维度而非 token 维度上计算注意力，使复杂度随序列长度呈线性增长。
• 用可学习的缩放参数替代标准自注意力中的 softmax 操作，以稳定训练并提升优化性能。
• 设计一种分层的 XFormer 架构，由高效的 CNN 块与 XFA 变压器块交替组成，以平衡局部与全局特征学习。
• 采用简单的全 MLP 解码器进行语义分割，在保持高性能的同时最小化额外计算量。
• 在 ImageNet-1K 上预训练时使用知识蒸馏与数据增强技术，以提升泛化能力。
• 将模型部署于 YOLOv3 和 SegFormer 框架中，评估其在检测与分割任务中的迁移性与效率。

实验结果

研究问题

• RQ1 能否重新设计一种自注意力机制，在保持视觉变换器全局建模能力的同时实现线性复杂度？
• RQ2 如何在轻量化架构中有效结合 CNN 与 ViT，以提升移动端视觉任务的性能？
• RQ3 混合 CNN-ViT 模型是否能在多个视觉基准上，同时超越最先进的 CNN 与 ViT，在准确率、参数效率与 FLOPs 方面表现更优？
• RQ4 所提出的 XFormer 主干网络是否能在无需复杂解码器的情况下，良好泛化至下游任务（如目标检测与语义分割）？

主要发现

• 在 ImageNet-1K 上，XFormer 仅用 550 万个参数与 17 亿 FLOPs 即达到 78.5% 的 top-1 准确率，优于参数量相近的 EfficientNet-B0 2.2%，也优于 DeiT 6.3%。
• 在 MS COCO 目标检测任务中，XFormer 将 YOLOv3 的性能提升至 33.2 mAP，较 MobileNetV2 提升 10.5 个百分点，同时参数量减少 56%，FLOPs 减少 22%。
• 在 Cityscapes 语义分割任务中，XFormer 使用简单全 MLP 解码器即达到 78.5 mIoU，优于基于 MobileViT 的 SegFormer 2.9 mIoU，也优于最先进模型 LVT 3.6 mIoU。
• XFormer 保持了高推理速度，在 Cityscapes 上达到 15.3 FPS，且在准确率与速度上均优于基于 ViT 的模型（如 YOLOS 与 MobileViT）。
• 该模型展现出强大的迁移能力，即使仅作最小程度的架构修改，作为主干网络也能显著提升检测与分割任务的性能。
• 消融研究证实，XFA 的线性复杂度与可学习缩放因子在低资源环境下对效率与性能的提升具有关键贡献。