Axial-DeepLab: 独立的轴向注意力用于全景分割（2003）

Axial-DeepLab: Stand-Alone Axial-Attention for Panoptic Segmentation

Axial-DeepLab: 独立的轴向注意力用于全景分割

3 Method

我们首先正式介绍我们的位置敏感自注意力机制。然后，我们讨论如何将其应用于轴向注意力，以及如何使用轴向注意力层构建独立的 Axial-ResNet 和 Axial-DeepLab。

3.1 Position-Sensitive Self-Attention

自注意力机制通常被应用于视觉模型中，作为增强卷积神经网络（CNN）输出的附加组件 [87,95,42]。给定一个输入特征图 $x\in {\mathbb{R}}^{h\times w\times d_{\text{in}}}$，其高度为 $h$，宽度为 $w$，通道数为 $d_{\text{in}}$，在位置 $o=(i,j)$ 的输出 $y_o\in {\mathbb{R}}^{d_{\text{out}}}$ 通过在投影输入上进行池化计算得到：


其中 $N$ 是整个位置格点, query $q_o=W_Q{x}_o$、key $k_o=W_K{x}_o$、value $v_o=W_V{x}_o$ 都是输入 $x_o$ 的线性投影，对于所有 $o\in N$。$W_Q\in {\mathbb{R}}^{d_q\times d_{\text{in}}}$、$W_K\in {\mathbb{R}}^{d_q\times d_{\text{in}}}$ 和 $W_V\in {\mathbb{R}}^{d_{\text{out}}\times d_{\text{in}}}$ 都是可学习的矩阵。softmax 函数应用于所有可能的位置 $p=(a,b)$，在这种情况下，也就是整个二维格点。

这种机制根据亲和力 $x_o^T{W}_Q^T{W}_K{x}_p$ 全局地汇聚值 $v_p$，使我们能够捕获整个特征图中相关但非局部的上下文信息，这与仅捕获局部关系的卷积操作形成对比。

然而，自注意力机制在计算上非常昂贵（时间复杂度为 $O(h^2{w}^2)$），当输入的空间维度较大时，这限制了它仅能用于卷积神经网络（CNN）的高层（即下采样的特征图）或小尺寸图像。另一个缺点是全局汇聚操作没有利用位置信息，而在视觉任务中，位置信息对于捕获空间结构或形状至关重要。

这两个问题在 [68] 中通过向自注意力机制添加局部约束和位置编码得到了缓解。对于每个位置 $o$，提取一个局部的 $m\times m$ 正方形区域作为计算输出 $y_o$ 的记忆库。这显著地将计算量减少到 $O(hw{m}^2)$，使得自注意力模块可以作为独立的层部署，从而形成一个完全基于自注意力的神经网络。此外，一个学习到的相对位置编码项被整合到亲和力中，产生了一个动态的先验，即在感受野（即局部的 $m\times m$ 正方形区域）中应该关注的位置。正式地，[68] 提出

位置信息（表明了联系），计算效率很重要
$q_o^T{r}_{p-o}$ 在下文看来是表示位置 p 和 o 之间相对位置的信息表示

其中 $N_{m\times m}(o)$ 是以位置 $o=(i,j)$ 为中心的局部 $m\times m$ 正方形区域，可学习的向量 $r_{p-o}\in {\mathbb{R}}^{d_q}$ 是添加的相对位置编码。内积 $q_o^T{r}_{p-o}$ 衡量了从位置 $p=(a,b)$ 到位置 $o=(i,j)$ 的兼容性。我们不考虑绝对位置编码 $q_o^T{r}_p$，因为它们相比相对位置编码的泛化能力较差 [68]。在接下来的段落中，为了简洁起见，我们省略“相对”这个词。

