HaloNets

https://arxiv.org/abs/2103.12731 CVPR2021 Oral

Scaling Local Self-Attention For Parameter Efficient Visual Backbones

Ashish Vaswani, Prajit Ramachandran, Aravind Srinivas, Niki Parmar, Blake Hechtman, Jonathon Shlens

https://hub.fastgit.org/lucidrains/halonet-pytorch unofficial

https://hub.fastgit.org/qwopqwop200/halonet-tf2 unofficial

timm中也有实现：

https://hub.fastgit.org/rwightman/pytorch-image-models/search?q=halo

https://hub.fastgit.org/rwightman/pytorch-image-models/blob/master/timm/models/layers/halo_attn.py

https://hub.fastgit.org/rwightman/pytorch-image-models/blob/master/timm/models/byoanet.py

自注意力由于感受野的缩放是与参数量无关的，并且是基于内容的交互（parameter-independent scaling of receptive fields and content-dependent interactions），而不像卷积那样是基于参数量的缩放、以及交互是不基于内容的（parameter-dependent scaling and content-independent interactions），自注意力有望改善计算机视觉系统。与基线卷积模型（如ResNet-50）相比，自注意模型在准确度参数权衡方面有令人鼓舞的改进。在这项工作中，我们的目标是发展自注意模型，它不仅可以超越规范基线模型，甚至可以超越高性能的卷积模型。我们提出了两个自注意的扩展，结合更高效的自注意实现，提高了这些模型的速度、内存使用率和准确性。我们利用这些改进来开发一个新的自注意模型族，即HaloNets，它在ImageNet分类基准的参数限制设置上达到了最先进的精度。在初步的迁移学习实验中，我们发现HaloNet模型的性能要比更大的模型好得多，并且具有更好的推理性能。在目标检测和实例分割等较难的任务上，我们简单的局部自注意力和卷积混合算法比非常强的基线有了改进。这些结果标志着在证明自注意模型在传统上由卷积模型支配的环境中的有效性efficacy方面又迈出了一步。

1. 引言

自注意是自然语言处理的主要工具，而卷积在几乎所有的视觉模型中都是普遍存在的。卷积算法体现了局部处理的原理，能够学习图像中丰富的边缘、纹理等局部空间特征。另一方面，Transformer[57]表明，自注意是一种有效的、计算效率高的机制，用于捕捉句子中单词之间的全局交互。自注意在自然语言处理中的成功推动了对自注意如何改善视觉系统的研究。自注意有几个特性使其非常适合视觉：（a）基于内容的交互作用，而不是卷积的内容无关的交互；（b）感受野大小与参数量无关，而不是卷积那样感受野取决于参数量（编者注：在卷积中，7×7卷积比3×3卷积感受野大，但是参数也更多，而在自注意力中，7×7的自注意力与3×3自注意力的参数量是相同的）；（c）已经被实验证明在大图像中可以捕获长程依赖性；（d）可以灵活地处理和利用视觉领域中不同类型的数据，例如像素[59、3、44、66]、点云[63]、序列条件信息[62]和图[32]。自注意也可以看作是一种自适应的、非线性计算，而计算机视觉中的非线性处理技术（如双边滤波[39]和非局部方法[4]）有着悠久的历史。

最近的几篇论文[3，43，12，66，50]尝试使用自注意来提高图像分类的准确性，而不是使用强大和常用的ResNet主干[17，18]。其中，独立自我注意（Stand-Alone Self-Attention ，SASA）[43]是一种完全自注意模型，它用局部自注意代替了每个空间卷积，在减少参数和浮点运算的同时提高了ResNet主干的性能。虽然这些模型在概念上很有前途，但在图像分类中却落后于最先进的卷积模型。最先进的卷积模型[55，67，42]使用多种缩放技术，在一系列计算和参数范围内实现强大的性能。在这项工作中，我们的目标是发展和了解缩放(scaling)局部自注意力模型的方法，以超过一些最好的卷积模型。扩展(scaling)自注意模型有一系列挑战。例如，卷积已经非常高效地映射到矩阵加速器，例如驱动大多数深度学习工作负载的tpu和gpu，但是目前还不存在局部2D自注意力的快速实现。为了弥补这一差距，我们引入了一个非中心版本(non-centered version)的局部注意力，通过haloing高效地映射到现有的硬件。虽然我们的公式打破了平移等变性，但相比SASA中使用的中心的局部自注意力(centered local self-attention)，它提高了吞吐量和准确度。本文还介绍了带步幅的自注意力下采样(trided self-attentive downsampling)，用于多尺度特征提取。

我们利用这些技术开发了一个新的局部自注意模型家族HaloNet，它在不同的参数范围内实现了最先进的性能。最大的HaloNet在ImageNet[47]分类基准（第4.1节）上达到84.9%的top-1精度。我们进行了一项详细的研究，以揭示自注意和卷积模型的scale如何不同。我们的自注意层在COCO基准上使用MaskR-CNN框架，在诸如目标检测和实例分割（第4.6节）等较难的任务上也显示了有希望的结果。最后，我们讨论了视觉自注意力目前的局限性和未来的工作思路。

2 模型与方法

