[論文メモ] HRFormer: High-Resolution Transformer for Dense Prediction

arxiv.org

NeurIPS 2021


セグメンテーションやポーズ検出でも使える高解像度な特徴量をオリジナルのViTに比べ低コストで学習できるHigh-Resolution Transformer(HRT)を提案。

手法

HRNetに倣い高解像度から初めて徐々に低解像度の枝を作っていく。アーキテクチャの全体像は図2を参照。
f:id:Ninhydrin:20211022090701p:plain
以下細かい部分について。

Local-window self-attention(LWSA)

特徴マップをオーバーラップなしのパッチに切り出して、各パッチそれぞれでself-attentionを行う。
特徴マップ \textbf{X}P個のパッチに切り出す( \textbf{X} \rightarrow \{\textbf{X}_1,...,\textbf{X}_P\})。各パッチのサイズは K \times K
 p番目のパッチに対するmulti-head self-attention(MHSA)は
f:id:Ninhydrin:20211022092014p:plain
となる。添字付きWは重み、 h \in \{1,...,H\} H個あるヘッドの番号、 Dはチャンネル数、 \hat{\textbf{X}}_p p番目のMHSAの出力。
またT5モデルで利用されたrelative positional embeddingを利用。
各パッチにMHSAを適用した出力を元の形に戻して最終的なモジュールの出力 \textbf{X}^{\verb|MHSA|}を得る。
図1の(a)の部分がそれ。

f:id:Ninhydrin:20211022093312p:plain

FFN with depth-wise convolution

LWSAはローカルなself-attentionなので他のパッチとのインタラクションが無い。
そこで3x3のdepth-wise convとpoint-wise MLPs(1x1 conv)を導入する。
図1の(b)の部分がそれ。depth-wiseをpoint-wiseで挟む。

このdepth-wiseを挟むことでパッチ同士のインタラクションが低コストでできる。
インタラクションのイメージが図3。
f:id:Ninhydrin:20211022094948p:plain

Representation head designs

図2の通りHRTは4つの解像度が得られる。これらの扱いはタスクごとに異なる。

ImageNet Classification

4つの解像度の特徴マップそれぞれにbottleneck(多分3x3 convolution)を適用してチャンネルサイズを高解像度から順に128, 256, 512, 1024にする。
そしてstrided convで解像度を低解像度のに合わせる and チャンネル数を256にして結合して2048チャンネルの特徴マップを得る。
それにglobal average pooling適用してクラス分類に利用する。

pose estimation

一番解像度の高いものだけ利用する

semantic segmentation

高解像度の特徴マップに合うように各解像度の特徴マップをupsamplingし、それをconcatしたものを利用する。

Instantiation

各ステージのブロック数やヘッド数などは表1、2を参照
f:id:Ninhydrin:20211022094724p:plain

実験

pose estimation

同じ解像度・パラメータ数のHRNetと比較しても優秀。
f:id:Ninhydrin:20211022095051p:plain

semantic segmentation

f:id:Ninhydrin:20211022095423p:plain

ImageNet-1K

Swin-Bが強い。ただパラメータ数やFLOPsが違うので単純比較はできない。
f:id:Ninhydrin:20211022095602p:plain

Ablation experiments

Influence of 3 × 3 depth-wise convolution within FFN

3x3 depth-wiseの有無。多少のオーバーヘッドでそこそこ大きな精度改善がある。
f:id:Ninhydrin:20211022095914p:plain

Influence of shifted window scheme & 3×3 depth-wise convolution within FFN based on SwinT

Swin Transformerのshifted windowとdepth-wise convの比較。
Intra-WindowはSwin-Tからshifted windowを取り除いたもの。
depth-wiseを追加すると精度が改善。Intra-Window+depth wiseに至ってはSwin-T+depth wiseよりも良いパフォーマンスになった。
f:id:Ninhydrin:20211022100017p:plain

Shifted window scheme v.s. 3×3 depth-wise convolution within FFN based on HRT-T

HRT-Tでshifted windowの効果検証.depth-wiseのが断然良い。
f:id:Ninhydrin:20211022101259p:plain

所感

最近流行りのdepth-wise convを使う系。
depth-wiseだけでなく、LWSAも区切られた空間でattentionをとるということで帰納バイアスがありそう。帰納バイアスって大事だなと感じる。
必ずしも全てTransformerというわけではなく適切なconvolutionとかを使うべきだと思う(全て〇〇でできました系は研究や知見としての意味はあるけど)。
Swin Transormerのshifted windowもなるほどと思ったけど、ニューラルネットワークは特別なことしないでdepth-wiseでも十分対応できる柔軟さがあるなと思ったり。
今後もdepth-wiseを使ったViTはたくさん出てきそう。