在实践中，$d_q$ 和 $d_{\text{out}}$ 通常远小于 $d_{\text{in}}$，并且可以将公式 (2) 中的单头注意力扩展为多头注意力，以捕捉多种亲和力的混合。具体来说，多头注意力是通过并行地对 $x_o$ 应用 $N$ 个单头注意力（每个头使用不同的 $W_{Qn}$、$W_{Kn}$、$W_{Vn}$，其中 n ∈ { 1 , 2 , … , N } n \in \{1, 2, \ldots, N\} n∈{1,2,…,N} 表示第 $n$ 个头），然后通过连接每个头的结果来获得最终输出 $z_o$，即 $z_o={\text{concat}}_n(y_{on})$。注意，位置编码通常在各个头之间共享，因此它们引入的额外参数非常少。

根据我之前看的多头注意力的机制是把嵌入进行分解，比如嵌入维度是d_dim，头数为N，那么每隔头处理的嵌入的维度为d_dim/N，应该没错吧，位置编码怎么在多头之间分享？分割嵌入维度的时候把位置编码也分割了啊？？？

位置敏感性：我们注意到，先前的位置偏差仅依赖于查询像素 $x_o$，而不是键像素 $x_p$。然而，键像素 $x_p$ 也可能包含有关应关注哪个位置的信息。因此，除了查询依赖的偏差 $q_o^T{r}_{p-o}^q$ 之外，我们还添加了一个键依赖的位置偏差项 $k_p^T{r}_{p-o}^k$。

键像素怎么理解呢？

同样地，在公式 (2) 中，值 $v_p$ 不包含任何位置信息。在感受野或记忆库较大时，$y_o$ 很可能不包含 $v_p$ 来自的确切位置。因此，以往的模型不得不在使用较小的感受野（即较小的 $m\times m$ 区域）和丢弃精确的空间结构之间做出权衡。在本工作中，我们使输出 $y_o$ 能够根据查询-键亲和力 $q_o^T{k}_p$ 检索相对位置 $r_{p-o}^v$，除了内容 $v_p$ 之外。正式地，


其中可学习的 $r_{p-o}^k\in {\mathbb{R}}^{d_q}$ 是键的位置编码，而 $r_{p-o}^v\in {\mathbb{R}}^{d_{\text{out}}}$ 是值的位置编码。这两个向量不会引入太多参数，因为它们在一个层中的各个注意力头之间是共享的，并且局部像素的数量 $|N_{m\times m}(o)|$ 通常较小。

我们称这种设计为位置敏感自注意力，它能够在合理的计算开销下捕捉长距离交互并保留精确的位置信息，这一点已在我们的实验中得到验证。

图 1. 非局部块（左） 与我们 沿宽度轴应用的位置敏感轴向注意力（右） 的对比。“$\otimes$” 表示矩阵乘法，“$\oplus$” 表示逐元素求和。softmax 操作是在最后一个轴上进行的。蓝色框表示 1×1 卷积，红色框表示相对位置编码。通道数 $d_{\text{in}}=128$，$d_q=8$，和 $d_{\text{out}}=16$ 是我们在 ResNet 的“stem”阶段之后第一阶段所使用的参数。

3.2 Axial-Attention

由独立自注意力模型 [68] 提出的局部约束显著降低了视觉任务中的计算成本，并使得构建完全基于自注意力的模型成为可能。然而，这种约束牺牲了全局连接，使得注意力的感受野不会大于具有相同核大小的深度可分离卷积。此外，局部自注意力在局部正方形区域内执行，其复杂度仍然与区域长度的平方成正比，引入了另一个超参数来在性能和计算复杂度之间进行权衡。在本工作中，我们提出在独立自注意力中采用轴向注意力 [42,33]，以确保全局连接和高效计算。具体来说，我们首先定义图像宽度轴上的轴向注意力层为简单的一维位置敏感自注意力，并对高度轴使用类似的定义。具体而言，沿宽度轴的轴向注意力层定义如下。


