[論文レビュー] Expectation-Maximization Attention Networks for Semantic Segmentation

Xia Li, Zhisheng Zhong+4 more|arXiv (Cornell University)|Jul 31, 2019

Advanced Neural Network Applications参考文献 38被引用数 113

ひとこと要約

EMAはEM反復を通じて注意を分析し、ピクセルごとの表現のコンパクトな基底を学習することで、計算量とメモリ使用を抑えた軽量で頑健なセマンティックセグメンテーションモジュール（EMAU）を生み出し、標準ベンチマークで性能を向上させる。

ABSTRACT

Self-attention mechanism has been widely used for various tasks. It is designed to compute the representation of each position by a weighted sum of the features at all positions. Thus, it can capture long-range relations for computer vision tasks. However, it is computationally consuming. Since the attention maps are computed w.r.t all other positions. In this paper, we formulate the attention mechanism into an expectation-maximization manner and iteratively estimate a much more compact set of bases upon which the attention maps are computed. By a weighted summation upon these bases, the resulting representation is low-rank and deprecates noisy information from the input. The proposed Expectation-Maximization Attention (EMA) module is robust to the variance of input and is also friendly in memory and computation. Moreover, we set up the bases maintenance and normalization methods to stabilize its training procedure. We conduct extensive experiments on popular semantic segmentation benchmarks including PASCAL VOC, PASCAL Context and COCO Stuff, on which we set new records.

研究の動機と目的

• 長距離依存性を活用しつつ注意の計算負荷を低減するためにセマンティックセグメンテーションを動機づける。
• 自己注意をEMプロセスとして再定式化し、注意マップのコンパクトな基底セットを学習する。
• CNNバックボーンに容易に組み込める軽量なEMAUモジュールを開発する。

提案手法

• 注意を期待値最大化（EM）プロセスとして再定式化し、注意マップを潜在変数とし、基底を学習するパラメータとする。
• EMAを用いて責任度（Eステップ）を反復的に推定し、基底を更新（Mステップ）して入力特徴の低ランク再構成を得る。
• 学習した基底からコンパクトでノイズに頑健な表現を生成するデータ再推定を導入する。
• EMAをニューラルネットワークモジュール（EMAU）に組み込み、EMAコアの周囲に2つの1x1畳み込みと残差接続を配置する。
• 初期基底をミニバッチ間で移動平均更新することで基底維持を実装し、訓練を安定化させるために基底にユークリッド正規化を適用する。
• 設計選択を検証するために反復回数、維持戦略、正規化のアブレーション研究を提供する。

実験結果

リサーチクエスチョン

• RQ1 EMスタイルの反復注意は、標準の自己注意やNon-localブロックと比較して計算を削減するコンパクトで頑健な基底セットを学習できるか。
• RQ2 EMAUモジュールは、FLOPsとメモリ使用量を低減しつつ標準ベンチマークでセグメンテーション精度を向上させるか。
• RQ3 基底の初期化、維持（移動平均）、正規化（L2Norm）は訓練の安定性と性能にどのような影響を与えるか。

主な発見

• EMAUはPASCAL VOC、PASCAL Context、COCO Stuffで競争力のあるまたは最先端のMean IoUを、いくつかのベースラインよりも低い計算コストで達成する。
• EMスタイルの注意は、計算量をO(N^2)からO(NK)（K << N）に低減し、通常は数回の反復（T ≈ 3）で収束する。
• 移動平均基底の維持とL2正規化は、勾配更新やLNなどの代替戦略よりも訓練安定性と性能を改善する。
• Non-localおよびA^2ブロックと比較して、EMAUはメモリとFLOPsを削減しつつ同等またはそれ以上の性能を提供する。
• 可視化により、学習された基底は単なる前景/背景分離を超えた意味的概念に対応していることが示される。