基于期望最大化EM思想的语义分割算法
-
论文标题: 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)更新为:
实验