虽然我们的模型使用自注意而不是卷积来捕捉像素之间的空间交互，但它们采用了现代卷积神经网络（CNNs）的一些重要结构特征。与CNNs一样，我们计算多尺度特征层次结构[34]，从而能够在定位和实例分割等任务中检测多种大小的物体。为此，我们开发了一个带步幅的（编者注：即下采样的）自注意层，这是步幅卷积的自然延伸（第2.2节）。为了解决在大分辨率下全局注意力的计算成本过大的问题，我们遵循相当普遍的局部处理原则，这是卷积和自然感知系统的核心[25，26]，并使用空间范围受限的自注意力形式。然而，与[43]的模型（即SASA，其也使用了局部自注意力）不同的是，我们放弃遵循平移等变性，以更好地利用硬件，这提高了速度-精度的权衡（第2.2节）。还要注意的是，当我们使用局部注意力时，我们每个像素的感受野相当大（高达18×18），我们在第4.2.2节中说明，较大的感受野在较大的图像上有帮助。在本节的剩余部分中，我们将阐述视觉任务的自注意力的动机，并描述如何放松平移等变性，以高效地将局部自注意映射到硬件。

2.1. 自注意可以产生空间变化的卷积滤波器

最近的研究[8]表明，有足够数量的头部和正确的几何偏置的自注意力可以模拟卷积，这表明自注意和卷积之间有更深的关系。自注意被认为是一种直接捕捉远处像素之间关系的方法[43，22，58]。它也被解释为非局部方法经典技术的具体实例[4，59]。我们在本节讨论的观点是，将自注意视为生成空间变化的滤波器，而不是像标准卷积那样在每个空间位置重复使用相同的滤波器[14]。为了观察这一点，我们将自注意和卷积看做是广义的空间维度聚合函数(pooling function)的具体实例。给定一个输入 $x\in R^{H\times W\times c_{in}}$ ，其中 $H$ 表示高度， $W$ 表示宽度， $c_{in}$ 表示输入通道数，我们定义一个二维的聚合函数计算在位置 $(i,j)$ 的输出 $y_{ij} \in R^{c_{out} }$ 为：

$y_{ij} =\sum_{a,b\in N(i,j)}f(i,j,a,b)x_{ab}$

其中 $f(i,j,a,b)$ 是一个函数，它在以 $(i,j)$ 为中心的大小为 $k×k$ 的2D窗口 $N(i,j)$ 中的每个位置返回权重矩阵 $W\in R^{{c_{in} }\times {c_{out}} }$ 。请注意，在本节后面的部分中，我们将介绍非居中窗口的自注意力，但为了便于解释，我们在此处使用居中窗口。对每个像素 $(i,j)$ 重复该计算。

对于卷积， $f(i,j,a,b)$ 为邻域中的每个相对距离返回不同的线性变换，并且这些权重在所有位置 $(i,j)$ 中共享。权重共享显著减少了参数，并鼓励学习在空间上重复出现的特征。在点积相对自注意[48，43，3]中（也就是公式(2)和公式(3)），邻域中的每个像素共享相同的线性变换，该线性变换乘以标量概率，标量概率是内容-内容交互和内容-几何交互的函数，从而导致权重可以在空间上变化。例如，对于图像中两个不同位置的球和橙子，球和橙子区域内部的像素可能会生成不同的 $p_{a-i,b-j}^{ij}$ ，因为它们周围的内容（例如颜色或纹理）不同。

对于自注意力， $W_Q$ 、 $W_K$ 和 $W_V$ 是在所有空间位置共享的习得的线性变换，当用于变换 $x$ 时，它们分别产生查询、键和值（queries, keys, and values）。空间几何由 $r_{a-i,b-j}$ 捕获，它是习得的基于相对位置的嵌入。 ${(W_Qx_{ij})}^TW_Kx_{ab}$ 捕获窗口中查询像素和关键像素之间的内容-内容交互。 ${(W_Qx_{ij})}^Tr_{a-i,b-j}$ 是内容-几何交互，它捕获了查询和键像素的相对位置之间的关系(the relationship between the query and the relative position of the key pixel)[48]。注意，这个公式保留了平移等变性。如果一个物体在图像中平移，对于物体内的任何像素，它周围的内容保持不变，产生相同的 $p_{a-i,b-j}^{ij}$ ，从而在自注意力后产生相同的输出。为了增加表现力，使用了多头注意力[57]，它以不同的参数并行地多次重复这种计算，类似于组卷积[30，61]。

在[43]的SASA模型中，局部窗口 $N(i,j)$ 是一个以 $(i,j)$ 为中心的 $k×k$ 窗口，类似于卷积。这个局部窗口 $k$ 的大小的设置对于使用自注意力是重要的。与密集卷积不同，将 $k$ 增大，模型的参数数量不会显著增大。由于投影 $(W_Q,W_K,W_V)$ 的参数是独立于 $k$ 的，随着 $k$ 增大，只有 $r_{a-i,b-j}$ 的参数量会增多。但是，与投影参数量相比， $r_{a-i,b-j}$ 只占参数量的一小部分，因此增加 $k$ 不会显著影响层的参数数量。（注释1：例如一个窗口大小为63，每个注意力头的维度为16， $r_{a-i,b-j}$ 只会在每层增加63×16=1008的参数量，因为 $r_{a-i,b-j}$ 在不同的注意力头之间是共享的。相反，如果注意力层的维度是512， $(W_Q,W_K,W_V)$ 会贡献786432的参数量。我们在附录中展示细节。)相比之下，卷积层中的参数数量与 $k$ 成二次比例（例如，5×5卷积的参数量是3×3卷积的25/9倍）。另一方面，自注意的计算成本与 $k$ 成二次增长，这阻止了 $k$ 取很大的值。

2.2. 通过放松平移等变性来改善速度-内存的权衡

