【GNN】NLNN：self-Attention 的泛化框架

今天學習的是 CMU 和 Facebook 的同學 2017 年的工作《Non-local Neural Networks》，發表於 CVPR 18，作者王小龍，導師是何凱明老師，這篇論文的截止目前的引用量有 1100 多次。

這篇論文的應用領域最初是在計算機視覺，作者提出的非局部神經網絡（non-local neural network，NLNN）來捕獲深度神經網絡的長依賴關係，這種非局部運算是對所有位置特徵的加權計算，因此也可以被看作是不同“self-attention”的方法的統一。

1.Introduction

卷積運算和遞歸運算是典型的局部運算（只考慮局部信息），長距離依賴關係通常是通過不斷迭代來形成更大的感受野。但這種方式有很多缺點，包括計算效率低，優化困難等。

作者受 CV 領域中經典的非局部均值濾波的啓發提出了 NLNN，將一個位置相應爲所有位置特徵的加權計算，這種位置位置可以是空間、時間或者時空，也就是說這種方法的應用領域適用於圖像、序列和視頻等問題。

這種非局部運算有幾大優點：

1. 與卷積運算和遞歸運算相比，非局部運算通過計算任意兩個位置之間的交互直接捕獲長依賴；
2. 即使神經網絡的層數不多也能達到最佳效果；
3. 非局部運算支持可變的輸入大小，並且可以輕鬆與其他運算（如卷積運算）相結合。

在介紹正文之前，先接單介紹下非局部均值法。

非局部均值濾波（Non-Local Means，NLM）是 Buades 等人 2005 年的論文《A non-local algorithm for image denoising》提出的一種對傳統領域濾波方法進行改進的一種濾波方法，其考慮圖像的自相似性，充分利用圖像中的非局部信息來輔助濾波降噪。

算法的主要原理是計算圖像中所有像素與當前像素的相似性，其公式定義爲：
$\begin{aligned} NL[v](i) &= \sum_{j∈I}^{}{}{} w(i,j)v(j) \\ w(i,j) &= \frac{1}{Z(i)}e^{−\frac{||v(Ni)−v(Nj)||^{2}_{2,a}}{h2} } \\ Z(i) &= \sum_{j}e^{−\frac{||v(Ni)−v(Nj)||^{2}_{2,a}}{h2} } \\ \end{aligned} \\$
其中，$w(i,j)$ 爲兩個像素的相似性；$N_i$ 表示以像素 i 爲中心固定大小的領域；$v(N_i)$ 爲灰度強度的映射向量；$Z(i)$ 爲歸一化常數；$h$ 爲平滑參數。

2.NLNN

2.1 Formulation

首先我們先來看下非局部運算的通用公式，受非局部均值濾波方法的影響，我們有：
$\mathbb{y}_i = \frac{1}{\mathcal{C}(\mathbb{x})}\sum_{\forall j} f(\mathbb{x}_i, \mathbb{x}_j) g(\mathbb{x}_j) \\$
其中，$\mathbb{x_i}$ 表示 i 位置的特徵向量；f 爲計算相似性的函數；g 爲將點映射爲向量的映射函數（可以理解爲計算點的特徵）；$\mathcal{C}(\mathbb{x})$ 爲一個歸一化因子。

非局部運算考慮的是所有位置，而局部運算如卷積操作只考慮節點的周圍信息，遞歸運算只考慮上一步的信息。

非局部運算與全連接運算的區別主要有兩點：

1. 非局部運算的多輸出值會受到輸入值之間關係的影響（因爲要計算 $f(\cdot,\cdot)$），而全連接運算通過學習到權重來計算輸入到輸出的映射，所以不同位置的關係是不會影響到輸出，一定程度上損失了位置的相關性；
2. 非局部運算支持可變大小的輸入，而全連接運算的輸入大小是固定的；
3. 非局部操作可以很容易地與卷積/遞歸層一起使用，而不像全連接層只能放到神經網絡的最後使用。這便使得我們能夠將非局部信息和局部信息結合起來，來構建一個更豐富的層次結構，

2.2 Instantiations

接下來我們介紹幾種具體的例子。

由於作者在實驗中發現不同形式的實例其實對結果的影響不太大（因爲這種非局部的行爲纔是主要的影響因素），所以爲了簡化操作，只考慮線性映射函數 $g(\mathbb{x}_i)=W_g \mathbb{x}_i$，其中，$W_g$ 是一個可學習的權重矩陣。

接下來我們考慮 f 的不通過中形式：

