[論文レビュー] Lightweight Vision Transformer with Cross Feature Attention

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

Advanced Neural Network Applications被引用数 8

ひとこと要約

本論文は、トークン系列ではなく特徴マップの次元で動作させることで、計算複雑度を二次関数的から線形に削減する軽量な自己注意機構であるCross Feature Attention (XFA)を提案する。この手法により、ハイブリッドCNN-ViTモデルXFormerが構築され、ImageNet-1K（78.5% top-1精度）、COCOオブジェクト検出（33.2 mAP）、Cityscapesセマンティックセグメンテーション（78.5 mIoU）で最先端の性能を達成し、パrameter数と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の局所的なインダクティブバイアスと効率的なViTのグローバルモデリング能力を組み合わせたハイブリッドアーキテクチャを設計すること。
• 分類、検出、セグメンテーションの複数のタスクに普遍的に適応し、高い性能を示す効率的なバックボーンを構築すること。

提案手法

• 自己注意機構の計算をトークン次元ではなく特徴マップ次元で行うことで、系列長に比例する線形の複雑度を実現する、新規の注意機構であるCross Feature Attention (XFA)を提案。
• 従来のソフトマックス演算を、学習可能なスケーリングパラメータに置き換えることで、訓練の安定性と最適化の改善を図る。
• 局所的特徴学習とグローバル特徴学習のバランスを取るために、効率的なCNNブロックとXFAトランスフォーマーブロックを交互に配置した階層的XFormerアーキテクチャを導入。
• 追加計算を最小限に抑えるために、セマンティックセグメンテーション用に単純なすべてMLPデコーダーを採用。
• ImageNet-1Kでの事前学習段階で知識蒸留とデータオーグメンテーション技術を用いることで、一般化性能を向上。
• オブジェクト検出とセグメンテーションタスクにおける転移性と効率性を評価するため、YOLOv3およびSegFormerフレームワークにモデルを統合。

実験結果

リサーチクエスチョン

• RQ1 ビジョントランスフォーマーにおいて、グローバルモデリング能力を保持しつつ、計算複雑度を線形に抑える自己注意機構を再設計することは可能か？
• RQ2 CNNとViTを効果的に統合した軽量アーキテクチャを設計することで、モバイルビジョンタスクの性能を向上させることは可能か？
• RQ3 ハイブリッドCNN-ViTモデルが、複数のビジョンベンチマークにおいて、精度、パrameter効率、FLOPsの観点で、最先端のCNNおよびViTを上回ることは可能か？
• 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におけるセマンティックセグメンテーションでは、単純なすべてMLPデコーダーを用いて78.5 mIoUを達成し、MobileViTベースのSegFormerを2.9 mIoU、最先端のLVTを3.6 mIoU上回った。
• XFormerは高い推論速度を維持しており、Cityscapesでは15.3 FPSを達成し、YOLOS や MobileViT などのViTベースモデルを精度と速度の両面で上回った。
• モデルは優れた転送性を示し、最小限のアーキテクチャ的変更で、検出およびセグメンテーションタスクの性能を顕著に向上させた。
• アブレーションスタディの結果、XFAの線形複雑度と学習可能なスケーリング因子が、特にリソースが限られた環境下で、効率性と性能の両面で寄与していることが確認された。