全局注意力，要计算所有位置与其它所有位置的相关性（all locations attend to each other）由于 $k$ 的二次计算成本，对于大多数图像尺度来说计算量都过高。因此，多尺度视觉主干需要使用局部注意力来限制 $k$ 的大小。我们遵循[43]中发展的局部注意力的直观形式，模仿卷积中使用的正方形邻域。这种形式的局部注意力需要在每个像素周围提取局部2D网格。不幸的是，虽然深度学习库自动处理用于卷积的邻域收集，但并没有用于局部自注意力的邻域收集函数（或任何通用的局部函数）。因此，实现局部自注意力需要在实际的自注意力运算之前显式地收集局部邻域。虽然这个局部邻域收集函数的实现最初看起来可能是一个相对较小的实现细节，但实际上，它必须仔细设计，以减少内存使用，同时避免不必要的额外计算。未优化的实现会因为导致内存不足或速度过慢，阻碍了自注意力模型的扩展。下面的讨论是给这个邻域收集函数的设计设定一个框架。

一种简单的方法是在每个像素周围分别收集 $k×k$ 大小的窗口。如表1（第1行）所总结的，由于要复制领域里 $k^2$ 个像素中的每一个的内容，因此该方法会占用 $k^2$ 因子的内存。此解决方案很快会导致内存不足错误。全局注意力（第4行）是另一个极端，所有像素共享同一个邻域，这样降低内存但也会导致更多的FLOPs（注释2：例如，128×128分辨率、64通道的输入，相比输入输出都是64通道的3×3卷积，全局注意力会导致28倍多的FLOPs）。这个解决方案大大减慢了模型的速度，同时由于注意力矩阵的规模过大，也带来了内存占用的问题。介于这两个极端之间的解决方案应该适当地权衡内存占用和计算量，要认识到要减少浪费。

表1.自注意机制的缩放表现scaling behavior。

f

是头的数目，

b

是块的大小，

c

是通道总数，

h

是halo的带宽的大小

一种折衷的解决方案可以通过以下想法来实现：相邻的像素，共享它们的大部分相邻像素(neighboring pixels share most of their neighborhood)（编者注：如果我们俩是邻居，那么我们的邻居的大部分是共同的）。例如，两个紧挨着的相邻像素（即两个同邻域内的像素），共享 $k\times (k-1)$ 个相邻像素。因此，像素块(block of pixels)的局部邻域可以一起提取一次，而不是每个像素提取单独的邻域。FLOPs可以通过改变构成块的像素数来控制。我们把这种策略命名为分块局部自注意力(blocked local self-attention)。上面讨论的两个极端是分块局部自注意力的特例。全局注意力是将块大小设置为整个空间范围，而逐像素提取邻域是将块大小设置为1。

图1描述了对高度为4、宽度为4、通道数为 $c$ 的图像执行步幅为1的成块局部自注意力所涉及的不同步骤。分块将图像划分成一 $\frac{H}{b} ,\frac{W}{b}$ 个非重叠的、尺寸为 $(b,b)$ 的块。每个块表现为一组查询像素，haloing操作将它们周围的一圈像素（带宽为 $h$ ）的像素组合进来（在图像边界处使用padding），以获得共享领域块，形状为 $(\frac{H}{b} ,\frac{W}{b},b+2h,b+2h，c)$ ，从中计算键和值(keys and values)。然后，对每个查询块及其对应的邻域进行并行的 $\frac{H}{b} \times \frac{W}{b}$ 个注意力运算，如图1中的不同颜色所示。SASA[43]使用了相同的分块策略3，将带宽设置为 $h=\lfloor \frac{k}{2} \rfloor$ ，并使用注意力掩膜，来模拟大小为 $k×k$ 的以像素为中心的邻域窗口。例如，以像素为中心的窗口大小为7×7，[43]设置 $h=3$ 。注意力掩膜的使用使操作具有平移等变性，因为每个像素只看它周围的一个正方形窗口。

图1.HaloNet局部自注意力结构：一个[4,4,c]图像的分块局部自注意力的不同阶段，块大小b=2，带宽h=1。首先将图像分块为非重叠的、形状为[2,2,c]的块，用于计算queries。haloing步骤提取围绕着块，提取[4,4,c]的区域，其线性变换为keys和values。注意力计算后的空间维度与queries是相同的。

然而，使用注意力掩膜的缺点是，无论算法如何实现，它都要浪费一些计算量。如果不使用注意力掩膜，感受野会增加，却不导致额外的计算量，如表1（第2行和第3行）所示。然而，由于非方形感受野意味着像素的输出取决于它落在哪个块中，因此像素级的平移等变性丢失。例如，在一个块的左边缘的一个像素，它可以看到位于其方形感受野右侧的其他像素。如果将整个图像向右移动一个像素，则该像素现在落入相邻块的右边缘，并且现在可以看到位于其方形感受野左侧的额外像素。因此，像素的输出取决于其在块中的位置，如果图像移动，该输出可能改变。另一个观点是，分块局部自我注意仅对平移量为 $b$ 是平移等变的。虽然像素级平移等效性被认为对实现良好性能很重要[65]，但我们发现，根据经验，使用无掩膜的分块局部自注意力实际上提高了模型的准确性（见第4.3节）。我们怀疑常见的数据增强策略中的图像偏移和裁剪扰动会减少对这种归纳偏置的依赖。因此，我们采用无掩膜的分块局部自注意力，因为它在不牺牲性能的情况下提高了准确性。

与SASA的另一个不同之处是我们实现了下采样。SASA中的下采样层是一个注意力层接一个池化层，我们的实现是用一个单独的步幅注意力层，其对queries下采样，类似于步幅卷积，如图2所示。请注意，步幅为2和步幅为1提取的相同的邻域（如图1）。这种变换不会影响精度，但在下采样层可以将FLOPs缩小4倍。我们还实现了一些重要的算法优化，主要通过避免reshapes 和data formatting操作来提高吞吐量。为了行文简洁，我们在附录D中列出了它们。综上所述，这些改进所产生的加速效果非常显著，如图3所示，step time提速2倍。这些改进可以用来训练以前过于昂贵的大型自注意力模型。我们将更多的优化留到以后的工作中，比如融合操作和更好的内存访问与计算的流水线。

