局所的な情報しか扱うことができないCNNやRNNに、大域的な情報を付与することが可能なnon-local blockというモジュールを提案した論文です。traditionalなノイズ除去手法のNon-local Mean Filterを、ニューラルネットワークで表現したような手法です。

Abstract

　CNNやRNNでは局所的な情報しか扱うことができない。CNNだとカーネルサイズを大きくするとか方法があるけど、限界があるし小さいカーネルの利点が失われる。そこで、本論文ではより大域的な情報を扱うことが可能なnon-local operatorを提案する。 non-local operatorはノイズ除去手法のNon-local Mean Filterから着想を得ており、着目点の値を特徴マップ中の全ての点の重み付き和として表現する。 non-local operatorはcomputer vision系の様々な構造のモデルに導入可能であり、導入することで動画像の分類や物体検出、領域分割において性能が向上した。

3. Non-local Neural Networks

3.1. Formulation

　Non-local Mean Filterの演算をニューラルネットワーク向けに以下のようなnon local operatorに定式化する。

$\displaystyle {\bf y}_i = \frac{1}{C({\bf x})} \sum_{\forall j} f({\bf x}_i, {\bf x}_j) g({\bf x}_j)$

$f({\bf x}_i, {\bf x}_j)$ は特徴マップ中の着目点 $i$ とある点 $j$ の類似度、 $g({\bf x}_j)$ はある点 $j$ での入力信号の表現を計算する関数を表す。 $C({\bf x})$ は正規化係数。

全結合層との違い

　non-local operatorにおいて全結合層の重みにあたる $f$ は学習パラメータではなく類似度である。また、non-local operatoのr入力サイズは任意でありCNNやRNNに簡単に導入できる。

3.2. Instantiations

　 $g({\bf x}_j)$ はカーネルサイズが1の畳み込みによるembeddingで表現する。類似度を表す $f({\bf x}_i, {\bf x}_j)$ には以下のようなものが考えられる。

Gaussian

　 ${\bf x}_i$ と ${\bf x}_j$ の内積にガウス関数をかける。

$\displaystyle f({\bf x}_i, {\bf x}_j) = e^{{\bf x}_i^T{\bf x}_j}$

ユークリッド距離ではなく内積を用いるのは、その方がDeep learningのプラットフォーム上で実装しやすいため。正規化係数は $f({\bf x}_i, {\bf x}_j)$ の総和 $C({\bf x}) = \sum_{\forall j} f({\bf x}_i, {\bf x}_j)$ 。

Embedded Gaussian

　 ${\bf x}_i$ と ${\bf x}_j$ のembeddingの内積にガウス関数をかける。

$\displaystyle f({\bf x}_i, {\bf x}_j) = e^{\theta({\bf x}_i)^T\phi({\bf x}_j)}$

\theta({\bf x}_i)と\phi({\bf x}_j)はembeddingを求める関数で $g({\bf x}_j)$ と同様に、カーネルサイズが1の畳み込みで表現する。正規化係数はGaussianと同様に $f({\bf x}_i, {\bf x}_j)$ の総和 $C({\bf x}) = \sum_{\forall j} f({\bf x}_i, {\bf x}_j)$ 。

　機械翻訳タスクにおいて用いられているself-attention moduleは、Embedded Gaussianを類似度関数に使ったnon-local operatorとして捉えることができる。 Embedded Gaussianを用いるとき、 $\frac{1}{C}({\bf x}) f({\bf x}_i, {\bf y}_i)$ は次元jに沿ってsoftmax関数をかけるのと同じとみなせるため、non-local operatorは ${\bf y} = softmax({\bf x}^T {\bf W}_{\theta}^T {\bf W}_{\phi} {\bf x})$ と表すことができる。これはself-attention moduleを表す式と同じ。non-local operatorは、シーケンスのみに対応していたself-attention moduleを画像や動画像データにも対応可能なように拡張したものといえる。

　self-attention moduleとの関係性とは反対に、softmaxは必ずしも必要ではない。

Dot product

　 ${\bf x}_i$ と ${\bf x}_j$ のembeddingの内積。Embedded Gaussianのガウス関数を除いたもの。

$\displaystyle f({\bf x}_i, {\bf x}_j) = \theta({\bf x}_i)^T\phi({\bf x}_j)$

正規化係数は ${\bf x}$ の要素数 $N$ 。 $f({\bf x}_i, {\bf x}_j)$ の総和を使わないのは、勾配計算を簡単にするため。

Concatenation

　visual reasoning(どういうタスク?)用のRelation Networksというモデルで使われている関数。

$\displaystyle f({\bf x}_i, {\bf x}_j) = {\rm ReLU}({\bf w}_f^T [\theta({\bf x}_i), \phi({\bf x}_j)])$

$[ \cdot , \cdot ]$ は連結を表す。 ${\bf w}_f$ はベクトルをスカラーに射影するための重みベクトル。正規化係数はDot productと同様に ${\bf x}$ の要素数 $N$ 。

3.3. Non-local Block

　non-local operatorを既存のモデルに導入するために、ラッパーとしてnon-local blockを用いる。non-local blockは以下の式で表される。

$\displaystyle {\bf z_i} = {\bf W}_z {\bf y}_i + {\bf x}_i$

${\bf y}_i$ はnon local operatorの出力。 $+ {\bf x}_i$ はresidual connectionであり、 ${\bf W}_z$ の初期値を0にすることで事前学習済みモデルに対しても、元々の振る舞いを妨げることなくnon-local operationを導入することができる。