【論文読解】Disentangled Non-Local Neural Networks

# 概要

Non-local Neural Networks[^nl]（以下、NL）を解釈が容易な形に変形する手法"Disentangle Non-Local Neural Networks"（以下、DNL）について紹介します。

NLは、遠く離れたピクセル間の関係を捉えることを可能にするself-attentionなモジュールで、セマンティックセグメンテーションや行動認識モデルなどで利用されることがあります。DLNでは、NLによって得られるAttentionをWhitened-Pairwise（あるピクセルが他のピクセルに対して持っている重要性）とUnary（そのピクセルが本質的に持っている重要性）とに分離する、ということを目的としたモジュールです。

以下の図をご覧ください。NL/DLNのAttentionをWhitened-PairwiseとUnaryに分離したときの状況を可視化しています。Pairwiseの列では、白い十字が描かれているピクセルに対し、画像中の各ピクセルがどの程度重要とみなせるかが示されてます。また、Unaryの列では、各ピクセルそれ自体の重要性を表しています。DNLのPairwiseでは白い十字が建物、車、歩道、道路などの位置にあるときに、同じカテゴリの領域に対してAttentionが集まっていることがわかります。NLのほうでも似たような傾向は見えますが、DNLよりはあいまいな結果になっています。また、Unaryでは物体カテゴリ間の境界に強くAttentionが出ているということがわかります。NLのほうでは、ここまで明確な区分にはなっておらず、Attentionが何を意味しているのかが分かりにくいです。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/4010/bed33bcd-701d-d53d-a7fa-4ffeb4ece2e5.png)

以上のように目に見えてDNLにおけるAttentionは解釈が容易になっています。

解釈の容易さだけではなくもちろん精度も向上します。NLが使用されているモデルをDNLに置き換えると、一貫した精度の向上がみられる、ということが実験で示されています。

[^nl]: Wang, X., Girshick, R., Gupta, A., He, K.: Non-local neural networks. In: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. pp. 7794–7803 (2018)

# 書誌情報
- Yin, Minghao, et al. "Disentangled Non-Local Neural Networks." arXiv preprint arXiv:2006.06668 (2020).
- https://arxiv.org/abs/2006.06668

# Non-Local Neural Networksの復習
まずは、NLについて簡単に復習していきます。この章では、NLの元論文の表記に従った数式を示します。

NLは以下のような数式で表現されます。

```math
\mathbf{y}_{i}=\frac{1}{\mathcal{C}(\mathbf{x})} \sum_{\forall j} f\left(\mathbf{x}_{i}, \mathbf{x}_{j}\right) g\left(\mathbf{x}_{j}\right)
```

上式は、特徴マップの各座標の特徴量$g\left(\mathbf{x}\_{j}\right)$を$f\left(\mathbf{x}_{i}, \mathbf{x}\_{j}\right)$によって重みづけして足し合わせるという計算です。この操作は、通常のConv層が近傍からのみ特徴量を集約するのに対して、NLは特徴マップ全体から特徴量を集約するという操作になっていると比較して考えるとわかりやすいです。

関数$g$の具体的な形として、単純な線形変換$g\left(\mathbf{x}\_{j}\right)=W_{g} \mathbf{x}_{j}$を考えます。これ以外の形ももちろんあり得ますが、NLの論文で考慮されているのはこの場合のみです。

関数$f$の具体的な形は4つの種類が考えられており、それぞれ以下のようになっています。$\theta, \phi$はいずれも単純な線形変換を表します。

- Gaussian: $f\left(\mathbf{x}\_{i}, \mathbf{x}\_{j}\right)=e^{\mathbf{x}\_{i}^{T} \mathbf{x}\_{j}}$
- Embedded Gaussian: $f\left(\mathbf{x}\_{i}, \mathbf{x}\_{j}\right)=e^{\theta\left(\mathbf{x}\_{i}\right)^{T} \phi\left(\mathbf{x}\_{j}\right)}$
- Dot product: $f\left(\mathbf{x}_{i}, \mathbf{x}\_{j}\right)=\theta\left(\mathbf{x}\_{i}\right)^{T} \phi\left(\mathbf{x}\_{j}\right)$
- Concatenation: $f\left(\mathbf{x}\_{i}, \mathbf{x}\_{j}\right)=\operatorname{ReLU}\left(\mathbf{w}\_{f}^{T}\left[\theta\left(\mathbf{x}\_{i}\right), \phi\left(\mathbf{x}\_{j}\right)\right]\right)$

