【速览】ICCV2019 | EMANet：期望最大化注意力网络

学会“成果速览”系列文章旨在将学会会员最新代表性成果进行传播，通过短篇文章让读者用母语快速了解相关学术动态，欢迎关注和投稿~

◆ ◆ ◆ ◆

EMANet：期望最大化注意力网络

李夏、钟之声、吴建龙、杨一博、林宙辰、刘宏

ICCV 2019, Oral

推荐理事：林宙辰

◆ ◆ ◆ ◆

背景介绍

语义分割是计算机视觉领域的一项基础任务，它的目标是为每个像素预测类别标签。由于类别多样繁杂，且类间表征相似度大，语义分割要求模型具有强大的区分能力。近年来，基于全卷积网络（FCN[2]）的一系列研究，在该任务上取得了卓越的成绩。

这些语义分割网络，由骨干网络和语义分割头组成。全卷积网络受制于较小的有效感知域，无法充分捕获长距离信息。为弥补这一缺陷，诸多工作提出提出了高效的多尺度上下文融合模块，例如全局池化层、Deeplab[3]的空洞空间卷积池化金字塔、PSPNet[4]的金字塔池化模块等。

近年来，自注意力机制在自然语言处理领域取得卓越成果。非局部网络（Nonlocal NN[5]）被提出后，在计算机视觉领域也受到了广泛的关注，并被一系列文章证明了在语义分割中的有效性。它使得每个像素可以充分捕获全局信息。然而，自注意力机制需要生成一个巨大的注意力图，其空间复杂度和时间复杂度巨大。其瓶颈在于，每一个像素的注意力图都需要对全图计算。

本文所提出的期望最大化注意力机制（EMA），摒弃了在全图上计算注意力图的流程，转而通过期望最大化（EM）算法迭代出一组紧凑的基，在这组基上运行注意力机制，从而大大降低了复杂度。其中，E步更新注意力图，M步更新这组基。E、M交替执行，收敛之后用来重建特征图。本文把这一机制嵌入网络中，构造出轻量且易实现的EMA Unit。其作为语义分割头，在多个数据集上取得了较高的精度。

前提知识

1、期望最大化算法

期望最大化（EM）算法旨在为隐变量模型寻找最大似然解。对于观测数据

，每一个数据点

都对应隐变量

。我们把

称为完整数据，其似然函数为

，

是模型的参数。

E步根据当前参数

计算隐变量

的后验分布，并以之寻找完整数据的似然

：

M步通过最大化似然函数来更新参数得到

：

EM算法被证明会收敛到局部最大值处，且迭代过程完整数据

的似然值单调递增。

高斯混合模型（GMM）是EM算法的一个范例，它把数据用多个高斯分布拟合。其

即为第

个高斯分布的参数

，隐变量

为第

个高斯分布对第

数据点的“责任”。E步更新“责任”，M步更新高斯参数。在实际应用中，

经常被简化为

。

2、非局部网络

非局部网络（Nonlocal NN[5]）率先将自注意力机制使用在计算机视觉任务中。其核心算子是：

其中

表示广义的核函数，

是归一化系数。它将第

个像素的特征

更新为其他所有像素特征经过

变换之后的加权平均

，权重通过归一化后的核函数计算，表征两个像素之间的相关度。这里

，所以视为像素特征被一组过完备的基进行了重构。这组基数目巨大，且存在大量信息冗余。

期望最大化注意力机制

期望最大化注意力机制由

三部分组成，前两者分别对应EM算法的E步和M步。假定输入的特征图为

，基初始值为

，

估计隐变量

，即每个基对像素的权责。具体地，第

个基对第

个像素的权责可以计算为：

在这里，内核

可以有多种选择。我们选择

的形式。在实现中，可以用如下的方式实现：

其中，

作为超参数来控制

的分布。

步更新基

。为了保证

和

处在同一表征空间内，此处

被计算作

的加权平均。具体地，第

个基被更新为：

值得注意的是，如果

，则公式(5)中，

会变成一组one-hot编码。在这种情形下，每个像素仅由一个基负责，而基被更新为其所负责的像素的均值，这便是标准的K-means算法。

和

交替执行

步。此后，近似收敛的

和

便可以被用来对

进行重估计得

：

相比于

，具有低秩的特性。从下图中

和

的对比中可看出，其在保持类间差异的同时，类别内部差异得到缩小。从图像角度来看，起到了类似保边滤波的效果。

