大可分离核注意力： 重新思考 CNN 中的大核注意力设计

今天刚读大可分离核注意力这篇论文，可与大家一起学习。感兴趣的可以下载论文看一下。
论文下载地址：https://arxiv.org/pdf/2309.01439.pdf
代码地址：https://github.com/StevenLauHKHK/Large-Separable-Kernel-Attention/blob/main/mmsegmentation/van.py
摘要
带有大内核注意力（LKA）模块的视觉注意力网络（VAN）已被证明在一系列基于视觉的任务中具有超越视觉转换器（ViT）的卓越性能。然而，随着卷积核大小的增加，这些 LKA 模块中的深度卷积层会导致计算和内存占用的二次方增加。为了缓解这些问题，并在 VAN 的注意力模块中使用超大卷积核，我们提出了大型可分离核注意力模块系列，称为 LSKA。LSKA 将深度卷积层的二维卷积核分解为层叠的水平和垂直一维核。与标准的 LKA 设计不同的是，所提出的分解方法可以在注意力模块中直接使用具有大内核的深度卷积层，而不需要任何额外的模块。我们证明，VAN 中的拟议 LSKA 模块可以达到与标准 LKA 模块相当的性能，而且计算复杂度和内存占用更低。我们还发现，随着内核大小的增加，拟议的 LSKA 设计会使 VAN 更偏向于物体的形状，而不是纹理。此外，我们还对 VAN、ViTs 和最近的 ConvNeXt 中的 LKA 和 LSKA 的鲁棒性进行了基准测试，测试对象是 ImageNet 数据集的五个损坏版本，这些版本在以前的工作中基本上都未进行过研究。我们的大量实验结果表明，随着内核大小的增加，VAN 中提出的 LSKA 模块显著降低了计算复杂度和内存占用，同时在对象识别、对象检测、语义分割和鲁棒性测试方面的表现优于 ViTs 和 ConvNeXt，与 VAN 中的 LKA 模块相比也具有相似的性能。
相关工作
在视觉的任务中，使用注意力机制可进一步提高 CNN 的性能，注意力机制能产生与人眼和大脑一致的突出区域，从而提高 CNN 的性能。尽管带有注意力机制的 CNN 非常有效，但最近在视觉领域，基于自注意力的网络（如VIT）及其变体的适应性激增，在图像分类、物体检测和语义分割方面的表现优于 CNN。ViT 的性能更优越，这要归功于 Transformers 使用多头自注意力（MHSA）作为其关键组件。然而，ViTs 在视觉领域的有效性是以高分辨率输入图像时计算和内存占用的四倍增长为代价的。
CNN 中的这种长距离依赖性可以通过使用具有注意机制的更大感受野来建模 。CNN 中的大感受野是通过堆叠许多卷积操作或使用更大的内核尺寸获得的。 前一种方法可能会导致模型体积增大，而后一种方法则因其内存和计算要求而被认为成本过高。
有研究表明，CNN 中的大内核可以通过级联深度卷积和扩张深度卷积来建模，而不会导致计算和内存占用的二次方上升 。如图 2c 所示，他们提出的视觉注意力网络（VAN）使用了具有简单注意力机制的内核堆栈，命名为大内核注意力（LKA）。LKA 模块采用标准的深度卷积和小感受野核来捕捉局部依赖性并补偿网格问题，然后采用扩张的深度卷积和大感受野核来模拟远距离依赖性。这种深度卷积和扩张深度卷积的组合相当于 大规模 CNN 内核。稀释深度卷积的输出被输入到 1 × 1 卷积中，以推断注意力图。然后将输出的注意力图与输入特征相乘，进行自适应特征提纯。
在图像分类、物体检测和语义分割方面，带有 LKA 的 VAN 已被证明优于最先进的 ViT 和 CNN。然而，大规模深度卷积核的设计仍然会产生高计算量和内存占用，随着核大小的增加，模型的有效性也会降低。当核大小达到 35 × 35 和 53 × 53 时，在 VAN 的 LKA 模块中设计深度卷积（不使用深度扩张卷积）的计算效率很低。