実験の結果、Gaussian以外の3つはどれも似たような結果になるのでどれでもよいという結論が得られています。Embedding Gaussianの場合を図示したのが以下の図になります。なおこの場合は、$\mathcal{C}(\mathbf{x})=\sum_{\forall j} f\left(\mathbf{x}\_{i}, \mathbf{x}\_{j}\right)$とすることで、$\frac{1}{\mathcal{C}(\mathbf{x})} f\left(\mathbf{x}\_{i}, \mathbf{x}\_{j}\right)$は単なるsoftmax操作となります。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/4010/68038b1a-59bf-a2d9-03da-d06665de88d1.png)

もともとNLは動画における行動認識モデルで使用することを想定しているので、時間を表す次元$T$が含まれていますが、静止画を対象とした場合でも同様の構造となります。

# NLの分解
DNLの論文では、NLの式は以下のように表現されます。

```math
\mathbf{y}_{i}=\sum_{j \in \Omega} \omega\left(\mathbf{x}_{i}, \mathbf{x}_{j}\right) g\left(\mathbf{x}_{j}\right)
```

DNLではEmbedded Gaussianを前提としており、$\omega\left(\mathbf{x}\_{i}, \mathbf{x}\_{j}\right)$は以下のように表現されます。

```math
\omega\left(\mathbf{x}_{i}, \mathbf{x}_{j}\right)=\sigma\left(\mathbf{q}_{i}^{T} \mathbf{k}_{j}\right)=\frac{\exp \left(\mathbf{q}_{i}^{T} \mathbf{k}_{j}\right)}{\sum_{t \in \Omega} \exp \left(\mathbf{q}_{i}^{T} \mathbf{k}_{t}\right)}
```

ここで、$\mathbf{q}\_{i}=W_{q} \mathbf{x}\_{i}, \mathbf{k}\_{j}=W_{k} \mathbf{x}\_{j}$であり、単純な線形変換後の特徴量を表しています。$q, k$という表記は、Attentionの文脈でよく使われるqueryとkeyに対応します。

本論文の大きな主張は、「$\omega\left(\mathbf{x}\_{i}, \mathbf{x}\_{j}\right)$には、query座標$i$とkey座標$j$の間の関係の重要度だけでなく、key座標$j$そのものの重要度も混ざっているので、これらを分離したい」というものです。

まず、白色化によって$\mathbf{q}\_{i}^{T} \mathbf{k}\_{t}$を以下のように変換します。

```math
\mathbf{q}_{i}^{T} \mathbf{k}_{j}=\left(\mathbf{q}_{i}-\boldsymbol{\mu}_{q}\right)^{T}\left(\mathbf{k}_{j}-\boldsymbol{\mu}_{k}\right)+\boldsymbol{\mu}_{q}^{T} \mathbf{k}_{j}+\mathbf{q}_{i}^{T} \boldsymbol{\mu}_{k}+\boldsymbol{\mu}_{q}^{T} \boldsymbol{\mu}_{k}
```

softmaxの操作をする前提であれば、$\mathbf{q}\_{i}^{T} \boldsymbol{\mu}\_{k}+\boldsymbol{\mu}\_{q}^{T} \boldsymbol{\mu}\_{k}$は消去して、$\omega\left(\mathbf{x}\_{i}, \mathbf{x}\_{j}\right)$は以下のような形にできます（導出は論文の付録を参照）。

```math
\omega\left(\mathbf{x}_{i}, \mathbf{x}_{j}\right)=\sigma\left(\mathbf{q}_{i}^{T} \mathbf{k}_{j}\right)=\sigma(\underbrace{\left(\mathbf{q}_{i}-\boldsymbol{\mu}_{q}\right)^{T}\left(\mathbf{k}_{j}-\boldsymbol{\mu}_{k}\right)}_{\text {pairwise }}+\underbrace{\boldsymbol{\mu}_{q}^{T} \mathbf{k}_{j}}_{\text {unary }})
```