图2。注意力下采样层对查询进行下采样，邻域与步幅为1的情况相同

图3.优化提高了性能。这里的改进是通过我们的注意力下采样和改进的局部自注意算法来减少FLOPs的结果，这些算法避免了reshape和数据格式化。在某些情况下，我们将TPUv3上计算的训练步长时间减半。

在本节结束时，需要注意的是，在多尺度体系结构的更深层中，更小的空间维度和更大的通道将使计算变成全局注意力。我们在第4节中介绍的模型也利用了这一点，通常在较高分辨率时使用局部注意，在图像分辨率最小时使用全局注意。

2.3. HaloNet

利用上面详细介绍的带haloing的局部二维自注意力的实现，我们提出了一个新的模型HaloNet，它在参数精度折衷曲线上匹配最新的卷积模型。我们利用ResNet[17]的结构，将多个残差瓶颈块堆叠在一起（见表2）。HaloNet使用了ResNet的一些小修改：（a）根据EfficientNet[55]，在较大模型的全局平均池化之前添加一个最终的1×1的卷积层，（b）修改瓶颈块宽度因子，传统上固定为4，（c）修改空间运算的输出宽度因子，传统上固定为1，（d）由于计算原因，将第三阶段的块数从4更改为3，因为在高分辨率层中，注意力更昂贵。我们还将四个阶段中每个阶段的注意力头数固定为（4，8，8，8），因为在更高的分辨率下注意力头数越多计算量越大。总之，HaloNet中的缩放维度是：图像大小 $s$ 、查询块大小 $b$ 、halo大小 $h$ 、注意力输出宽度因子 $r_v$ 、瓶颈输出宽度因子 $r_b$ 、第三组中的瓶颈块数 $l_3$ 和最终的1×1卷积宽度 $d_f$ 。我们的注意力范围从14×14（b=8；h=3）到18×18（b=14；h=2）。

表2 HaloNet模型家族的一个实例

由于ResNet结构最初是为卷积而设计的，因此我们怀疑设计专门用于注意力的架构可能会改进HaloNet。在我们的工作中，我们保持了模型所有层超参数的均匀性，如块（ $b$ ）和带宽（ $h$ ）的大小。我们还希望，使用自动架构搜索方法[55]对特定加速器来优化这些超参数将导致更好的局部注意力架构。在我们的工作中，我们使用可比较的图像大小，例如EfficientNet，来确定注意力模型是否可缩放到更大图像上。

3. 相关工作

近年来，自注意力在视觉模型上的应用稳步上升。首先是在各种形式的序列建模方法[15，1，57]中被提出，在图像描述模型的文本生成模块中用来关注图像特征[62]。注意力也与非局部均值密切相关[4]，非局部均值是最初用于图像去噪的一种像素成对加权、全局求和的方法。[59]在空间下采样卷积特征的基础上应用非局部方法来改进视频分类。然而，由于这些方法与感受野大小成二次比例，因此不能使用，因为空间大小太大。为了应用自注意，[40]在图像上应用局部注意力来生成图像。[3] 在空间上对注意力特征进行降采样，并将注意力输出连接到卷积特征。相反，我们直接建立在[43]方法的基础上，后者计算局部区域的注意，以便建立一个完全自注意的视觉模型，用于分类和目标检测。对于纯自注意力的视觉模型，也提出了不同的注意形式[22，66]，这是正交的，是对本研究中缩放注意力的补充。除了我们关注的空间维度外，通道注意力也被用于增强卷积模型[23，33]。在最近和同时进行的工作中，Vision Transformer[12]表明，在非常大（JFT-300M[52]）和中型（ImageNet-21k[11]）分类数据集上预训练时，将Transformer应用于非重叠图像块的投影可以获得与SOTA相当的精度。然而，他们的模型并没有采用多尺度体系结构，我们在这项工作中的重点是从头开始对ImageNet进行训练[47]。在第4.5节中，我们进行了迁移实验，并与ViT和BiT[29]进行了比较。

一般来说，由于OP的算法变化和更好的软件实现，性能往往会随着时间的推移而提高。特别是，在过去的十年中，卷积通过以下方面的变化得到了改进：（a）OP的计算[5，28，37，56，60，31]；（b）软件实现[6]；（c）OP本身的结构，例如，通过分组卷积[61]和深度可分离卷积[49]。注意力正处于这种性能改进轨迹的开始阶段，鉴于其在序列建模中的重要性[57]，很可能会看到持续的努力来提高性能。如果更广泛地采用局部注意力来解决处理大量输入的一般问题，也可以提高性能。我们的工作引入了分块局部注意力来高效地处理近邻。其他形式的非全局像素交互也可以有效地实现[7，21，58，2]。

4 实验

每个HaloNet模型（H0–H7）都是通过依次增加表2中定义的超参数值来设计的。为了节省空间，我们把模型的精确配置留给附录C.2。我们还将较大的HaloNet模型的训练和评估与较大的EfficientNet模型相比较留待将来的工作。

4.1. 与最先进的卷积模型相比，HaloNet具有竞争力

我们在ImageNet[47]（ILSVRC-2012）基准上训练我们的HaloNet模型，批量大小为4096，学习率为1.6，线性预热10个epoch，然后进行余弦衰减[36]。用Nesterov的加速梯度[38,53]对模型进行350个epoch的训练，并用dropout[51]、权重衰减、随机增强[10]和随机深度[24]进行正则化。

