[論文レビュー] MaxViT: Multi-Axis Vision Transformer

Zhengzhong Tu, Hossein Talebi+5 more|arXiv (Cornell University)|Apr 4, 2022

Visual Attention and Saliency Detection被引用数 26

ひとこと要約

MaxViT は、ブロック状ローカル注意と拡張グローバル注意を組み合わせ、畳み込みと組み合わせてスケーラブルな階層型のビジョンバックボーンを形成し、ImageNetとCOCOで最先端の結果を達成します。

ABSTRACT

Transformers have recently gained significant attention in the computer vision community. However, the lack of scalability of self-attention mechanisms with respect to image size has limited their wide adoption in state-of-the-art vision backbones. In this paper we introduce an efficient and scalable attention model we call multi-axis attention, which consists of two aspects: blocked local and dilated global attention. These design choices allow global-local spatial interactions on arbitrary input resolutions with only linear complexity. We also present a new architectural element by effectively blending our proposed attention model with convolutions, and accordingly propose a simple hierarchical vision backbone, dubbed MaxViT, by simply repeating the basic building block over multiple stages. Notably, MaxViT is able to ''see'' globally throughout the entire network, even in earlier, high-resolution stages. We demonstrate the effectiveness of our model on a broad spectrum of vision tasks. On image classification, MaxViT achieves state-of-the-art performance under various settings: without extra data, MaxViT attains 86.5% ImageNet-1K top-1 accuracy; with ImageNet-21K pre-training, our model achieves 88.7% top-1 accuracy. For downstream tasks, MaxViT as a backbone delivers favorable performance on object detection as well as visual aesthetic assessment. We also show that our proposed model expresses strong generative modeling capability on ImageNet, demonstrating the superior potential of MaxViT blocks as a universal vision module. The source code and trained models will be available at https://github.com/google-research/maxvit.

研究の動機と目的

• 局所とグローバルな相互作用の両方を捉えるスケーラブルなビジョンアーキテクチャを動機づける。
• ブロックローカルとグリッドグローバルなアテンションを畳み込みと組み合わせたトランスフォーマーブロックを開発する。
• MaxViT ブロックを各ステージで繰り返すことで、シンプルな階層バックボーンを構築する。
• 分類、検出、美学、生成タスクの各領域で強力な性能を示す。

提案手法

• 全注意をブロックローカルとグリッドグローバル注意に分解し、線形の計算量で実行するマルチアクシス自己注意（Max-SA）を導入する。
• Max-SA を MBConv ブロックと SE でブレンドし、MBConv を条件付き位置エンコーディングとして使用して一般化を向上させる。
• 4つのステージ（S0-S4）にわたって繰り返し MaxViT ブロックを積み上げて、階層バックボーン MaxViT を構築する。
• ステージごとのブロック数とチャネルサイズを変えた構成可能なバリアント（MaxViT-T、-S、-B、-L、-XL）を提供する。
• Max-SA が Swin に類似したアテンションのドロップインとして、同じパラメータ数と FLOPs で挿入可能であり、全ステージでグローバルな相互作用を提供することを示す。

実験結果

リサーチクエスチョン

• RQ1 高解像度のビジョンタスクにおいて、マルチアクシスアテンション（局所ブロック + グローバルグリッド）は線形計算量でグローバルコンテキストを提供できるか。
• RQ2 単純な階層型バックボーンで Max-SA と畳み込みを組み合わせることで、既存の Vision Transformer やハイブリッドより精度と効率が向上するか。
• RQ3 様々なアーキテクチャの選択（ブロック順序、逐次 vs 並列のアテンション、MBConv の包含、垂直レイアウト）がビジョンタスク全体の性能にどう影響するか。
• RQ4 MaxViT はデータ（ImageNet-1K、ImageNet-21K、JFT-300M）にどの程度スケールし、検出や美学などの下流タスクにどの程度適用可能か。
• RQ5 MaxViT は画像生成設定で強力な生成性能を発揮できるか。

主な発見

• MaxViT は設定全体で ImageNet-1K top-1 の最先端精度を達成し、MaxViT-L の 512x512 ファインチューニングで 86.7%、追加データなしで 224x224 の MaxViT-L は 85.17% を達成。
• ImageNet-21K のプリトレーニングでは、MaxViT-B が 88.38% の top-1 精度、MaxViT-XL が 512x512 で 88.70% に達し、同等または大きいサイズの従来モデルを上回る。
• JFT-300M規模データでは、MaxViT-XL が 89.53% top-1 精度を達成し、巨大データセットへの強いスケーリングを示す。
• COCO の物体検出/インスタンスセグメンテーションでは、MaxViT バックボーンはさまざまなサイズで Swin、ConvNeXt、UViT を上回り、基礎レベルで特筆すべき差を示す（例: MaxViT-S が同等の FLOPs で Swin-B および UViT-B を上回る）。
• 画像美学（AVA）では、MaxViT-T が競争力のある PLCC/SRCC スコアを示し、従来法より解像度が高いほど改善する。
• 無条件の 128x128 画像生成では、MaxViT は HiT や他のベースラインより少ないパラメータでより良い FID/IS を示す。