[論文メモ] MaxViT: Multi-Axis Vision Transformer

Google Research

入力画像サイズにスケーラブルなVision Transformer(ViT)の提案

ViTはモデルの能力が高く過学習しやすい。それを抑えるため大量の学習データを必要とした。
Swin Transformerはwindow-baseにしてうまく制御したが、window-baseにしたことで大域へのattentionが失われた。
global attentionとlocal attentionの相互作用のためには高解像度(階層ネットワークの場合は入力に近い層)でglobal attentionを行う必要があるがattentionは $N^2$ の計算コストがかかる。どうにか軽量なglobal attentionをしたいというお気持ち。

手法

localとglobalのインタラクションを可能にするmulti-axis self-attention(Max-SA)をつかったMaxViTを提案。
MaxViTのアーキテクチャの全体像は図2参照。
f:id:Ninhydrin:20220406095459p:plain

はじめにConvしたあと(Stem)は基本Max-SAブロックを積むだけ。
Max-SAブロックはMBConv、BlockAttention、GridAttentionの3つのモジュールからなる。

MBConvはMobileNetv2で提案されたInverted Residual Block。Conv1x1でDepthwiseConv3x3とSqueeze-and-Excitationモジュールを挟んだもの。
MBConvを前に入れたのは、入れたほうが汎化性能が高くなり学習もしやすくなったから(実験による観測)。
またDepthwiseConvにはconditional position encoding (CPE)とみなすことができるので明示的なPEは入れてない。

メインはBlockAttention、GridAttention。
BlockAttentionがlocalなattentionでGridAttentionがglobalなattentionを担う。

まずBlockAttentionについて
画像 $X=\mathbb{R}^{H \times W \times C}$ を普通のViTなら $HW$ のシーケンスとしてattentionを行うが、BlockAttentionでは $(\frac{H}{P} \times \frac{W}{P}, P\times P, C)$ にreshapeして $P \times P$ の次元、つまり $P \times P$ のwindowに分割してそのwindow内でattentionを行う。これはSwin Transformerと同じ。
window内でのattentionになるのでlocal。

globalなattentionを担うGridAttentionについて。
画像 $X$ を $(G \times G, \frac{H}{G} \times \frac{W}{G}, C)$ にreshapeして $G \times G$ の次元でattentionを行う。
これ故にMulti axis。