[論文メモ] Masked-attention Mask Transformer for Universal Image Segmentation

arxiv.org

panoptic・instance・semantic segmentation用のMasked-attention Mask Transformer(Mask2Former)を提案。
f:id:Ninhydrin:20211223093746p:plain

手法

Mask2Formerは3つのコンポーネントから成る。
1) 特徴量を抽出するバックボーン
2) 低解像度の特徴量を高解像にするため徐々にupsamplingするpixel decoder
3) 最終的な予測を行うTransformer Decoder

この中でもTransformer Decoderの改良がメイン。

Masked attention

Transformer-based modelは収束が遅いといわれており、cross-attentionの学習に

彼らは局所的な特徴はクエリ特徴の更新で十分で、コンテキスト情報はself-attentionで吸収出来ると仮定。
一般的なcross-attentionは以下。 $l$ は層インデックス、 $\textbf{X}_l \in \mathbb{R}^{N \times C}$ は $l$ 番目の層の特徴量で $N$ がシーケンス長、 $C$ が特徴量の次元。あとはよしなに。
f:id:Ninhydrin:20211222094051p:plain

提案するmasked attentionは以下の式(2)。
f:id:Ninhydrin:20211222094508p:plain
$\mathcal{M}_{l-1}$ は以下の式(3)で定義されるattentionマスク。
f:id:Ninhydrin:20211222094554p:plain
$\textbf{M}_{l-1} \in \{0,1\}^{N \times H_l W_l}$ は( $l - 1$ )番目(1つ前)のdecoder layerから予測したバイナリマスクで、各座標( $x, y$ )毎。
$\mathcal{M}$ と $\textbf{M}$ なので注意。
また $\textbf{M}_0$ は $\textbf{X}_0$ から予測するらしいが、 $\textbf{X}_0$ は入力画像に非依存な気がするのだがどういうこと？

High-resolution features

高解像度の特徴で処理したほうが精度は良くなるが計算コスト等の問題がある。
そこでpixel decoderにfeature pyramidの戦略を導入する。
pixel decoderの各解像度の特徴量には正弦波positional encoding( $e_{\verb|pos|} \in \mathbb{R}^{H_l \times W_l}$ )と学習可能な解像度毎のembbeding $e_{\verb|lvl|} \in \mathbb{R}^{1 \times C}$ を加算する。
$H_l, W_l$ は層 $l$ でのheightとwidth。

Optimization improvements

一般的なTransformer decoderはself-attention(SA)、cross-attention(CA)、feed-forward network(NNF)の3つから構成される。
このdecoderに対して3つの変更を加える

1つ目はSAとCAの入れ替え。
一番はじめのSAはpixel decoderからの特徴が入ってきておらずSAの意味がない。
なので一般的にはSA->CAの順番だがこれをCA->SAにする。

2つ目は一番はじめのクエリ特徴 $\textbf{X}_0$ を学習可能パラメータにする。

3つ目はdropoutの削除。試してみてdropoutは必要なくパフォーマンスが下がることがわかった。

Improving training efficiency

高解像度のマスクを扱うとメモリ消費が大きい(MaskFormerなんかは画像1枚でGPUメモリを32GB消費する)。
PointRendという手法ではセグメンテーションのマスクのlossを全体ではなくK個のランダムサンプリングした点で行うことでメモリ消費を抑えており、それに倣った方法で18~6GBの消費に抑えた。

実験・結果

詳しくは論文参照。
f:id:Ninhydrin:20211223092620p:plain
f:id:Ninhydrin:20211223092636p:plain
f:id:Ninhydrin:20211223092657p:plain

所感

実装はdetectron2らしく誰でも触れそう。
このあとに出た「Mask2Former for Video Instance Segmentation」の事前知識として読んだ。
NLPでのattention-maskとは意味は異なるが、画像でもマスクしてattentionするというのは面白い発想。
ただ、そのマスクがどれだけ信頼できるのかなどが少々気になる。masked-attention自体は効果があることがAblationで示されてはいる。
$\textbf{M}_0$ を $\textbf{X}_0$ から予測するのは謎。解像度も小さそうだし、画像非依存なら別にマスクいらないのでは？