解讀北大提出的期望最大化注意力網路EMANet

機器之心專欄

作者：李夏

ICCV 2019 將在10月27於韓國首爾舉行，本文作者介紹了一篇Oral論文，它利用EM 演算法優化注意力機制，並在語義分割等任務上獲得更好的效果。

本文介紹筆者被 ICCV 2019 接受為 Oral 的論文 Expectation-Maximization Attention Networks for Semantic Segmentation[1]。論文作者為：李夏、鍾之聲、吳建龍、楊一博、林宙辰、劉宏。

背景介紹

語義分割是計算機視覺領域的一項基礎任務，它的目標是為每個像素預測類別標籤。由於類別多樣繁雜，且類間表徵相似度大，語義分割要求模型具有強大的區分能力。近年來，基於全卷積網路（FCN[2]）的一系列研究，在該任務上取得了卓越的成績。

這些語義分割網路，由骨幹網路和語義分割頭組成。全卷積網路受制於較小的有效感知域，無法充分捕獲長距離信息。為彌補這一缺陷，諸多工作提出提出了高效的多尺度上下文融合模塊，例如全局池化層、Deeplab[3] 的空洞空間卷積池化金字塔、PSPNet[4] 的金字塔池化模塊等。

近年來，自注意力機制在自然語言處理領域取得卓越成果。Nonlocal[5] 被提出後，在計算機視覺領域也受到了廣泛的關注，並被一系列文章證明了在語義分割中的有效性。它使得每個像素可以充分捕獲全局信息。然而，自注意力機制需要生成一個巨大的注意力圖，其空間複雜度和時間複雜度巨大。其瓶頸在於，每一個像素的注意力圖都需要對全圖計算。

本文所提出的期望最大化注意力機制（EMA），摒棄了在全圖上計算注意力圖的流程，轉而通過期望最大化（EM）演算法迭代出一組緊湊的基，在這組基上運行注意力機制，從而大大降低了複雜度。其中，E 步更新注意力圖，M 步更新這組基。E、M 交替執行，收斂之後用來重建特徵圖。本文把這一機制嵌入網路中，構造出輕量且易實現的 EMA Unit。其作為語義分割頭，在多個數據集上取得了較高的精度。

期望最大化注意力

前提知識

期望最大化演算法

期望最大化（EM）演算法旨在為隱變數模型尋找最大似然解。對於觀測數據 X=，每一個數據點 x_i 都對應隱變數 z_i。我們把 稱為完整數據，其似然函數為 ln p(X, Z|θ)，θ 是模型的參數。

E 步根據當前參數θ^old 計算隱變數 Z 的後驗分布，並以之尋找完整數據的似然 Q(θ, θ^old)：

M 步通過最大化似然函數來更新參數得到θ^new：

EM 演算法被證明會收斂到局部最大值處，且迭代過程完整數據似然值單調遞增。

高斯混合模型（GMM）是 EM 演算法的一個範例，它把數據用多個高斯分布擬合。其 θ_k 即為第 k 個高斯分布的參數μ_k, Σ_k，隱變數 z_nk 為第 k 個高斯分布對第 n 數據點的「責任」。E 步更新「責任」，M 步更新高斯參數。在實際應用中，Σ_k 經常被簡化為 I。

非局部網路

非局部網路（Nonlocal[5]）率先將自注意力機制使用在計算機視覺任務中。其核心運算元是：

其中 f(., .) 表示廣義的核函數，C(x) 是歸一化係數。它將第 i 個像素的特徵 x_i 更新為其他所有像素特徵經過 g 變換之後的加權平均 y_i，權重通過歸一化後的核函數計算，表徵兩個像素之間的相關度。這裡 1

期望最大化注意力機制

期望最大化注意力機制由 A_E, A_M, A_R 三部分組成，前兩者分別對應 EM 演算法的 E 步和 M 步。假定輸入的特徵圖為，基初始值為，A_E 估計隱變數，即每個基對像素的權責。具體地，第 k 個基對第 n 個像素的權責可以計算為：

在這裡，內核 K(a, b) 可以有多種選擇。我們選擇的形式。在實現中，可以用如下的方式實現：

其中，λ 作為超參數來控制 Z 的分布。

A_M 步更新基 μ。為了保證μ和 X 處在同一表徵空間內，此處μ被計算作 X 的加權平均。具體地，第 k 個基被更新為：

值得注意的是，如果λ趨向於無窮，則公式 (5) 中， 會變成一組 one-hot 編碼。在這種情形下，每個像素僅由一個基負責，而基被更新為其所負責的像素的均值，這便是標準的 K-means 演算法。

A_E 和 A_M 交替執行 T 步。此後，近似收斂的μ和 Z 便可以被用來對 X 進行重估計得 X tilde：

X tilde 相比於 X，具有低秩的特性。從下圖中可看出，其在保持類間差異的同時，類別內部差異得到縮小。從圖像角度來看，起到了類似保邊濾波的效果。

綜上，EMA 在獲得低秩重構特性的同時，將複雜度從 Nonlocal 的 O(N^2) 降低至 O(NKT)。實驗中，EMA 僅需 3 步就可達到近似收斂，因此 T 作為一個小常數，可以被省去。至此，EMA 複雜度僅為 O(NK)。考慮到 k 遠小於 N，其複雜度得到顯著的降低。