综上，EMA在获得低秩重构特性的同时，将复杂度从Nonlocal的

降低至

。实验中，EMA仅需

步就可达到近似收敛，因此

作为一个小常数，可以被省去。至此，EMA复杂度仅为

。考虑到

，其复杂度得到显著的降低。

期望最大化注意力模块

期望最大化注意力模块（EMAU）的结构如上图所示。除了核心的EMA之外，两个

卷积分别放置于EMA前后。前者将输入的值域从

映射到

；后者将

映射到

的残差空间。囊括进两个卷积的额外负荷，EMAU的FLOPs仅相当于同样输入输出大小时

卷积的1/3，参数量仅为

。

对于EM算法而言，参数的初始化会影响到最终收敛时的效果。上一节中讨论了EMA如何在单张图像的特征图上进行迭代运算。而对于深度网络训练过程中的大量图片，在逐个批次训练的同时，EM参数的迭代初值

理应得到不断优化。本文中，迭代初值

的维护参考BN中running_mean和running_std的滑动平均更新方式，即：

其中，

表示动量的权重；

表示

在一个mini-batch上的平均。

此外，EMA的迭代过程可以展开为一个RNN，其反向传播也会面临梯度爆炸或消失等问题。此外，公式(8) 也要求

和

的差异不宜过大，不然初值

的更新也会出现不稳定。RNN中采取LayerNorm（LN）来进行归一化是一个合理的选择。但在EMA中，LN会改变基的方向，进而影响其语义。因为，本文选择L2Norm来对基进行归一化。这样，

的更新轨迹便处在一个高维球面上。

此处，我们可以考虑下EMA和A2Net[6]的关联。A2Net的核心算子如下：

其中

代表三个

卷积，它们的参数分别为

、

和

。如果我们将

和

的参数共享，并将和记作

。那么，

和公式(5)无异；而

内则在更新

，即相当于

和

迭代一次。因此，A2-Block可以看作EMAU的特殊例子，它只迭代一次EM，且

由反向传播来更新。而EMAU迭代

步，用滑动平均来更新

。

实验

首先是在PASCOL VOC上的消融实验。这里对比了不同的

更新方法和归一化方法的影响。

可以清楚地看到，EMA使用滑动均值（Moving average）和L2Norm最为有效。作为对比，Nonlocal和A2Net的模块作为语义分割头，在同样设置下分别达到 77.78%和77.34%的分数，而EMANet仅迭代一次时分数为77.34%，三者无显著差异，符合上文对Nonlocal和A2Net的分析和对比。接下来是不同训练和测试中迭代次数

的对比实验。

可以发现，EMA仅需三步即可近似收敛（精度不再增益）。而随着训练时迭代次数的继续增长，精度有所下降，这是由EMA的RNN特性引起的。

接下来，是EMANet和DeeplabV3[3]、DeeplabV3+[7]和PSANet[8]的详细对比。

可以发现，EMANet无论在精度还是在计算代价上，都显著优于表中几个经典算法。

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

最后是学习到的注意力图的可视化。如下图，

表示四个随机选择的基的下标。右边四列绘出的是它们各自对应的注意力图。可以看到，不同的基会收敛到一些特定的语义概念。

相关论文

[1] Xia Li, Zhisheng Zhong, Jianlong Wu, Yibo Yang, Zhouchen Lin, and Hong Liu, Expectation Maximization Attention Networks for Semantic Segmentation, ICCV 2019.

[2] Jonathan Long, Evan Shelhamer, and Trevor Darrell, Fully Convolutional Networks for Semantic Segmentation, CVPR 2015.

[3] Liang-Chieh Chen, et al., Deeplab: Semantic Image Segmentation with Deep Convolutional Nets, Atrous Convolution, and Fully Connected CRFs, IEEE TPAMI 40(4): 834-848, 2017.

[4] Hengshuang Zhao, et al., Pyramid Scene Parsing Network, CVPR 2017.

[5] Xiaolong Wang, et al., Non-local Neural Networks, CVPR 2018.

[6] Yunpeng Chen, et al., A^2-Nets: Double Attention Networks, NeurIPS 2018.

[7] Liang-Chieh Chen, et al., Encoder-decoder with Atrous Separable Convolution for Semantic Image Segmentation, ECCV 2018.

[8] Hengshuang Zhao, et al., PSANet: Point-wise Spatial Attention Network for Scene Parsing, ECCV 2018.

往期精选

通知