我们发现，在相同的参数下，与EfficientNet模型相比，HaloNet的性能相当或稍好一些（图4），优于其他模型族。我们的最佳模型H7在ImageNet上实现了84.9%的top-1验证集准确度，以及在ImageNet V2[45]有74.7%的top-1准确度（与[45]中的线性拟合有-0.5%的差距）。对于我们的每一个HaloNet模型，我们使用与相应的EfficientNet模型相当的图像大小来训练，图像大小最大达到600×600（表A2）。要比较我们与EfficientNet的延迟，读者可以参考第5节。据我们所知，这些结果首次表明，当从头开始在Imagenet上训练时，基于自注意的视觉模型与SOTA的图像分类性能相当。请注意，对于我们的所有实验，我们报告的是训练结束时的准确度，以及，我们调试了正则化参数，例如基线和HaloNet模型的增强超参数。

图4。在精确度与参数的权衡上，HaloNets 可以与EfficientNets相匹配。EfficientNets B5 和 B7的准确度使用了RandAugment

4.2. 模型研究1：比较自注意和卷积

在下面的部分中，我们将着重于模型研究，以区分自注意力相对于视觉卷积的优势，并了解如何最佳地设计自注意力视觉体系结构。这方面的知识非常重要，因为卷积网络的大部分进展来自于架构设计的改进，同时保持核心卷积运算不变[30，54，17]。我们相信，我们的研究是第一次明确检查最佳自注意力视觉架构的设计。

在实验部分的剩余部分，我们将与ResNet-50[18]进行比较，ResNet-50是一种典型的视觉模型，因为我们烧蚀的许多组件已经为ResNet-50进行了很好的研究，使我们能够使用基线模型的最佳实践。我们调整了基线ResNet-50实现，与文献中通常报告的数字相比，达到了更好的准确率77.6%。例如，[17]报告为76.3%。然后我们创建了一个新的HaloNet架构HaloNet-50，它通过用局部自注意力代替空间卷积来精确匹配ResNet-50架构。HaloNet-50和ResNet-50的参数分别约为18million和25.5million。我们在256×256的图像大小上训练两者150个epoch。我们在附录中分享了消融设置的其他培训细节。

4.2.1 将卷积转移至自注意力

在核心OP之外利用正则化和体系结构模块是获得强大结果的关键[19]。在本节中，我们将研究这些附加成分对自注意模型的影响。我们研究的组件都是为卷积模型设计的，因为它们是通过实验（人工或自动搜索）在卷积模型上开发的。我们检验这些成分是否能成功地转移到新的自注意网络模型家族中。

基于EfficientNet[55]的设计，我们关注4个不同的组件，2个架构模块和2个正则化：Squeeze-and-Excitation（SE）[23]，其是在空间卷积后使用通道注意力；SiLU/Swish-1[44，13，20]，形式为x·sigmoid（x）的激活函数；RandAugment（RA）[10]，一种简化自动增强的数据增强方案[9]；以及标签平滑（LS）[54]，一种标签分布的平滑。

将这些不同成分添加到基线的结果如表3所示。令人惊讶的是，相同强度的正则化对HaloNet带来的提升要比ResNet显著多得多，尽管HaloNet的参数比ResNet少30%左右。添加标签平滑和随机增强后，HaloNet提高了1.3%，ResNet提高了0.8%。这一结果表明，自注意力模型可能需要正则化，这是较大卷积模型的典型，这可能是由于自注意力模型表示能力更强。

然而，为卷积模型开发的体系结构模块只对注意力模型进行了少量改进。当添加SE块和SiLU/Swish-1时，ResNet提高了1.3%，而HaloNet仅提高了0.4%。我们推测，HaloNet模型受益于门控和乘法相互作用（包括自注意），不需要显式门控，如SE。为了发现能够持续改进各种自注意模型的架构模块，必须进行进一步的研究。受这些发现的启发，我们决定在HaloNet H0-h7模型中使用标签平滑、SiLU/Swish-1和Rand-Augment。我们还将随机深度用于更大的模型[24，55]。

表3。正则化对HaloNet帮助更多。而对ResNet有很多提升的体系结构的改进，却没有给HaloNet带来显著提升。从基线模型开始，添加标签平滑（LS）、随机增强（RA）、压缩和激励（SE）和SiLU/Swish-1（SiLU/Sw1）。

4.2.2 增大图像尺寸可提高精确度

自注意的一个有益特性是，注意力感受野的大小可以随着图像的大小而变化，而不会显著影响参数的数量（见第2.1节）。如图6所示，当使用较大的图像时，HaloNet持续改进。虽然我们也看到了卷积模型的改进，但HaloNet和ResNet的准确度的差距仍然是存在的。

图6.随着图像尺寸的增加，HaloNet-50和ResNet-50之间的精度差距仍然存在。HaloNet上标示的是块大小(b)和带宽(h)。

4.3 模型研究2：HaloNet结构的研究

在本节中，我们将研究放松平移等变的影响以及邻域窗口和带宽大小的关系。为行文简洁，我们在附录B中详细研究了我们模型的各个组成部分，如 $r_v$ 和 $r_{qk}$ 等。

放松平移等变性： 在图5中，我们看到HaloNet-50（b=8，h=3）实现了相比[43]更好的精度，两者使用相同的block和带宽，[43]的感受野是7×7并且具有注意力掩膜。并且随着更多的正则化，差距变大。这表明更大的感受野比诸如平移等变性的归纳偏置更重要。

图5 放松平移等变性提高准确度

窗口和带宽大小：

当使用分块的输入时，有两种更改注意力窗口大小的方法：更改查询块大小或带宽大小。对于相同的窗口大小 $w$ ，较小的查询块（编者注：分得细）和较大的带宽会占用更多的内存，如第2.2节所述。