Gaussian：從非局部均值濾波和雙邊濾波的角度來看，顯而易見的首先想到高斯濾波：
$f(\mathbb{x}_i, \mathbb{x}_j)=e^{\mathbb{x}_i^T \cdot \mathbb{x}_j} \\ where \quad \mathcal{C}(x)=\sum_{\forall j} f(\mathbb{x}_i, \mathbb{x}_j) \\$
Embedded Gaussian：考慮到 Embedding 的向量空間，我們有高斯濾波的變種：
$f(\mathbb{x}_i, \mathbb{x}_j)=e^{\theta(\mathbb{x}_i)^T \cdot \phi(\mathbb{x}_j)}\\ where\quad\theta(\mathbb{x}_i)=W_{\theta}\mathbb{x}_i \quad \phi ( \mathbb{x}_i)=W_{\phi}\mathbb{x}_i \quad \mathcal{C}(\mathbb{x})=\sum_{\forall j} f(\mathbb{x}_i, \mathbb{x}_j) \\$
self-attention 其實就是 non-local 的 Embedded Gaussian 的一種特殊情況。對於所有 i，我們把 $\frac{1}{\mathcal{C}(\mathbb{x})} f(\mathbb{x}_i, \mathbb{x}_j)$ 變成計算所有 j 的 Softmax，得到：
$y = sofmax(\mathbb{x}^T W_{\theta}^T W_{\phi} \mathbb{x})g(\mathbb{x}) \\$
這便是《Attention is all your need》中的表達式。

此時，作者便將 self-attention 與傳統的非局部均值聯繫到了一起，並將基於序列的 self-attention 推廣空間和時空領域中。

此外，與《Attention》那篇論文不同的是，作者發現 Attention 的形式對於研究 case 來說並不必須的，爲了說明這點，作者有給出了以下幾種表達形式。

Dot product：考慮點乘相似度：
$f(\mathbb{x}_i, \mathbb{x}_j)=\theta(\mathbb{x}_i)^T \cdot \phi(\mathbb{x}_j) \\ \mathcal{C}(\mathbb{x})=N \\$
N 爲輸入向量 x 的數量，而不是 f 的和，這樣做的目的在於簡化梯度的計算。

這個版本與 Embedded Gaussian 的區別在於是否做 softmax。

Concatenationn：考慮連接運算：
$f(\mathbb{x}_i, \mathbb{x}_j)=\text{ReLU}(\mathbb{w}_f^T \cdot [\theta(\mathbb{x}_i), \phi(\mathbb{x}_j)]) \\ \mathcal{C}(x)= N \\$
其中，$[\cdot,\cdot]$ 爲連接運算，$\mathbb{w}_f$ 爲將 concat 向量轉換成標量的權重向量。

當然，也可以是其他變種，這一定程度上也說明了 non-local 的靈活性。

2.3 Non-local Block

這一節作者將非局部運算變成一個非局部模塊，以便插入到現有的神經網絡結構中。

作者將其設計成殘差模塊，讓其學習 x 的殘差：
$\mathbb{z}_i =W_z \mathbb{y}_i + \mathbb{x}_i \\$
這種殘差結構可以使得在任意的預訓練模型中直接插入一個新的非局部模塊而不用改變原有的網絡結構。考慮計算量問題，在輸入比較大的情況下可以只在高階語義層中引入 non-local layer，此外也可以通過 pooling 層來減少計算量。

看一個簡單的實現：

$\oplus$ 爲元素點和，$\otimes$ 爲矩陣相乘。

我們通過利用 $1\times1\times1$ 的卷積將 $W_g,W_\theta,W_\phi$ 的 channel 數量控制爲 X 的一半，形成一個 bottleneck 的結構，可以大大減少計算量。在輸出時利用 $W_z$ 可以將 channel 還原，從保證了輸入輸出的維度一致。

此外，還有一個下采樣（subsampling）的 trick 可以使用，比如說，我們可以把通用公式修改一下：
$y_i=\frac {1}{C(\hat{\mathbb{x}})}\sum_{\forall j}f(\mathbb{x}_i,\hat{\mathbb{x}}_j)g(\hat {\mathbb{x}}_j) \\$
其中，$\hat{\mathbb{x}}_j$ 爲 $\mathbb{x}_j$ 的下采樣，比如說可以利用 pooling，這樣就可以減少 1/4 的 pair-wise 的計算量了。

這個 trick 不會改變 non-local 的行爲，但是可以使得計算更加稀疏，通常會在上圖中 $\phi, g$ 的後面增加一層 max pooling 層來實現。

3.Experiments

由於視頻分類的實驗，所以簡單看一下就行。

首先是動力學行爲分類的多個實驗：

然後是模型與 baseline 的對比：

然後是不同模型之間的對比：

4.Conclusion

總結：作者根據非局部均值濾波方法提出了非局部神經網絡（NLNN），綜合考慮節點間不同位置之間的關係，同時也具有較強的靈活性，支持自定義相似函數，並支持在不改變網絡結構的情況下以殘差的形式接入到現有網絡中，最後通過實驗表明，這種非局部運算的框架具有非常顯著的效果。

這篇文章思路總的來說還是比較清晰，模塊也非常簡單，由於 NLNN 支持可變輸入大小，所以也可以無縫銜接到 GNN 中。

5.Reference

1. Wang X, Girshick R, Gupta A, et al. Non-local neural networks[C]//Proceedings of the IEEE conference on computer vision and pattern recognition. 2018: 7794-7803.
2. 《Non-local neural networks》
3. 《video-nonlocal-net》

關注公衆號跟蹤最新內容：阿澤的學習筆記。