在本文首先研究了在 VAN 中使用简单的注意力模块进行深度卷积对大内核的影响。将这种带有大内核和注意力模块的深度卷积设计称为 LKA-trivial。其次，本文提出了 VAN 深度卷积的可分离版本。深度卷积的可分离版本将给定的 k × k 卷积核平均分为 1 × k 和 k × 1 可分离卷积核，这些卷积核以级联方式作用于输入特征。在其他因素保持不变的情况下，VAN 的 LKA 三倍模块中的深度卷积的可分离版本大大减少了随着内核大小增加而参数数量的二次增长。将这种改进的 LKA-trivial 设计称为 LSKA-trivial。此外，这种内核分解也与深度扩张卷积兼容，因此本文可以在 VAN 中提出一个完全可分离的 LKA 模块版本 LSKA。研究表明，本文所提出的 LSKA 版本的 LKA 在 VAN 中获得了类似的性能，同时即使在更大的内核尺寸下也具有很高的计算效率。此外，拟议的 LSKA 模块在更大的内核尺寸下增强了输入图像的远距离依赖性，而不会产生高计算量和内存占用。
论文工作的贡献总结如下：

1. 我们解决了深度卷积核的计算效率低下问题，即随着核大小的增加，LKA-trivial 和 LKA in VAN 的计算效率低下。我们的研究表明，用级联的 1 × k 和 k × 1 卷积核取代深度卷积中的 k × k 卷积核，可以有效减少 LKA-trivial 和 LKA in VAN 随着核大小增加而产生的参数数量二次增长，而性能却没有任何下降。
2. 我们在各种基于视觉的任务（包括图像分类、物体检测和语义分割）中通过实验验证了 LSKA 在 VAN 中的有效性。我们证明，与原始 LKA 中的小内核相比，LSKA 可以从大内核中获益，同时保持相同的推理时间成本。
3. 我们在 5 个不同的 ImageNet 数据集上对基于 LKA 的 VAN、基于 LSKA 的 VAN、ConvNeXt 和最先进的 ViT 的鲁棒性进行了基准测试，这些数据集包含应用于图像的各种类型的扰动。结果表明，与之前的大核 CNN 和 ViT 相比，基于 LSKA 的 VAN 是一种稳健的学习器。
4. 我们提供的定量证据表明，与 ViT 和以前的大核 CNN 相比，基于 LSKA 的 VAN 中的大核学习到的特征编码了更多的形状信息和更少的纹理信息。此外，特征表征中编码的形状信息量与对不同图像失真的鲁棒性之间存在高度相关性。这些证据有助于我们解释为什么基于 LSKA 的 VAN 是一种鲁棒学习器。

A. 具有大内核的 CNN
过去十年间，为改进用于一般图像识别任务的卷积神经网络（CNN）架构，人们开展了大量研究工作。除 AlexNet 和 Inception网络外，这些 CNN 架构的内核设计主要限于 3×3（如 VGG、ResNet 、MobileNets ），原因是随着权重层深度的增加，计算效率也会提高。为了尝试使用大内核尺寸，提出了一种由局部关系层组成的局部关系网络（LR-Net），其内核尺寸相对较大，为 7×7。虽然 LR-Net 的性能优于传统的 ResNets（内核大小相同），但当内核大小进一步增大时，LR-Net 的性能就会下降。为了缩小分层变换器和 CNN 之间的性能差距，ConvNeXt 通过将变换器设计逐步借入 ResNet 进行了实证研究。他们发现了提高 CNN 性能的几个关键要素，如改变 Swin Transformer的训练程序、改变阶段计算比、使用更少的激活层和归一化层以及使用更大的核大小。与 LR-Net 类似，他们发现当核大小超过 7 × 7 时，性能就会达到饱和。
最近，作者在论文中重新审视了 CNN 中的大内核设计，这种设计曾长期处于阴影之下。他们证明，在 MobileNetV2 中用几个重新参数化的大重量内核取代小重量内核集，可以产生大的有效感受场（ERF），并部分模拟人类对物体形状的理解。他们提出的 RepLKNet 采用 31 × 31 的大内核尺寸，在 ImageNet 分类上比 Swin Transformers 优胜 0.3%，在 MS-COCO 检测上比 ResNet-101 优胜 4.4%。然而，RepLKNet 的计算量很大，因此限制了它在其他领域（如分割）的有效性。例如，作者在文献中指出，大核卷积会损害分割任务的性能，因为当参数大小随核大小增加时，模型会出现过拟合问题。为了克服这一问题，他们提出了使用 1×k 和 k×1 大卷积核的全局卷积网络（GCN），以提高语义分割性能。
另一项最新研究成果 SLaK 观察到，当内核大小超过 31、51 和 61 时，RepLKNet 的性能开始趋于稳定。为了解决超大内核的可训练性问题，SLaK 将大内核分解为两个矩形内核（即 51 × 5 和 5 × 51），并使用动态稀疏技术来减少可学习参数。与这些方法不同的是，我们在基于 CNN 的视觉注意力网络（VAN）的深度卷积和深度稀释卷积中使用了可分离内核，以进一步提高其计算效率。

