[論文レビュー] SegViT: Semantic Segmentation with Plain Vision Transformers

Bowen Zhang, Zhi Tian+5 more|arXiv (Cornell University)|Oct 12, 2022

Advanced Neural Network Applications被引用数 75

ひとこと要約

SegViT は Attention-to-Mask (ATM) デコーダを導入し、Plain Vision Transformer を用いたセマンティックセグメンテーションを実現。Shrunk バックボーン設計により計算量を削減しつつ、最先端または競合的な結果を達成。

ABSTRACT

We explore the capability of plain Vision Transformers (ViTs) for semantic segmentation and propose the SegVit. Previous ViT-based segmentation networks usually learn a pixel-level representation from the output of the ViT. Differently, we make use of the fundamental component -- attention mechanism, to generate masks for semantic segmentation. Specifically, we propose the Attention-to-Mask (ATM) module, in which the similarity maps between a set of learnable class tokens and the spatial feature maps are transferred to the segmentation masks. Experiments show that our proposed SegVit using the ATM module outperforms its counterparts using the plain ViT backbone on the ADE20K dataset and achieves new state-of-the-art performance on COCO-Stuff-10K and PASCAL-Context datasets. Furthermore, to reduce the computational cost of the ViT backbone, we propose query-based down-sampling (QD) and query-based up-sampling (QU) to build a Shrunk structure. With the proposed Shrunk structure, the model can save up to $40\%$ computations while maintaining competitive performance.

研究の動機と目的

• プレーンな Vision Transformers (ViTs) をセマンティックセグメンテーションに活用する可能性を探る。
• Attention-to-Mask (ATM) モジュールを提案し、アテンションマップからマスクを導出する。
• ViT 層に ATM をカスケードさせて多層情報を統合し、セグメンテーションを実現する。
• 計算量を削減する Shrunk バックボーン（クエリベースのダウンサンプリングとアップサンプリング）を導入する。
• ADE20K、COCO-Stuff-10K、PASCAL-Context で最先端または競争力のある結果を示す。

提案手法

• クラス・トークンのクエリを定義し、バックボーンの特徴マップとクロスアテンションを用いて、類似度マップのシグモイドを介してクラスごとのマスクを生成する。
• 更新されたクラス・トークンに対して線形変換と Softmax を適用してクラス予測を計算する。
• 複数の ViT 層からの ATM 出力を統合して最終的なセグメンテーション予測を形成する。
• 計算量を節約する Shrunk（QD: クエリベースのダウンサンプリング、QU: クエリベースのアップサンプリング）を導入し、GFLOPs を最大約 40% 程度低減する。
• マルチターム損失で訓練する：L_overall = L_cls + lambda_focal L_IoU + lambda_dice L_dice、層間でクラス・トークンとマスクの両方を監視する。

実験結果

リサーチクエスチョン

• RQ1 プレーンな ViT バックボーンを、アテンション駆動のマスク推論アプローチで密なセマンティックセグメンテーションに効果的に使用できるか。
• RQ2 クロスアテンションの類似性マップをマスクとして活用することは、ViT の特徴上の画素ごとのデコーディングよりセグメンテーション品質を改善するか。
• RQ3 複数層 ATM のカスケードと Shrunk バックボーンは、ViT を用いたセグメンテーションで計算量を削減しつつ精度を損なわないか。

主な発見

• ATM を備えた SegViT は ADE20K で 55.2% mIoU（ViT-Large バックボーン）を達成し、Shrunk 版は 55.1% を達成し、コストを削減しつつ競争力を確保。
• ADE20K では ViT-Large を用いた SegViT はいくつかの ViT ベース手法を上回り、最先端に近いまたは上回る設定もある。
• SegViT-Shrunk は計算コストを約 40% 削減（373.5 GFLOPs vs 637.9 GFLOPs）し、性能低下は小幅。
• 複数層の ATM 入力は一貫した mIoU の利得を生み出す（例えば 3 層の使用で ADE20K で最大 +1.7%）。
• SegViT は PASCAL-Context（60 クラスで 65.3% mIoU）および COCO-Stuff-10K（ViT-Large で 50.3% mIoU）で強力な結果を示す。