一个轴向注意力层沿一个特定轴传播信息。为了捕获全局信息，我们分别对高度轴和宽度轴连续使用两个轴向注意力层。这两个轴向注意力层都采用了上述描述的多头注意力机制。

轴向注意力将复杂度降低到 $O(hwm)$。这使得全局感受野成为可能，通过直接将跨度 $m$ 设置为整个输入特征的大小来实现。可选地，也可以使用一个固定的 $m$ 值，以减少在巨大特征图上的内存占用。

轴向残差网络（Axial-ResNet）：要将残差网络（ResNet）[32]转换为轴向残差网络，我们将残差瓶颈块中的3×3卷积替换为两个多头轴向注意力层（一个用于高度轴，另一个用于宽度轴）。在每个轴向注意力层之后，可选地在相应轴上执行步长操作。保留两个1×1卷积用于洗牌特征。这就形成了我们的（残差）轴向注意力块，如 图2 所示，该块多次堆叠以获得轴向残差网络。请注意，我们在两个轴向注意力层之间不使用1×1卷积，因为矩阵乘法（$W_Q,W_K,W_V$）紧随其后。此外，原始残差网络中的主干（即，第一个步长为7×7的卷积和3×3的最大池化）被保留，从而得到一个卷积茎模型，在该模型中，第一层使用卷积，其他地方都使用注意力层。在卷积茎模型中，我们将跨度 m 设置为从第一个块开始的整个输入，此时特征图的大小为56×56。

图2. 一个轴向注意力块，由两个轴向注意力层组成，这两个层依次沿着高度轴和宽度轴操作。通道数 $d_{in}=128、d_{out}=16$ 是我们使用在残差网络（ResNet）“茎”之后第一阶段的参数。我们采用了 $N=8$ 个注意力头。

在我们的实验中，我们还构建了一个完整的轴向注意力模型，称为全轴向残差网络（Full Axial-ResNet），它进一步将轴向注意力应用于主干。我们没有设计一个特殊的空间变化的注意力主干[68]，而是简单地堆叠了三个轴向注意力瓶颈块。此外，为了降低计算成本，我们在全轴向残差网络的前几个块中采用了局部约束（即，如[68]中的局部 m×m 方形区域）。

stem怎么理解？主干？

轴向DeepLab：为了将轴向残差网络（Axial-ResNet）进一步转换为用于分割任务的轴向DeepLab，我们进行了以下讨论的几项更改。

首先，为了提取密集的特征图，DeepLab [13] 改变了残差网络（ResNet）[32] 最后一个或两个阶段的步长和扩张率。类似地，我们去除了最后一个阶段的步长，但我们没有实现“扩张”注意力模块，因为我们的轴向注意力已经能够捕获整个输入的全局信息。在这项工作中，我们提取输出步长（即输入分辨率与最终主干特征分辨率的比率）为16的特征图。我们没有追求输出步长为8，因为这在计算上过于昂贵。

“最后一个或两个阶段”指的是该网络结构中最后的部分，“步长”是指在卷积操作中，卷积核移动的步数****，“扩张率”是指卷积核感受野的扩张程度。

其次，我们没有采用扩张空间金字塔池化模块（ASPP）[14,15]，因为我们的轴向注意力块也能够高效地编码多尺度或全局信息。我们在实验中表明，我们的轴向DeepLab在没有ASPP的情况下，性能超过了有和无的ASPP的全景DeepLab [20]。

最后，遵循全景DeepLab [20]，我们采用了完全相同的三个卷积的茎部 [81]、双解码器和预测头。这些头产生语义分割和类别无关的实例分割，并通过多数投票 [93] 合并以形成最终的全景分割。

在输入非常大（例如，2177×2177）且内存受限的情况下，我们在所有的轴向注意力块中采用了较大的跨度 m = 65。请注意，我们不将轴向跨度视为超参数，因为它已经足够覆盖几个数据集上的长距离甚至全局上下文，而且设置较小的跨度并不会显著减少乘加运算次数（M-Adds）。