B. 大内核注意力机制
注意机制用于选择图像中最重要的区域。一般来说，可分为四类：空间注意 ；通道注意；时间注意 和分支注意 。在此，我们更关注通道注意和空间注意，因为它们与我们的工作更为相关。通道注意力侧重于特定模型层的 "什么 "语义属性。由于特征图的每个通道都是一个检测器（也称为滤波器）的响应图，通道注意力机制允许模型将注意力集中在各通道中物体的特定属性上。与通道注意力不同，空间注意力关注的是模型应该关注的语义相关区域的 “位置”。STN 、GENet 和非局部神经网络是涉及不同类型空间注意力方法的一些代表性作品。自注意是空间注意的一种类型，已被用于 CNN 和 ViT 。视觉注意力网络（VAN）提出了一种新的 CNN 骨干网络，它结合了 CNN 的特性和自注意力模块。
作者采用了带有大核注意力（LKA）模块的 CNN 网络架构，以实现 CNN 的局部性、长程依赖性和自注意力模块的空间自适应特性，类似于 ViT。此外，LKA 模块还具有通道适应性，这是标准 CNN 和Transformer中的自注意模块所不具备的。为了提高计算效率，LKA 采用了带深度卷积（DW-D-Conv）的扩张卷积，以获得更大的 ERFs。
在图像分类、物体检测和语义分割方面，VAN 甚至比 PVT-V2、Swin Transformer 和 Twins-SVT 等一系列变压器网络取得了更好的性能。

方法
在本节中，我们首先讨论如何通过使用一维卷积核来重组 LKA 模块（使用和不使用扩张深度卷积），从而设计出 LSKA 模块。然后，我们总结了 LSKA 模块的几个关键特性，接着分析了 LSKA 的复杂性。
A. 配方
如图 2a 所示，我们首先设计基本的 LKA 模块，但不使用扩张深度卷积。给定一个输入特征图 $F\in R^{C\times H\times W}$，其中 C 是输入通道的数量，H 和 W 分别代表特征图的高度和宽度。LKA 的输出可通过公式 1-3 得出。
$$Z^C=\sum _{H,W}{W}_{k\ast k}^C\ast F^C\text{\hspace{1em}(1)}$$
$$A^C=W_{1\ast 1}\ast Z^C\text{\hspace{1em}(2)}$$
$$\overline{F}=A^C⨂F^C\text{\hspace{1em}(3)}$$
其中 ∗ 和 ⊗ 分别表示卷积和哈达玛乘积。Z^C 是深度卷积的输出，由大小为 k * k 的核 W 与输入特征图 F 卷积得到。需要注意的是，F 中的每个通道 C 都要与核 W 中的相应通道进行卷积[41]。公式 1 中的 k 也代表核 W 的最大感受野。然后，深度卷积的输出与 1 × 1 内核卷积，得到注意力图 AC。LKA 的输出̄ F C 是注意力图谱 AC 与输入特征图谱 F C 的哈达玛乘积。可以看出，随着内核大小的增加，LKA 模块中的深度卷积将导致计算成本的二次增长。我们在本文中将这种设计称为 LKAtrivial，以区别于 [34] 中提到的 LKA 实际设计。我们很快就会发现，增加 LKA-trivial 内核的大小会导致 VAN 的计算复杂度二次方（k2）增长（见图 3）。
为了缓解 LKA-trivial（LKA-trivial）中较大核大小的深度卷积计算成本高的问题，作者在 [34] 中将大核的深度卷积分解为小核的深度卷积和相当大核的扩张深度卷积（图 2c）。这种类型的大核分解有助于缓解单独的深度卷积因核尺寸大而导致计算成本二次增加的问题。正如文献[34]所述，LKA 的输出结果可按如下方式获得。
$${\overline{Z}}^C=\sum _{H,W}{W}_{(2d-1)\ast (2d-1)}^C\ast F^C\text{\hspace{1em}(4)}$$
$$Z^C=\sum _{H,W}{W}_{\lfloor \frac{k}{d}\rfloor \ast \lfloor \frac{k}{d}\rfloor }^C\ast {\overline{Z}}^C\text{\hspace{1em}(5)}$$
$$A^C=W_{1\ast 1}\ast Z^C\text{\hspace{1em}(6)}$$
$${\overline{F}}^C=A^C⨂F^C\text{\hspace{1em}(7)}$$