$\omega\left(\mathbf{x}\_{i}, \mathbf{x}\_{j}\right)$のsoftmaxの中身を、queryとkeyの関係によって決まるpairwise項とkeyそのものの重要性を表すunary項とに分離できました。このように意味的に異なる項に分離できるからには、それぞれが何らかの役割分担をしていると期待したいところですが、冒頭の図で示したように通常のNLではあまりうまく役割分担ができるとは言えません。

なぜうまくいかないのか、勾配を確認することで理解できます。その準備として、softmax部分を以下のように分解します。

```math
\begin{aligned}
\sigma\left(\mathbf{q}_{i} \cdot \mathbf{k}_{j}\right) &=\sigma\left(\left(\mathbf{q}_{i}-\boldsymbol{\mu}_{q}\right)^{T}\left(\mathbf{k}_{j}-\boldsymbol{\mu}_{k}\right)+\boldsymbol{\mu}_{q}^{T} \mathbf{k}_{j}\right) \\
&=\frac{1}{\lambda_{i}} \sigma\left(\left(\mathbf{q}_{i}-\boldsymbol{\mu}_{q}\right)^{T}\left(\mathbf{k}_{j}-\boldsymbol{\mu}_{k}\right)\right) \cdot \sigma\left(\boldsymbol{\mu}_{q}^{T} \mathbf{k}_{j}\right)=\frac{1}{\lambda_{i}} \omega_{\mathrm{p}}\left(\mathbf{x}_{i}, \mathbf{x}_{j}\right) \cdot \omega_{\mathbf{u}}\left(\mathbf{x}_{i}, \mathbf{x}_{j}\right)
\end{aligned}
```

ここで、$\omega\_{\mathrm{p}}\left(\mathbf{x}\_{i}, \mathbf{x}\_{j}\right)=\sigma\left(\left(\mathbf{q}\_{i}-\boldsymbol{\mu}\_{q}\right)^{T}\left(\mathbf{k}\_{j}-\boldsymbol{\mu}\_{k}\right)\right)$、 $\omega\_{\mathbf{u}}\left(\mathbf{x}\_{i}, \mathbf{x}\_{j}\right)=\sigma\left(\boldsymbol{\mu}\_{q}^{T} \mathbf{k}\_{j}\right)$で、$\lambda_i$は正規化項です。


そのうえで損失に対する勾配を確認してみると、以下のようになります。片方の項の勾配に、もう片方の項が出てきいることがわかります。このようなとき、片方の項が小さい場合に勾配が消失してしまい、もう片方の項の訓練が進まなくなります。

```math
\begin{aligned} \frac{\partial L}{\partial \sigma\left(\omega_{\mathrm{p}}\right)} &=\frac{\partial L}{\partial \sigma(\omega)} \cdot \frac{\partial \sigma(\omega)}{\partial \sigma\left(\omega_{\mathrm{p}}\right)}=\frac{\partial L}{\partial \sigma(\omega)} \cdot \sigma\left(\omega_{\mathrm{u}}\right), \\\frac{\partial L}{\partial \sigma\left(\omega_{\mathrm{u}}\right)} &=\frac{\partial L}{\partial \sigma(\omega)} \cdot \frac{\partial \sigma(\omega)}{\partial \sigma\left(\omega_{\mathrm{u}}\right)}=\frac{\partial L}{\partial \sigma(\omega)} \cdot \sigma\left(\omega_{\mathrm{p}}\right)\end{aligned}
```
また、pairwiseとunaryに共通して現れる$\mathbf{k}_j$を作るための$W_k$の存在が、両者の関係をさらに複雑にしてしまっています。

# DNLの導入
ここまでの議論を踏まえ、NLのpairwiseとunaryの項の役割分担を実現できるDNLを導入します。DNLはNLにおける$\omega\left(\mathbf{x}\_{i}, \mathbf{x}\_{j}\right)$にたいして以下の2点の変更を加えます。

1. softmaxの乗算を加算にする
2. $W_k$をpairwiseとunaryで個別に用意する

まず、各項の勾配にもう片方の項が出てきてしまうという問題を解決するために、乗算になっているところを加算にします。

```math
\omega\left(\mathbf{x}_{i}, \mathbf{x}_{j}\right)=\omega_{\mathrm{p}}\left(\mathbf{x}_{i}, \mathbf{x}_{j}\right) \cdot \omega_{\mathrm{u}}\left(\mathbf{x}_{i}, \mathbf{x}_{j}\right) \Rightarrow \omega\left(\mathbf{x}_{i}, \mathbf{x}_{j}\right)=\omega_{\mathrm{p}}\left(\mathbf{x}_{i}, \mathbf{x}_{j}\right)+\omega_{\mathrm{u}}\left(\mathbf{x}_{i}, \mathbf{x}_{j}\right)
```

