基于期望最大化EM思想的语义分割算法

189****6672

论文标题： Expectation-Maximization Attention Networks for Semantic Segmentation

论文地址：http://arxiv.org/pdf/1907.13426

代码地址：https://xialipku.github.io/EMANet

摘要

引出主题：自注意力机制已广泛应用于各种任务。它旨在通过所有位置的特征加权和来计算每个位置的表示。因此，它可以捕捉计算机视觉任务的长远关系。
现存问题：然而，自注意力具有较大的计算消耗。因为注意力图是在所有位置上进行计算的。
解决方法：该论文将注意力机制表述为一种期望最大化方式，并自动估计出一组更紧凑的基，在此基础上计算注意力图。通过对这些基的加权求和，得到的表示是低秩的，并且有效消除来自输入的噪声信息。所提出的期望最大化注意力（EMA）模块对输入方差具有鲁棒性，并且在内存和计算方面也很友好。此外，还建立了规范化方法，以稳定其训练过程。
实验结果：在流行的语义分割基准上进行了广泛的实验，包括PAS-CAL VOC、PASCAL Context和COCO Stuff，并获得了新的记录。

前提知识

Expectation-Maximization Algorithm

期望最大化（EM）算法旨在寻找潜在变量模型的最大似然解。

将X={x1，x2，···，xN}表示为由N个观察样本组成的数据集，每个数据点xi有其相应的潜在变量zi。将{X，Z}称为完整数据，其似然函数的形式为ln p(X，Z |θ)，其中θ是模型的所有参数集。实际上，Z中潜在变量的唯一知识由后验分布p（Z | X，θ）给出。

EM算法设计为通过两个步骤最大化似然ln p(X，Z |θ)，即E步骤和M步骤。

在E步中，使用当前参数θold来找到由p(X，Z |θ)给出的Z的后验分布。
然后，使用Z的后验分布来找到完整数据似然，Q(θ，θold)的期望值，其由下式给出

然后在M步中，通过最大化函数来确定修正参数θnew

EM算法交替执行E步和M步，直到满足收敛准则。

Gaussian Mixture Model

高斯混合模型（GMM）是EM算法的特例。它将数据xn的分布视为高斯分布的线性叠加：

其中平均μk和协方差∑k是第k个高斯基的参数。这里省略了先验πk。完整数据的概率为：

其中∑kznk=1。znk可以被视为第k个个高斯分布对第基对xn的“责任”。对于GMM，在E步骤中，znk的预期值由下式得出

在M步中，参数重新估计如下

其中，

在GMM参数收敛后，重新估计的xnnew可以公式化为：

在实际应用中，可以简单地将∑k替换为单位矩阵I，并在上述等式中省略$∑_k

Non-local

非局部模块的功能与自注意力机制相同。可以用公式表示

其中，f（·，·）表示一般核函数，C（x）是归一化因子，xi表示位置i的特征向量。该模块应用于卷积神经网络（CNN）的特征图上。

考虑到等式（5）中的N(xn |μk，∑k)是介于xn和μk之间的特定核函数，等式（8）只是等式（9）的特定设计。

在GMM中，高斯基的数量是手动选择的，通常满足K << N。但在非局部模块中，基被选择为数据本身，因此它具有K=N。

非局部模块有两个明显的缺点。首先，数据位于低维流形中，因此基过于完整。其次，计算开销大，内存成本也大。

期望最大化注意力

鉴于注意机制的高计算复杂度和非局部模型的局限性，首先提出了期望最大化注意力（EMA）方法，这是自注意的一个增强版本。与选择所有数据点作为基的非局部模块不同，使用EM迭代来寻找紧凑的基集。

为了简单起见，考虑来自单个样本的大小为C×H×W的输入特征映射X。X是CNN的中间激活。为了简化符号，将X重塑大小为N×C，其中N=H×W。

EMA由三个操作组成：

包括责任估计（AE）
似然最大化（AM）
数据重新估计（AR）。

简而言之，给定输入X∈RN×C和初始基μ∈RK×C，AE估计潜在变量（或“responsibility责任”）Z∈RN×K，用作E步。AM使用估计来更新基μ，其作为M步。AE和AM步骤交替执行预先指定的迭代次数。然后，利用收敛的μ和Z，AR将原始X重构为Y并输出。