其中，d 是扩张率。 式 4 中的̄ ZC 表示核大小为 (2d - 1) × (2d - 1) 的深度卷积输出，它捕捉了局部空间信息，并补偿了接下来核大小为 ( k d⌋×⌊k d ⌋)的深度扩张卷积（式 5）的网格效应。注意，⌊.⌋ 代表下限运算。扩张深度卷积负责捕捉深度卷积输出̄ ZC 的全局空间信息。虽然 LKA 设计大大改进了 LKA-trivial，但当内核大小超过 23 × 23 时，在 VAN 中仍会产生较高的计算复杂度和内存占用（见图 3）。

B. 有注意力的大型可分离内核
通过将深度卷积和深度扩张卷积的二维权重核拆分为两个级联的一维可分离权重核，可以获得等效的改进型 LKA 配置。如图 2d 所示，我们将这种改进的 LKA 模块配置称为 LSKA。根据文献[34]，LSKA 的输出结果如下。
$${\overline{Z}}^C=\sum _{H,W}{W}_{(2d-1)\ast 1}^C\ast (\sum _{H,W}{W}_{1\ast (2d-1)}^C\ast F^C)\text{\hspace{1em}(8)}$$
$${\overline{Z}}^C=\sum _{H,W}{W}_{\lfloor \frac{k}{d}\rfloor \ast 1}^C\ast (\sum _{H,W}{W}_{1\ast \lfloor \frac{k}{d}\rfloor )}^C\ast {\overline{Z}}^C)\text{\hspace{1em}(9)}$$
$$A^C=W_{1\ast 1}\ast Z^C\text{\hspace{1em}(10)}$$
$${\overline{F}}^C=A^C⨂F^C\text{\hspace{1em}(11)}$$

如图 2b 所示，LKA-trivial 的可分离版本也可以用类似的方法得到，我们将其命名为 LSKA-trivial。从图 3 可以看出，与 LKA-trivial 和 LKA 相比，LSKA-trivial 和 LSKA 都大大降低了 VAN 的计算复杂度。在下面的小节中，我们将报告 LSKA 区别于一般卷积、自注意和 LKA 的特性。
C. LSKA 的特性
通过重新审视以往的注意机制 [27]、[28]、[31]、[32]，LSKA 模块与一般卷积、自我注意和 LKA 模块相比有四个重要特性，如表一所示。长程依赖性。如第二节 B 所述，自注意机制 [52] 是转换器建立长程依赖关系模型的关键组件。不过，最近的研究 [33] 和 [35] 表明，大内核是捕捉全局信息的另一种方法。为了实现这一目标，我们采用了文献[34]中的 LKA 设计，即把一个大内核分解成两个小内核，而不是使用文献[33]、[35]中的琐碎大内核设计，因为它的计算量大，优化难度高[43]。为了验证我们提出的 LSKA 的长程依赖性，我们利用[?]、[35]、[44]中描述的有效感受野（ERF）生成方法，生成了内核大小从 7 到 65 的 VAN-LSKA-Tiny 的 ERF 图，如图 4 所示。热图上分布的暗区范围越大，表示 ERF 越大。从图 4a 到图 4f，我们观察到暗区从内核大小 7 扩展到 65，这表明 LSKA 方法可以有效捕捉图像中的长程依赖性。
空间和通道适应性。空间注意力和信道注意力是两种常见的策略，它们可以根据上下文依赖关系自适应地重新校准特征的权重，这一点在第二节 B 部分已经提到。我们的工作继承了 LKA 的设计，与自我注意相比，LKA 包括这两种特性，但参数和计算复杂度更低。如图 2d 所示，LKA 和 LSKA 的不同之处在于我们采用了层叠式水平和垂直内核，以进一步降低内存和计算复杂度。
对超大内核的可扩展性。如图 3 所示，随着核大小的增加，VAN 中的 LKA 三倍计算成本呈二次方增长。LKA 设计大大减少了计算占用空间，但当内核大小超过 23 × 23 时，模型参数的数量也会增加。如果在 VAN 中引入最近最先进的 SLaK-Decomposed 和 SLaK-SparseDecomposed 方法，当内核大小超过 100 时，它们产生的参数数和计算足迹都低于 LKA。请注意，图 3 中报告的是 VAN-Small 网络的结果。令人惊讶的是，与 LKA 和 SLaK 相比，分别提出的 LSKA-trivial 和 LSKA 版本的 LKA-trivial 和 LKA 不仅降低了计算成本，还使 VAN 的模型参数数保持相对不变。请注意，内核大小也指最大接收场（MRF）。
在精确度方面，如表 VIII 所示，LSKA-Base 在内核大小从 23 增大到 53 时表现出持续增长。相反，当内核大小超过 23 时，LSKA-Base 就开始饱和。这些结果表明，LSKA 在参数大小、GFLOPs 和精度方面都表现出了对超大内核的可扩展性。
D. LSKA 的复杂性分析
在本小节中，我们将计算拟议的 LSKAtrivial、LSKA、LKA-trivial 和 LKA 注意模块的浮点运算（FLOP）次数和参数，如图 3 所示。请注意，为简化计算，以下分析中忽略了偏置项。我们还假设 LSKA 和 LKA 的特征图输入大小和输出大小相同（即 H × W × C）。为简洁起见，我们只提供 LSKA 和 LKA 的 FLOPs 和参数计算公式。不过，同样的方程也可以用来计算 LKA-trivial 和 LSKA-trivial 的参数和 FLOP。原始 LKA 的参数和 FLOP 计算如下：
$$Param=(2d-1{)}^2\ast C+\lfloor \frac{k}{d}{\rfloor }^2\ast C+C\ast C\text{\hspace{1em}(12)}$$
$$FLOPs=((2d-1{)}^2\ast C+\lfloor \frac{k}{d}{\rfloor }^2\ast C+C\ast C)\ast H\ast W\text{\hspace{1em}(13)}$$
其中，k 为内核大小，d 为扩张率。LSKA 注意模块的 FLOP 总数和参数可按下式计算：
$$Param=(2d-1{)}^2\ast C+\lfloor \frac{k}{d}\rfloor \ast C\ast 2+C\ast C\text{\hspace{1em}(14)}$$
$$FLOPs=((2d-1{)}^2\ast C+\lfloor \frac{k}{d}\rfloor \ast C\ast 2+C\ast C)\ast H\ast W\text{\hspace{1em}(15)}$$