我们在图7中看到，随着窗口大小的增加，精确度不断提高。特别是，将窗口大小从6×6加倍到12×12，可获得1.3%的精度增益。这些结果表明，在不增加参数数量的情况下，增加窗口大小可以成功地用于缩放模型，这可能对生产环境有利。此外，对于固定的窗口大小，查询块大小的选择不会影响结果，因此可以使用较大的查询块大小来减少内存。图7还显示，带宽为0的话，会显著降低准确率，除非块非常大。例如，使用4的块大小和1的带宽比使用8的块大小和0的带宽精度更好，尽管邻域大小更小。

图7.增加窗口大小在一定程度上提高了精确度。标示是块大小(b)和带宽(h)。h=0表示对非重叠块的注意力计算

4.4. 卷积注意混合模型提高了速度精度的权衡

在我们的最后一组烧蚀中，我们用卷积来代替自注意，以了解注意层目前最有益的地方。在表4中，我们展示了在我们表现最好的模型（HaloNet H7）的每个阶段，用带SE块的卷积代替注意力层的结果。最后一个阶段使用注意力，其余所有阶段都用卷积，是最快的模型，但是top-1准确度会有1%的下降。在第1-2阶段用卷积，第3-4阶段用注意力，在显著改善训练与速度的同时将准确度性能下降的影响降至最小。我们将混合模型的详细研究留给将来的工作。

表4。在第1阶段和第2阶段用卷积代替注意层显示出最佳的速度和准确性折衷。所有的模型都有大约67百万个参数，训练和推理时间是与EfficientNet B7之比。请参见图8，了解step time与其他卤网模型的比较

4.5. 从ImageNet-21k迁移

到目前为止，我们的实验主要集中在ImageNet-ILSVRC-2012[47]上从头开始的训练，其中正则化和更长时间的训练对于良好的精确度至关重要。文献[12，29]已经表明，在较大的标记数据集（如ImageNet-21k[11]或JFT-300M[52]）上预训练模型后，短时间的微调步骤可以获得更好的精度，而无需正则化。为了了解HaloNet模型的迁移特性，我们将HaloNet-H4的基本宽度增加到128，并对[29]中的迁移协议进行评估，在公共ImageNet-21k数据集上进行预训练，在ImageNet上进行微调。根据我们在表4中的观察，我们训练了这个模型的混合版本，也就是在前两个阶段用卷积。请注意，我们的混合版本可以看做是用一系列卷积对图像做下采样。由于我们需要与[29]进行公平的比较，所以在卷积阶段我们不使用SE块[23]。我们还研究了类似于视觉变换器的4×4非重叠patches的线性投影替换卷积干的效果。模型的细节见附录E。ImageNet-21k包含1420万个带注释的图像和21k个标签，两者都比ImageNet大一个数量级。和[29]一样，我们进行了90个阶段的预训练，批量大小为4096，基本学习率为0.16，线性预热2个epoch，然后进行余弦衰减[36]。我们还使用0.00008的权重衰减，并在预训练和微调期间使用Nesterov的加速梯度[38，53]进行训练。我们对256×256大小的图像进行预训练，并对不同大小的图像进行微调，如表6所示。我们更宽的H4和混合-H4模型比Vision Transformer和[29]中的4×wide ResNet-152实现了更好的精度，并且在较大图像上的推理速度也更快。我们在ImageNet上进行了8个阶段的微调，除了标签嵌入矩阵初始化为零外，其他参数都是通过预训练学习的。我们的训练批量为512，学习率为0.016，线性升温0.5个epoch后余弦衰减。尽管在更大的数据集上进行了预训练，但在微调过程中，我们还是受益于标签平滑度为0.1的微调。在微调过程中，我们不使用Polyak平均[41]和其他正则化。

我们在迁移学习上的初步结果是有希望的，因为我们在这个数据集上获得了比其他模型更好的参数-精度和速度-精度折衷。我们将更大模型的HaloNet与混合HaloNet的迁移学习的研究留给未来工作。我们的模型在大图像上的速度优势使得它们适合于具有挑战性的大图像结构化预测任务，如目标检测和实例分割，我们将在下一节简要探讨。

表6 当在ImageNet上进行微调时，在ImagetNet-21k上预训练的HaloNet模型表现良好。对于HaloNet和ViT，我们对384×384和512×512大小的图像进行了微调。pretraining step time报告每个核,批量大小为32的TPUv3计算时间。推理速度也在单个TPUv3核上计算

4.6 检测与实例分割

为了了解我们的OP是否可以推广到更大图像上的结构化预测任务，我们使用MaskR-CNN[16]框架对简单的注意力-卷积混合模型检测和实例分割进行了初步研究。与纯注意力模型相比，这些混合模型的速度更快，消耗的内存更少，从而实现了更快的实验周期。我们仅将ResNet-50和ResNet-101主干中的最后3个卷积层替换为两个b=8和h=3的halo层（表5中的第3行和第6行）。对于ResNet-50，我们还研究了使用b=32，h=3来了解更大感受野的好处。我们还使用带SE块的卷积，并使用第4.2.1节提到的正则化手段：标签平滑、随机增强和随机深度，对512×512图像进行预训练。我们在COCO数据集[35]上使用Cloud-TPU-Detection-Codebase （注释4：https://github.com/tensorflow/tpu/tree/master/models/official/detection）训练我们的模型，在该数据集训练32个epoch，训练尺寸是1024×1024。我们在附录C.4中提供了更多的训练细节。