已经证明，随着EM步骤的迭代，完全数据似然ln P(X，Z)将单调增加。因此，随着AE和AM的迭代，更新的Z和μ具有更好的重构原始数据X的能力。重构的X可以尽可能从原始X中捕捉重要语义。此外，与非局部模块相比，EMA为输入图像找到了一组紧凑的基。该机制消除了许多不必要的噪声，并使每个像素的最终分类更易于处理。此外，该操作将复杂度（在空间和时间上）从O（N2）减少到O（NKT），其中T是AE和AM的迭代次数。同时保证了EM算法的收敛性。值得注意的是，EMA只需要三次迭代就可以在实验中获得可观的结果。因此T可以被视为一个小常数，这意味着复杂度仅为O（NK）。

Responsibility Estimation

责任估计（AE）作为EM算法中的E步骤。该步骤计算znk的预期值，其对应于第k个基μ到xn的责任，其中1≤ k≤ K和1≤ n≤ N、将给定μk的xn的后验概率公式如下：

其中K表示一般的核函数。现在，等式（5）可以重新表述为更一般的形式：

K(a，b)有几种选择，例如内积aTb、指数内积exp(aTb)、欧几里得距离，RBF核经验等。与非局部模块相比，这些函数的选择会在最终结果中产生微小的差异。因此，在本文中选择指数内积exp(aTb)。

在实验中，等式（11）可以实现为矩阵乘法加上一个softmax层。总之，第t次迭代中的AE操作公式如下：

其中λ是控制Z分布的超参数。

Likelihood Maximization

似然最大化（AM）作为EM算法的M步。根据估计的Z，AM通过最大化完整数据的概率来更新μ。

为了使基μ与X位于相同的嵌入空间中，使用X的加权求和来更新基μ。因此μk更新为

值得注意的是，如果在等式（12）中设置λ→ ∞，则{zn1，zn2，··，znK}将成为一个one-hot嵌入。在这种情况下，每个像素仅被分配到一个基。基由分配给它的像素的平均值更新。这就是K-means聚类算法所做的。因此，AE和AM的迭代也可以被视为K-均值聚类的软版本。

Data Re-estimation

EMA交替运行AE和AM T次。然后，使用最终μ（T）和Z（T）重新估计X。采用等式（8）构造新的X，即X~，其公式为

由于X~由紧凑的基集构造而成，与输入X相比，它具有低秩属性。

在图2中描述了示例。很明显，从AR输出的X~在特征空间中非常紧凑，并且对象内部的特征方差小于输入。

EMA Unit

为了更好地将所提出的EMA与深度神经网络相结合，进一步提出了期望最大化注意单元（EMAU），并将其应用于语义分割任务。在本节中，将详细描述EMAU。首先介绍了EMAU的总体结构，然后讨论了基的维护和规范化机制。

Structure of EMA Unit

EMAU的整体结构如图2所示。EMAU乍看起来像是ResNet的bottleneck，只是它用EMA操作取代了沉重的3×3卷积。第一个不带ReLU的卷积被预先将值为从（0+∞) 转换至(−∞,+∞)，这种转换非常重要，否则估计的μ（T）也将位于[0+∞)，这与一般卷积参数相比，容量减半。插入最后1×1卷积，将重新估计X~转换为X的残余空间。

对于每个AE、AM和AR步骤，计算复杂度为O（NKC）。

Bases Maintenance

EM算法的另一个问题是基的初始化。由于完整数据的可能性有限，EM算法保证收敛，并且在每次迭代中，E和M步都会提升其当前下限。然而，不能保证收敛到全局最大值。因此，迭代前基的初始值非常重要。

我们只描述了如何使用EMA处理上述一幅图像。然而，对于计算机视觉任务，一个数据集中有数千个图像。由于每个图像X具有不同于其他图像的像素特征分布，因此不适合使用在图像上计算的μ来重构其他图像的特征图。所以我们在每个图像上运行EMA。

对于第一个小批量，我们使用Kaiming来初始化μ(0)，其中我们将矩阵乘法视为1×1卷积。对于以下小批量，一个简单的选择是使用标准反向传播更新μ(0)。然而，由于AE和AM的迭代可以展开为递归神经网络（RNN），梯度在其中传播将遇到消失或爆炸问题。因此，μ(0)的更新是不稳定的，EMA单元的训练过程可能会失效。

论文在训练过程中使用移动平均更新μ(0)。迭代一张图像后，生成的μ(T)可以视为μ(0)的有偏更新，其中偏差来自图像采样过程。为了减少偏差，首先在小批量上平均μ(T)，得μ¯(T)，然后将μ(0)更新为：