将公式 14 和 12 中的第一项相等，我们可以发现，所提出的 LSKA 可以在原始 LKA 设计的深度卷积层中节省 2d-1 2 个参数。同样，通过比较式 14 和式 12 的第二项，我们可以发现，所提出的 LSKA 可以在原始 LKA 设计的扩张深度卷积层中节省 1 2 k d 个参数。所节省的触发次数与参数相同。我们还可以看到，LSKA 在计算上比 LSKAtrivial 更有效。因此，除非另有说明，在与 LKA 和最先进的方法进行比较时，我们报告的是 LSKA 的性能。
E. 模型架构
在这项工作中，我们沿用了 VAN [34] 的结构设计，如表 II 所示。模型的超参数如下：

• Si：第 i 阶段输入干和下采样卷积层的步长；
  
  - Ki：第 i 阶段输入干和下采样卷积层的核大小；
• Ci：第 i 阶段输出通道数；
• Ei：第 i 阶段卷积前馈层的扩展比；
• Li：第 i 阶段的块数；
  按照 VAN 的设计，我们的模型由一个输入干层和四个后续阶段组成。输入干线的第一层包含一个步长为 4 的 7 × 7 卷积层，然后是一个批处理归一化层。这一层将输入分辨率降低 4，并根据模型容量将通道增加到 32 或 64。
  除第 1 阶段外，每个阶段都以 3 × 3 卷积层和跨度为 2 的降采样层开始。如图 5 所示，随后是一个卷积块，其中包含批量归一化、LSKA 模块和一个卷积前馈网络（CFFN）。我们的模型在深度卷积之前和之后都包含一个 1 × 1 卷积层，用于通道交互。为了提供更多的非线性，在 LSKA 之前和 CFFN 内部附加了 GELU 激活层。VAN 中的 LKA 与我们的工作的主要区别在于，我们在每个卷积块中都用我们的 LSKA 层取代了 LKA 层。
  为了提供更多讨论实例，我们设计了三种不同容量的带 LSKA 的 VAN，即 VAN-LSKATiny、VAN-LSKA-Small 和 VAN-LSKA-Base。这些模型包含的卷积块数、通道数和 CFFN 的扩展率与带 LSKA 的 VAN 相同。
  实验