期望最大化注意力网络用于语义分割——Expectation-Maximization Attention Networks for Semantic Segmentation

0.摘要

        自注意机制已广泛应用于各种任务。它的设计是通过所有位置的特征的加权和来计算每个位置的表示。因此，它可以捕捉计算机视觉任务的长程关系。然而，这种方法在计算上是耗费资源的，因为注意力图是相对于所有其他位置计算的。在本文中，我们将注意机制形式化为期望最大化的方式，并迭代地估计一组更紧凑的基础，基于这些基础计算注意力图。通过对这些基础的加权求和，得到的表示是低秩的，并且将输入中的噪声信息降低。所提出的期望最大化注意力（EMA）模块对输入的方差具有鲁棒性，并且在内存和计算方面也很友好。此外，我们建立了基础维护和规范化方法来稳定其训练过程。我们在流行的语义分割基准数据集，包括PASCAL VOC、PASCAL Context和COCO Stuff上进行了广泛的实验，在这些数据集上创造了新的记录。

1.引言

        语义分割是计算机视觉中的一个基本且具有挑战性的问题，其目标是为图像的每个像素分配一个语义类别。它对于各种任务非常重要，例如自动驾驶、图像编辑和机器人感知。为了有效地完成语义分割任务，我们需要区分一些混淆的类别并考虑不同对象的外观。例如，“草”和“地面”在某些情况下颜色相似，“人”可能在图像的不同位置具有不同的比例、形状和服装。同时，输出的标签空间非常紧凑，特定数据集的类别数量有限。因此，这个任务可以被视为将高维噪声空间中的数据点投影到一个紧凑的子空间中。其实质在于去除这些变化中的噪声并捕捉最重要的语义概念。

最近，许多基于全卷积网络（FCNs）[22]的最先进方法已被提出来解决上述问题。由于固定的几何结构，它们本质上受到局部感受野和短程上下文信息的限制。为了捕捉长程依赖关系，一些作品采用了多尺度上下文融合[17]，例如astrous卷积[4]、空间金字塔[37]、大卷积核卷积[25]等等。此外，为了保留更详细的信息，提出了编码器解码器结构[34,5]以融合中等级和高级语义特征。为了聚合所有空间位置的信息，使用了注意机制[29,38,31]，这使得单个像素的特征可以融合来自所有其他位置的信息。然而，原始的基于注意力的方法需要生成一个大的注意力图，这具有高计算复杂度并占用大量GPU内存。瓶颈在于注意力图的生成和使用都是相对于所有位置计算的。

针对上述问题，本文从期望最大化（EM）算法[7]的角度重新思考注意力机制，并提出了一种新的基于注意力的方法，即期望最大化注意力（EMA）。我们不像之前的方法[38,31]将所有像素本身视为重构基础，而是使用EM算法找到一个更紧凑的基础集，可以大大减少计算复杂度。具体而言，我们将构建基础视为在EM算法中要学习的参数，并将注意力图视为潜在变量。在这种设置下，EM算法旨在找到参数（基础）的最大似然估计。给定当前参数，Expectation（E）步骤用于估计注意力图的期望，而Maximization（M）步骤用于通过最大化完整数据似然函数来更新参数（基础）。E步骤和M步骤交替执行，直到收敛。在收敛后，输出可以计算为基础的加权和，其中权重为规范化后的最终注意力图。EMA的流程如图1所示。我们进一步将提出的EMA方法嵌入到神经网络模块中，命名为EMA单元。EMA单元可以简单地通过常用操作实现。它还非常轻量级，可以轻松嵌入到现有的神经网络中。此外，为了充分利用其容量，我们还提出了两种方法来稳定EMA单元的训练过程。我们还在三个具有挑战性的数据集上评估了其性能。

本文的主要贡献如下： •我们将自注意力机制重新定义为期望最大化迭代方式，可以学习一个更紧凑的基础集，并大大减少计算复杂度。据我们所知，这是第一次将EM迭代引入到注意力机制中。 •我们将所提出的期望最大化注意力构建为神经网络的轻量级模块，并建立了特定的基础维护和规范化方式。 •在三个具有挑战性的语义分割数据集（包括PASCAL VOC、PASCAL Context和COCO Stuff）上进行了广泛的实验，证明了我们的方法优于其他最先进的方法。

图1：所提出的期望最大化注意力方法的流程。

2.相关工作

        语义分割。基于全卷积网络（FCN）[22]的方法通过利用预先在大规模数据上进行预训练的分类网络[14,15,33]强大的卷积特征，在图像语义分割方面取得了巨大的进展。为了增强多尺度上下文聚合，提出了几种模型变种。例如，DeeplabV2 [4]利用星型空间金字塔池化（ASPP）嵌入上下文信息，其中包括具有不同扩张率的并行扩张卷积。DeeplabV3 [4]通过图像级特征扩展ASPP以进一步捕捉全局上下文。同时，PSPNet [37]提出了金字塔池化模块来收集不同尺度的上下文信息。GCN [25]采用大卷积核卷积的解耦来获得特征图的大接受野，并捕获长距离信息。对于另一类变种，它们主要关注于预测更详细的输出。这些方法基于U-Net [27]，结合高级特征和中级特征的优点。RefineNet [21]利用拉普拉斯图像金字塔明确地捕获下采样过程中可用信息，并从粗到细输出预测。DeeplabV3+ [5]在DeeplabV3上添加了一个解码器，以特别改善沿着物体边界的分割结果。Exfuse [36]提出了一个新的框架，以弥合低级特征和高级特征之间的差距，从而提高了分割质量。

        注意力模型。注意力在各种任务中得到了广泛应用，如机器翻译、视觉问答和视频分类。自注意力方法[2,29]通过对句子中所有位置的嵌入进行加权求和，计算出一个位置的上下文编码。Non-local[31]首次将自注意力机制作为计算机视觉任务的模块，如视频分类、物体检测和实例分割。PSANet [38]通过预测的注意力图来学习聚合每个位置的上下文信息。A2Net [6]提出了双重注意块，从整个时空图像空间中分发和收集信息丰富的全局特征。D