Polarized Self-Attention: Towards High-quality Pixel-wise Regression

arXiv

概要

この論文では、精密なコンピュータビジョンタスク（例えば、キーポイントヒートマップの推定やセグメンテーションマスクの生成）におけるピクセル単位の回帰を改善するために、Polarized Self-Attention (PSA)ブロックを提案している．
PSAは、チャンネルと空間の両方の注意計算において高い内部解像度を維持しながら、入力テンソルをそれぞれの対角方向に完全に圧縮するポラライズドフィルタリングを行う．さらに、PSAは典型的な細かい回帰の出力分布に直接適合する非線形性を組み合わせており、2Dガウス分布（キーポイントヒートマップ）や2D二項分布（二値セグメンテーションマスク）を直接フィットする．

従来手法との比較

従来の注意メカニズムでは、要素ごとの注意（例：Nonlocalブロック）が非常に複雑でノイズに敏感であり、多くの単純化された注意ハイブリッドが複数のタスク間で最適な妥協点を探していた．PSAは、チャンネルと空間の両方の注意計算において高い内部解像度を維持し、典型的な細かい回帰の出力分布に直接適合する非線形性を組み合わせることで、これらの課題を克服している．

新規性

PSAの新規性は以下の点にある：

ポラライズドフィルタリング: 入力テンソルをそれぞれの対角方向に完全に圧縮し、内部解像度を維持する
高解像度の注意計算: チャンネルと空間の両方の注意計算において高い内部解像度を維持する
非線形性の組み合わせ: 2Dガウス分布や2D二項分布を直接フィットするために、ソフトマックスとシグモイドの組み合わせを使用する

方法論

PSAブロックは、次の2つの主要なコンポーネントで構成されている：

チャンネル専用ブランチ (Ach): 入力テンソルのチャンネル方向に沿った特徴を強調または抑制する
空間専用ブランチ (Asp): 入力テンソルの空間方向に沿った特徴を強調または抑制する

これらのブランチは、並列レイアウトまたはシーケンシャルレイアウトで組み合わせることができる．

PSAの構造

Polarized Filetering：入力テンソルをチャネル方向と空間方向に分割し、それぞれ高解像度を保ちながら対角方向を完全に圧縮するフィルタリング手法
- 従来の手法では、計算効率を高めるために内部の解像度を低下させているが、このタイミングで高解像度の情報を失っている。
- PSAでは、チャネル専用ブランチと空間専用ブランチを作成し、チャネル専用ブランチでは、チャネル方向のチャネル解像度のみを高く保ち、空間方向は高解像度のまま計算を行う。
Nonlinearity Enhancement：出力分布に直接適合する非線形性を組み込み、典型的な回帰タスクの出力分布に適合させる
- 従来、Self-Attentionブロックは特定の出力分布に最適化されていない場合が多い
- PSAでは、Softmax-Sigmoidの組み合わせを導入し、2Dガウス分布や2D二項分布などの典型的な回帰タスクの出力分布に直接適合する。

結果と評価

PSAは、MS-COCOデータセットの2Dポーズ推定とPascal VOC2012 Augデータセットのセマンティックセグメンテーションのベンチマークにおいて、ベースライン手法を向上させることが実験で示されている．

結論

PSAは、様々なコンピュータビジョンタスクにおいて、より高精度な結果を得ることが期待される．