表5第2行中的ResNet-50基线显著优于文献中通常报告的结果（第1行）。相比我们的更强基线，我们的注意力变体在边界框检测上至少获得了0.7的mAP增益，在实例分割上至少获得了0.6的mAP增益（在表5的第3、4和6行中用++表示）。块大小b=8的局部注意力模型带来的提升，在检测中缩小了R50和R101基线之间的差距的接近一半，在实例分割中缩小了70%，尽管在系统时间上小于差距的三分之一（despite being less than a third of the gap in terms of wall-clock time）。b=8和h=3的局部注意力在R101主干上也有所提高。有趣的是，当注意被使用时，大物体的定位 ${AP}_{l}^\ast$ 显示出最大的改进。较大的块大小（第4行中的b=32）可以实现与b=8非常接近的性能，但速度较慢。然而，我们看到在小物体（ ${AP}_{s}^\ast$ ）上，b=32比b=8做得好得多。未来的工作可以结合这两种设置的最佳。注意，在b=32的情况下，最后两个注意层进行全局注意，因为图像在每个空间维度下采样到1024/32=32像素。与本文同期的工作，BoTNet[50]在ResNet结构的卷积-注意力混合模型中使用全局注意力以执行结构化预测任务和分类。有关定位任务的全局注意力效果(efficacy)的更多详细信息，请参见[50]。（[50]：Aravind Srinivas, Tsung-Yi Lin, Niki Parmar, Jonathon Shlens, Pieter Abbeel, and Ashish Vaswani. Bottleneck transformers for visual recognition. arXiv:2101.11605）

这些模型只有三层的自注意力，更多层的自注意力的结果可能会不一样。我们将纯注意模型的检测和实例分割的研究留给以后的工作。

表5。目标检测和实例分割的准确性。在最后一个阶段，我们尝试了两种自注意力的设置：在ResNet50上b=8、h=3，以及b=32、h=3。bb（bounding box）表示检测，mk（mask）表示分割。s、m和l分别表示小物体、中物体和大物体。速度的测量是，在2个TPUv3核上，batchsize为32，只测主干（不包括FPN），计算milliseconds。训练时间是指在批处理大小为64的8个TPUv3核上运行的Mask RCNN训练的图像/秒的峰值来计算出的总训练时间

5 讨论

在这项工作中，我们建立了多尺度的自我注意模型，与最好的卷积模型相竞争。为了达到这个结果，我们开发了两个注意改进：分块局部注意了和注意力下采样。我们还进行了多次烧蚀，以了解如何提高自注意模型的缩放。

我们的研究结果表明，当在ImageNet上从头训练时时，完全的自注意力模型在准确度方面是有竞争力的。图8显示，基于HaloNet的完全注意力模型要比对应的EfficientNet要慢一些，需要进一步优化以进行大批量训练。然而，我们的卷积-注意力的混合模型与EfficientNets具有相同的速度精度折衷。在从ImageNet-21k迁移时，我们的模型在精度和速度上都优于BiT[29]和ViT[12]等非常强大的模型。对于机器学习研究人员来说，模型优化，例如使用架构搜索方法来寻找更好的速度-精度折衷，或者搜索更快或更强的注意力不同形式[66，46]，是一个很有前途的方向。模型的优化，例如更好的内存管理，可以提高这些模型的实用性。此外，将模型扩展到更大的宽度，我们的操作实现了从内存占用限制转移为计算限制，并导致了更好的速度-精度折衷（Also, scaling up our models to larger widths might cause our operations to transition from being memory bound to compute bound, and lead to better speed-accuracy tradeoffs）。我们把这项研究留给将来的工作。总的来说，我们的研究表明，在传统上由计算机视觉主导的体制下，自注意具有竞争性。未来的工作可以在规模和效率方面进一步推动这些界限。

图8。纯注意力的HaloNet要比EfficientNet训练得慢。时间是在每个TPUv3核上处理32的batchsize的计算时间。带注释C1、C12和C123的绿色点分别对应于在第1阶段、第1到第2阶段、第1到第3阶段使用卷积的混合模型。（见表4）

6 致谢

We would like to thank David Fleet for valuable discus-sions. We would also like to thank Irwan Bello, Barret Zoph, Mingxing Tan, and Lucas Beyer for valuable commentary on earlier drafts of the paper.

附录

A 相对嵌入只会给模型增加很少的参数

随感受野的增长，我们的模型的参数量增长非常缓慢。在这一节中，我们将展示相对嵌入中的参数数量（唯一的空间相关参数）是非常小的。如文中所述，位置 $(i,j)$ 处的二维局部自注意力计算为：

其中queries $q_{ij}=W_Qx_{ij}$ 、keys $k_{ab}=W_Kx_{ab}$ 和values $v_{ab}=W_V x_{ab}$ 是像素的线性变换，且 $r_{a-i,b-j}$ 是学到的基于嵌入的相对位置编码。和Transformer[57]一样，我们还使用了多头注意，在这里我们使用不同的参数并行运行多个自注意力。然而，每个头共享相对嵌入 $r_{a-i,b-j}$ 的参数。对于每个像素周围大小为 $k$ 的注意力窗口，我们将沿高度和宽度方向的相对嵌入分解[43]，并将一个注意力头的通道数，一半分配给高度，一半分配给宽度。如文中所述，如果每个头部的维度固定为16，这样每个注意力层的 $r_{a-i,b-j}$ 的参数量就是 $2(k-1)\times 16$ 。相反，如果一个注意力层的通道数为 $d$ ，则3个线性变换的每一个都具有 $d^2$ 的参数量。因此，相对嵌入的参数量与线性投影的参数量之比为 $\frac{2(k-1)\times 16}{d^2}$ 。对于k和d的典型值来说，它很小

B 扩大自注意模型的研究