期望最大化注意力模塊

EMA Unit

期望最大化注意力模塊（EMAU）的結構如上圖所示。除了核心的 EMA 之外，兩個 1×1 卷積分別放置於 EMA 前後。前者將輸入的值域從 R 映射到 R；後者將 X tilde 映射到 X 的殘差空間。囊括進兩個卷積的額外負荷，EMAU 的 FLOPs 僅相當於同樣輸入輸出大小時 3×3 卷積的 1/3，參數量僅為 2C^2 KC。

對於 EM 演算法而言，參數的初始化會影響到最終收斂時的效果。上一節中討論了 EMA 如何在單張圖像的特徵圖上進行迭代運算。而對於深度網路訓練過程中的大量圖片，在逐個批次訓練的同時，EM 參數的迭代初值理應得到不斷優化。本文中，迭代初值的維護參考 BN 中 running_mean 和 running_std 的滑動平均更新方式，即：

其中，α∈[0,1] 表示動量；表示在一個 mini-batch 上的平均。

此外，EMA 的迭代過程可以展開為一個 RNN，其反向傳播也會面臨梯度爆炸或消失等問題。此外，公式 (8) 也要求和的差異不宜過大，不然初值的更新也會出現不穩定。RNN 中採取 LayerNorm（LN）來進行歸一化是一個合理的選擇。但在 EMA 中，LN 會改變基的方向，進而影響其語義。因為，本文選擇 L2Norm 來對基進行歸一化。這樣，的更新軌跡便處在一個高維球面上。

此處，我們可以考慮下 EMA 和 A2Net[6] 的關聯。A2Net 的核心運算元如下：

其中 θ, φ, ρ 代表三個 1×1 卷積，它們的參數分別為 W_θ、W_φ和 W_ρ。如果我們將θ和φ的參數共享，並將 W_θ和 W_φ記作 μ。那麼，softmax(θ(X, W_θ)) 和公式 (5) 無異；而 [.] 內則在更新 μ，即相當於 A_E 和 A_M 迭代一次。因此，A2-Block 可以看作 EMAU 的特殊例子，它只迭代一次 EM，且 μ 由反向傳播來更新。而 EMAU 迭代 T 步，用滑動平均來更新 μ。

實驗

首先是在 PASCOL VOC 上的消融實驗。這裡對比了不同的 μ 更新方法和歸一化方法的影響。

可以清楚地看到，EMA 使用滑動均值（Moving average）和 L2Norm 最為有效。作為對比，Nonlocal 和 A2Net 的模塊作為語義分割頭，在同樣設置下分別達到 77.78% 和 77.34% 的分數，而 EMANet 僅迭代一次時分數為 77.34%，三者無顯著差異，符合上文對 Nonlocal 和 A2Net 的分析和對比。接下來是不同訓練和測試中迭代次數 T 的對比實驗。

可以發現，EMA 僅需三步即可近似收斂（精度不再增益）。而隨著訓練時迭代次數的繼續增長，精度有所下降，這是由 EMA 的 RNN 特性引起的。

接下來，是 EMANet 和 DeeplabV3[3]、DeeplabV3 [7] 和 PSANet[8] 的詳細對比。

可以發現，EMANet 無論在精度還是在計算代價上，都顯著高於表中幾個經典演算法。

在 VOC test server 上，EMANet 在所有使用 ResNet-101 的演算法中，取得了最高的分數。此外，在 PASCAL Context 和 COCO stuff 數據集上也表現卓越。

最後是學習到的注意力圖的可視化。如下圖，I,j,k,l 表示四個隨機選擇的基的下標。右邊四列繪出的是它們各自對應的注意力圖。可以看到，不同的基會收斂到一些特定的語義概念。

參考

[1]Li, Xia, Zhong, Zhisheng, et al. " Expectation Maximization Attention Networks for Semantic Segmentation." Proceedings of the IEEE conference on computer vision. 2019.

[2]Long, Jonathan, Evan Shelhamer, and Trevor Darrell. "Fully convolutional networks for semantic segmentation." Proceedings of the IEEE conference on computer vision and pattern recognition. 2015.

[3]Chen, Liang-Chieh, et al. "Deeplab: Semantic image segmentation with deep convolutional nets, atrous convolution, and fully connected crfs." IEEE transactions on pattern analysis and machine intelligence 40.4 (2017): 834-848.

[4]Zhao, Hengshuang, et al. "Pyramid scene parsing network." Proceedings of the IEEE conference on computer vision and pattern recognition. 2017.

[5]Wang, Xiaolong, et al. "Non-local neural networks." Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2018.

[6]Chen, Yunpeng, et al. "A^ 2-Nets: Double Attention Networks." Advances in Neural Information Processing Systems. 2018.

[7]Chen, Liang-Chieh, et al. "Encoder-decoder with atrous separable convolution for semantic image segmentation." Proceedings of the European conference on computer vision (ECCV). 2018.

[8]Zhao, Hengshuang, et al. "Psanet: Point-wise spatial attention network for scene parsing." Proceedings of the European Conference on Computer Vision (ECCV). 2018.

本文為機器之心專欄，轉載請聯繫本原作者獲得授權。

------------------------------------------------

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！