[論文メモ] MixFormer: Mixing Features across Windows and Dimensions

arxiv.org

CVPR2022 Oral

window baseのattentionとdepth-wise convで双方向インタラクションすることで計算量を抑えつつlocal globalの情報を扱えるようにした。
f:id:Ninhydrin:20220411085746p:plain

Vision Transformer(ViT)の一つの成功例としてSwin Transformerがある。
Swin Transformerはオーバーラップ無しのwindowを切り出し、その中でattentionを行うというwindow baseなattention(local window self-attention)を採用している。
local window self-attentionは受容野を絞ることで計算量を抑えまた画像タスクの帰納バイアスとなる。しかし受容野を絞ったことで元のself-attentionより能力が制限されている。
これを改善するためにdepth-wise convと組み合わせる方法も提案されている(HRFormerとか)。これはlocal window self-attentionとdepth-wise convを続けて行うことでlocal globalを組み合わせているが、これだと情報がうまく混ざらず、またのlocal window self-attentionとdepth-wise convそれぞれの弱点を無視しているため特徴表現があまり良くない。
これらをどうにかしたいというお気持ち。

手法

local window self-attentionとdepth-wise convを並列し、かつ双方向インタラクションを行うMixing Blockを提案する(図1)。

並列化はそのままの通りでlocal window self-attentionとdepth-wise convを並列で行い、それらの出力をconcat後 feed-forward network(FFN)に通す。

local window self-attentionのサイズは7x7、depth-wise convのカーネルサイズは3x3。これらは実験的に決定。

次に双方向インタラクションについて(図1の水色の矢印)。

local window self-attentionはチャンネル方向の重みを共有して空間方向の位置で動的に重みを決定するためチャンネル方向のモデリング能力が低い。
逆にdepth-wise convはチャンネルに注目して空間方向の重みを共有している。
そこでこれらを双方向にインタラクション(bi-directional interactions)することでお互いの弱点を補う。
f:id:Ninhydrin:20220411092624p:plain

depth-wise convの出力はchannel interactionモジュールを通してlocal window self-attention(図だとW-Attention)に伝えられる。
local window self-attentionの出力はSpartial Interactionモジュールを通してdepth-wise convの出力に混ぜられる。

channel interactionモジュールは Squeeze-and-excitation NetのSEレイヤー。global average pooling(GAP)レイヤーで1x1の解像度に落として最後にsigmoidに通してチャンネル方向のattention的なものにする。そしてそれをself-attentionのvalue部分に掛ける。

spatial interactionモジュールはconv1x1を組み合わせてチャンネル方向を1にし、最後にsigmoidを通して空間方向のattention的なものにする。そしてdepth-wise convの出力に掛ける。

これがMixing Blockでこれを使ったViTをMixFormerとする(図3)。
f:id:Ninhydrin:20220411094340p:plain

例のごとく複数のサイズを用意(表2)。
f:id:Ninhydrin:20220411094432p:plain

実験・結果

224x224入力のImageNet-1Kのクラス分類。AdamWで300epochの学習。結果が表3。
f:id:Ninhydrin:20220411094535p:plain

これの前に読んだMaxViTとモデルサイズ等で良さげなペアがなく比較は難しかった。MixFormerのほうがFLOPsが小さそう。

Object DetectionやSemantic Segmentationのbackboneとしたときの比較。
f:id:Ninhydrin:20220411095254p:plain
f:id:Ninhydrin:20220411095329p:plain

Ablations。
並列化とインタラクションの影響。並列化の効果が高そう?
f:id:Ninhydrin:20220411095447p:plain

depth-wise convのカーネルサイズの違い。
f:id:Ninhydrin:20220411095539p:plain

アーキテクチャデザイン。
f:id:Ninhydrin:20220411095638p:plain


他のアーキテクチャにMixing Blockを適用。
f:id:Ninhydrin:20220411095747p:plain

所感

最近では「Swin Transformerのlocal window self-attentionは局所的な特徴しか扱えないのでなんとかしたい」から「local window self-attentionとグローバル情報をどうやって組み合わせるか」という問題設定に変化している。
前回読んだMaxViTはchannel shufflerの逆操作で、今回のMixFormerでは並列したdepth-wise convとの双方向インタラクション。
MixFormerのほうが1ブロックに付きself-attentionが1回少ない分計算量が少なそう。その代わり精度がわずかに劣っている印象。
正直self-attention使わなくても、pixel shufflerの逆操作とconvで十分な気がする。attention使わないからViTにはなれないけど。