次に、$\mathbf{k}\_j$を作るための$W_k$をpairwiseとunaryで個別に用意します。ここでは、unary項である$\boldsymbol{\mu}\_{q}^{T} W\_{k} \mathbf{x}\_{j}$を単純化し、$W_{m} \mathbf{x}\_{j}$とします。

```math
\boldsymbol{\mu}_{q}^{T} \mathbf{k}_{j}=\boldsymbol{\mu}_{q}^{T} W_{k} \mathbf{x}_{j} \Rightarrow m_{j}=W_{m} \mathbf{x}_{j}
```

以上を施すことで、$\omega\left(\mathbf{x}\_{i}, \mathbf{x}\_{j}\right)$のDisentangledバージョンである$\omega^{\mathrm{D}}\left(\mathbf{x}\_{i}, \mathbf{x}\_{j}\right)$は、以下のように表現されます。

```math
\omega^{\mathrm{D}}\left(\mathbf{x}_{i}, \mathbf{x}_{j}\right)=\sigma\left(\left(\mathbf{q}_{i}-\boldsymbol{\mu}_{q}\right)^{T}\left(\mathbf{k}_{j}-\boldsymbol{\mu}_{k}\right)\right)+\sigma\left(m_{j}\right)
```

DNLを図示すると、下図の(d)のような形になります。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/4010/80d4abbb-3ab7-45f8-013c-80b4ed3db0fc.png)

# 実験
論文ではセマンティックセグメンテーションや行動認識モデルの実験を通して本手法の有効性を示していますが、最も端的に本手法の有効性を示している実験を見ていきます。

セマンティックセグメンテーションのデータセットであるCityscapesを用いて、以下の2点を確認しています。

- pairwiseのAttentionが同一カテゴリの領域をどの程度表しているか
- pairwiseのAttentionがカテゴリ間の境界領域とはどの程度無関か
- unaryのAttentionがカテゴリ間の境界領域をどの程度表しているか

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/4010/25fc4917-03a9-3dac-e5ec-2a3541aa321d.png)

randomはランダムなAttentionにしたときのベースラインです。Eq.6と7は、それぞれpairwise項のみを残したバージョンとunary項のみを残したバージョンのことです。また、Eq.13と14は以下の2つのバージョンのことで、それぞれNDLに導入されている工夫の片方を取り除いたものになっています。

```math
\begin{aligned}
\omega^{\mathrm{D} *}\left(\mathbf{x}_{i}, \mathbf{x}_{j}\right) &=\sigma\left(\left(\mathbf{q}_{i}-\boldsymbol{\mu}_{q}\right)^{T}\left(\mathbf{k}_{j}-\boldsymbol{\mu}_{k}\right)+m_{j}\right) \\
\omega^{\mathrm{D} \dagger}\left(\mathbf{x}_{i}, \mathbf{x}_{j}\right) &=\sigma\left(\left(\mathbf{q}_{i}-\boldsymbol{\mu}_{q}\right)^{T}\left(\mathbf{k}_{j}-\boldsymbol{\mu}_{k}\right)\right)+\sigma\left(\boldsymbol{\mu}_{q}^{T} \mathbf{k}_{j}\right)
\end{aligned}
```

こうして比較すると、NLのAttentionは、カテゴリ領域との関係が不明瞭ですが、DNLでは工夫を加えていくにしたがってこれらの値が改善されることがわかります。また、pairwiseだけ、unaryだけというバージョンよりもDNLが優れた結果になっていることから、相互に補完的に訓練が進むことが確認できています。

# まとめ
NLをピクセル間の関係であるpairwiseとピクセルそのものの重要性であるunaryに分解し、それぞれを適切に訓練できるようにしたDNLを紹介しました。

NLに限らず、Attentionを含むモデルは解釈可能と喧伝される一方で必ずしも直感的な解釈が困難なことがあります。本手法も完ぺきではないにせよ一歩改善が見られます。そのうえで精度も向上するという点も興味深い結果だと思います。