MaxViT:マルチアキスビジョンTransformer

最近、Transformerはコンピュータビジョン分野において大きな注目を集めている。しかし、自己注意(self-attention)機構が画像サイズに対してスケーラビリティに欠けることから、最先端のビジョンバックボーンにおける広範な採用が制限されてきた。本論文では、効率的かつスケーラブルな注目モデルとして「マルチアキス・アテンション(multi-axis attention)」を提案する。このモデルは、2つの側面から構成されている:ブロッキングされた局所的アテンションと拡張されたグローバルアテンション。これらの設計により、任意の入力解像度においてグローバルとローカルの空間的相互作用を実現しつつ、計算量は線形の複雑さに抑えることが可能となる。さらに、提案したアテンションモデルを畳み込みと効果的に融合する新たなアーキテクチャ要素を導入し、複数段階にわたり基本ブロックを単純に繰り返すことで、シンプルな階層型ビジョンバックボーン「MaxViT(Maximal Vision Transformer)」を構築した。特に注目すべきは、MaxViTがネットワーク全体にわたり、初期の高解像度段階でもグローバルな視覚的洞察を維持できることである。我々は、幅広いビジョンタスクにおいて本モデルの有効性を実証した。画像分類において、追加データなしの設定下でMaxViTはImageNet-1Kデータセットでトップ1精度86.5%を達成。ImageNet-21Kでの事前学習を用いた場合、トップ1精度は88.7%に達した。下流タスクにおいても、物体検出および視覚的美的評価の分野で、優れた性能を示した。また、ImageNet上で生成モデルとしての強力な表現能力を示すことで、MaxViTブロックが汎用的なビジョンモジュールとしての優れた潜在能力を持つことを示した。本研究のソースコードおよび学習済みモデルは、https://github.com/google-research/maxvit にて公開される予定である。