在第4.3节中，我们介绍了模型的一些缩放特性。在表A1中，我们试图了解模型的哪些其他部分对精度的影响最大。在我们的研究中，我们通过缩放不同的超参数来增加HaloNet-50的尺寸，以达到3000万的参数预算。我们发现，通过增加 $r_v$ 和增加更多的层来增加注意力计算量是提高精度最有效的。

表A.1 增加值项的通道数和层数对精度的影响最大。

C 实验细节，超参数

这一节中，我们列出了实验细节和模型配置，这些都是为了行文简洁而在正文里省略的

C.1 模型研究的实验细节

在第4.2节和第4.3节中，所有HaloNet-50模型使用与标准ResNet-50[17]模型相同的层分配和通道宽度。ResNet-50和HaloNet-50模型都经过了150个epochs，训练尺寸为256×256，学习率为0.1。在使用RandAugment的实验中，我们使用0:00004的权重衰减，这是RandAugment[10]的设置，0.00008作为其他设置。使用RandAugment并使用0.00008的权重衰减似乎对ResNet-50的精确度有负面影响。在这些部分中，RandAugment的幅度(magnitude)为10。对于HaloNet-50，我们使用了块大小b=8，h=3。我们把每个头的通道数固定为16。对于SASA模型，我们按照[43]使用了一个以像素为中心的7×7窗口。

C.2 HaloNet模型

在表A2中，我们描述了我们的HaloNet模型H1到H7的配置。HaloNet族中的超参数是：图像大小 $s$ 、查询块大小 $b$ 、halo大小 $h$ 、注意力量输出通道数(width)因子 $r_v$ 、瓶颈输出通道数因子 $r_b$ 、第三阶段(the third group )中的瓶颈块数 $l_{3}$ 和最终1× 1卷积通道数(width) $d _f$ 。我们的每一个HaloNet模型都在与相应的EfficientNet[55]模型相当的图像大小上进行训练，如表A2所示。

表A2.HaloNet模型的配置，每个模型在参数方面与EfficientNet系列的模型相匹配。四个阶段的头数为4,8,8,8.

C.3 分类超参数

在本节中，我们将详细介绍我们的训练和正则化设置。我们使用了权重衰减2e5和学习率为0.1的余弦衰减方案[36]。最大的模型总是在训练结束时过拟合，我们将其归因于训练结束时的学习率变为0[64]。为了解决这个问题，我们将余弦退火的结束时间设置为原始学习速率的1.0/128，而不是0。对于RandAugment[10]，最小的H0到最大的H7的RandAugment幅度为6，8，10，14，17，19，24和31。请注意，我们还没有对RandAugment幅度做广泛的调试。

C.4 检测与实例分割的超参数

我们使用Mask-RCNN[16]进行所有的检测和实例分割实验。我们在ImageNet上预训练主干，超参数与C.3节基本相同。主干使用512的图像大小进行350个epoch的预训练，选择的图像大小更接近检测设置中使用的1024个图像大小。模型用RandAug正则化，随机深度为15，概率为0:1，压缩激励为1/8。检测代码和超参数直接采用开源的TPU检测与分割框架6。在检测/实例分割阶段，用预先训练好的权值初始化主干，其他参数则从头开始初始化。该模型训练67500步，学习率为0:1，在60000和65000步衰减，在SGD中使用0:1的学习率，动量为0:9，预热500步，固定学习率为2300，批量大小为64，分布在32TPUv3核上，1024× 1024图像大小，L2权重衰减为4e5，多尺度抖动，大小介于[4/5，5/4]之间。

D 优化

我们尽可能避免数据格式化操作( formatting operations)，这会减慢模型的速度。优化有两个关键的地方：

恒定的分块Persistent blocking： 一旦图像被分块，我们就将大小为 $(b,b)$ 的块展平成长度为 $b^2$ 的序列。在网络结束之前，我们都不会将其reshape回4D的数据格式，实现batch normalization[27]等操作来处理被分块的格式。这样，图像是以5D的格式被处理，即 $(batch,\frac{H}{b} ,\frac{W}{b} ,b^2 ,c)$ ，而不是 $(batch,H,W,c)$ 。

Gathers with convolutions：2.2节的haloing操作也是在5D的数据格式里被执行的，因此领域是展平的。为了提高速度，我们用3D卷积（used as gathering operations）来代替切片和拼接( slices and concatenations.)来实现haloing。

E ImageNet-21k模型

对于我们的ImageNet-21k转移实验（表6），我们对我们的HaloNet H4模型进行了4项更改（H4模型的规格见表A2）。为了增加模型主体中的参数数量，我们将底部宽度rw增加到2.0（使底部宽度为128，是正常宽度的两倍），并且我们还将b从3:0更改为默认的4:0。我们去掉最后一个多余的1× 1卷积，使标签嵌入有大量的滤波核，以说明标签数量较大。最后，我们将第二阶段的层数从3层增加到4层。对于混合模型，我们在前两个阶段使用卷积。用256× 256的图片分辨率预训练，我们设置b=8和h=2。如果用384×384图像分辨率微调，我们使用b=12，h=2。如果用512×512图像分辨率微调，我们使用b=16，h=1。在迁移预训练模型时，我们用预训练最后一步的pretrained checkpoint的所有参数来初始化模型，除了标签嵌入，其初始化为0，也除了相对嵌入，其初始化为预训练的相对嵌入的线性插值。

F 相对位置编码的影响

[43]表明，使用相对位置对于获得良好的精度非常重要。我们发现HaloNet也有同样的结果。使用绝对因子分解的绝对位置编码，在每一层的局部自注意之前添加到激活中，精确度从78.6%（表3第一行）下降到77.5%

HaloNets

推荐阅读更多精彩内容