原論文
Axial-deeplab: Standalone axial-attention for panoptic segmentation.
https://arxiv.org/pdf/2003.07853.pdf
結論
軸別にSelf-Attentionすることで計算量を削減できた．
概要
画像認識モデルのCNNは，カーネルの範囲の情報を用いて計算するため，受容野が小さく，遠くにある画素の情報を考慮できない．ViTは，CNNの受容野を解決するように，画像全体の情報を考慮して計算が可能になった．しかし，ViTの構造上，CNNと比較して膨大な計算コストが必要になる．ViTの派生モデルで，Swinは，画像全体の受容野を制限して計算量を抑えた．けど，これは根本的な解決になっていない．
グローバルな認識ができて，さらに計算量を抑えたい．そこで，2次元で動作するViTを2つの1次元SAに因数分解することで，計算効率が良いViTを提案する．また，これに付随して，位置に敏感な構造も提案する．
モデル構造
axial-attention block
axial-attention blockを下図に示す．入力された($H \times W \times 256$)の3次元の特徴量を1x1Convでチャンネル数を圧縮し（$H \times W \times 128$）の特徴量を得る．これは，その後に処理されるAttentionの計算量を削減するためにやる．その後，縦方向のAttentionと横方向のAttentionをして，1x1Convでチャンネル数を$256$に戻して，入力された特徴量に要素和する．
例えば，ResNetの畳み込みblockとこのaxial-attention blockを入れ替えることで，いい精度が達成する．
Multi-Head Attention （Height-Axis，Width-Axis）
下図に1headのAttention構造を示す．
$H$が特徴量の縦の長さ，$W$が横の長さ，$8$や$16$はチャンネル数を示す．
出力が（$H \times W \times 16$）なので，この出力が8種類（8head）あって，上図の（$H \times W \times 16$)$ \times 8$でConcatして（$H \times W \times 128$）に戻ってくる．
つまり，1headに（$H \times W \times 16$）が入力され，QとKは1x1Convで（$H \times W \times 8$），Vは（$H \times W \times 16$）の特徴量を得る．
軸のAttention
下図の左に示す．QKの特徴量をから，Hを引き離して（$W \times 8$）の部分で内積計算する．$8$の部分が相殺されて，（$W \times W$）になる．これとVの（$W \times 16$）で$W$が相殺して，（$W \times 16$）になる．
※Hの部分を外して考えてみるとわかりやすい．
Height-Axisの場合は，Hを外して考える．
Width-Axisの場合は，Wを外して考える．
軸＋位置情報を考慮したAttention
下図の右に示す．
主に，$r$で示される位置情報をかけることで位置を考慮した計算ができるようになる．Swinでも位置情報を考慮した構造を提案しているが，SwinはQのみに位置情報を付与しており，それは十分でない．なので，たくさんつけました．
QKの内積と，Qの位置情報，Kの位置情報の要素和して，Vと内積して，Vの位置情報を足して出力って感じ．
実験
1x1convで計算量が減って，軸別でAttentionをしたので，計算量が小さくなったのに，精度は向上した．
まとめ
今回は，軸別にSelf-Attentionするモデル【Axial-ResNet】について解説した．ViTのSAの計算量が大きいので，削減したいけど，Swinのようにローカルになるのは違うよね．グローバルで削減できたのは良い．