【读点论文】Learned Queries for Efficient Local Attention（QnA）Transformer新理解，学习好Query,更好的结合局部理解，提升语义理解的效率

Learned Queries for Efficient Local Attention

Abstract

• (Vision transformer, ViT)是一种强大的视觉模型。与前几年主导视觉研究的卷积神经网络不同，ViT具有捕获数据中的长期依赖关系的能力。尽管如此，任何transformer体系结构的一个组成部分，即自注意力机制，都存在高延迟和低效的内存利用问题，这使得它不太适合高分辨率输入图像。为了克服这些缺点，分层视觉模型局部地在非交错窗口上使用注意力注。这种松弛将复杂度降低到输入大小的线性;然而，它限制了跨窗口的交互，损害了模型的性能。

• 在本文中，提出了一个新的移位不变的局部注意层，称query和attend(QnA)，它以重叠的方式在局部聚集输入，很像卷积。QnA背后的关键思想是引入学习query，这允许快速有效地实现。本文通过将该层合并到一个分层视觉Transformer模型中来验证该层的有效性。本文展示了速度和内存复杂性的改进，同时实现了与最先进的模型相当的准确性。

• 最后，本文的层可以很好地扩展窗口大小，需要的内存比现有方法少10倍，而速度比现有方法快5倍。该代码可在 moabarar/qna: [CVPR2022 - Oral] Official Jax Implementation of Learned Queries for Efficient Local Attention (github.com) 上获得。

• 论文地址：[2112.11435] Learned Queries for Efficient Local Attention (arxiv.org)

• 一种具有线性复杂性且平移不变的局部注意力层，比卷积具有更强的表征能力，实验展示了速度和内存复杂性的改进，优于Focal、Swin等。由特拉维夫大学, 提赛德大学发表在 CVPR 2022 Oral。

• 分层视觉模型 —— Query and Attention （QnA）的核心是一种具有线性复杂性且平移不变的局部注意力层。其结构类似卷积，但是比卷积具有更强的表征能力。实验证明，引入学习query的视觉模型不仅性能媲美其他 SoTA 模型，还显著地改进了内存复杂度，加速了模型推理。

Introduction

• 在考虑图像处理的数据聚合机制时，有两个关键因素。卷积是最直接的选择。它们提供了局部性(这是图像处理的先决条件)和效率。然而，卷积捕获局部模式，将它们扩展到全局上下文即使不是不切实际，也是很困难的。另一方面，基于注意力的模型提供了一种自适应聚合机制，其中聚合方案本身依赖于输入，或空间动态。这些模型[Language models are few-shot learners,BERT]是自然语言处理领域事实上的选择，最近也在视觉任务中蓬勃发展。

• (ViT)的早期变体通过处理非交错图像补丁作为单词标记来提供全局上下文。为了使这些模型有效，它们通常需要大量的数据，大量的正则化或修改优化目标。更重要的是，观察到大规模训练驱动模型注意力局部，特别是对于早期的层，这鼓励了局部性是一个强先验的概念。

• 局部注意机制是目前选择的更好的视觉骨干的方法。这些主干遵循类似于卷积神经网络(cnn)的金字塔结构，并通过将自注意力限制在较小的窗口来处理高分辨率输入，最好是有一些重叠或其他形式的相互通信。后一种方法自然地诱导局部性，同时受益于空间动态聚集。另一方面，这些体系结构是以计算开销为代价的，更重要的是，它们不是位移等变的。

• 在本文中，重新审视了局部注意力的设计，并引入了一个新的聚合层，称为query和attend(QnA)。关键思想是利用卷积的局部性和移位不变性以及注意机制的表达能力。

• 在局部自注意中，计算组成窗口的所有元素之间的注意分数。这是一个代价高昂的操作，复杂度为窗口大小的二次。建议使用学习query来计算聚合权值，无论选择的窗口大小如何，都允许线性内存复杂度。我们的层也是灵活的，表明它可以作为一个有效的上采样或下采样操作。进一步观察到，组合不同的query允许以最小的计算开销捕获更丰富的特征子空间。得出的结论是，QnA层与普通transformer块交织形成了一个分层vit家族，与SOTA模型相比，它实现了相当或更好的精度，同时受益于高达2倍的高吞吐量、更少的参数和浮点操作(参见下图)。

  • 

  • $224^2$输入大小下的性能效率比较。与最先进的基线相比，QnA-ViT展示了更好的准确性和效率权衡。本文报告了ImageNet-1k在参数计数(左)、浮点运算(中)和推理吞吐量(右)方面的Top-1精度(y轴)权衡。吞吐量是使用timm库来测量的，在16GB内存的NVIDIA V100上进行测试。

• 通过严格的实验，我们证明了本文的新聚合层具有以下优点:

  • QnA施加局部性，在不影响准确性的情况下提高效率。

  • QnA可以作为通用层。例如，跨行QnA允许有效的下采样，而多个query可以用于有效的上采样，这表明了替代基线的改进。

  • QnA自然地将局部性合并到现有的基于transformer的框架中。例如，展示了在基于注意力的对象检测框架[DETR]中，用QnA层替换自注意层如何有助于提高精度，特别是对于小尺寸对象。

• 视觉Transformer能够捕获数据中的长期依赖关系。尽管如此，任何 Transformer 架构的一个组成部分，即自注意力机制，都存在高延迟和内存利用率低的问题，使其不太适合高分辨率输入图像。为了减轻这些缺点，分层视觉模型在非交错窗口上局部使用自注意力。这种松弛降低了输入大小的线性复杂性；然而，它限制了跨窗口交互，损害了模型性能。提出了一个新的平移不变局部注意力层，称为query and attend (QnA)，它以重叠的方式在局部聚合输入，很像卷积。

Related Work

• 卷积网络:基于cnn的网络已经统治了计算机视觉世界。几年来，计算机视觉社区通过设计强大的架构进行了实质性的改进。一个特别相关的基于cnn的工作是RedNet，它引入了一个对合操作。该操作通过线性投影提取每个像素的卷积核，实现自适应卷积操作。尽管具有自适应特性，但RedNet使用的线性投影缺乏自注意机制的表达能力。

• ViT:自注意的自适应在图像识别、图像生成、目标检测和语义分割等各种视觉任务中显示出良好的效果。然而，这些模型并没有将纯粹的自我注意作为视觉模型的主要工具。相比之下，视觉transformer带来了概念上的转变。这些模型最初是为图像分类而设计的，在标记化的图像补丁上使用全局自注意力，其中每个令牌注意力注所有其他令牌。T2T-ViT通过早期层的轻量级注意力进一步改进了标记化过程。同样，仔细设计基于卷积的stem块可以提高收敛速度和准确性。CrossViT建议在粗粒度和细粒度的补丁级别进行处理。另一方面，TNT-ViT将粗斑块分成局部参与的部分。这些信息随后被融合到补丁之间注意力局关注中。ConViT通过仔细注意力化自关注块以鼓励局部性来提高性能。LeViT通过精心设计提供了一种高效的视觉transformer，它结合了卷积和极端降采样。所有这些模型的共同点是，由于内存方面的考虑，在非常低的分辨率下提取具有表现力的特征映射，这不利于下游任务，如目标检测。

• 局部自注意力:密集预测任务涉及处理高分辨率图像。由于二次内存和计算需求，全局注意力在这种情况下是不可处理的。相反，采用了利用局部注意力的金字塔结构。对于此类方法，通常在每个窗口内执行自我注意，通常对全局上下文应用降采样。Liu等提出了移位窗口，表明窗口之间的通信优于独立窗口。Halo-Net扩展了每个窗口的邻域，以增加上下文和窗口间的通信。Twins采用了两阶段的自我注意。在第一阶段，采用局部注意力，而在第二阶段，对子采样窗口应用全局注意力注。然而，这些模型不是平移不变的，这是保持的一个属性。与本文的工作最接近的是独立注意力关注层(SASA)[Stand-alone self-attention in vision models]。正如后面章节所详述的那样，这一层施加了限制性的内存开销，并且速度明显较慢，与我们的层相比具有相似的精度(参见下图)。

  • 

  • 向前传递过程中的单层计算复杂度。QnA在速度和内存消耗方面优于SASA、HaloNet和本地自注意力基线。特别是，在正向传递期间，HaloNet需要至少x3的额外内存分配，而速度要慢x5。对于较大的内核，计算开销变得非常大，需要多达x10的额外内存分配。卷积层是内存效率最高的，但是对于大型内核，它们比QnA慢x1.8。所有实验都使用PyTorch进行测试，输入大小为256 × 256 × 64。

• 学习query的概念已经在其他设置的文献中进行了探索[Coordination among neural modules through a shared global workspace,Perceiver IO: A general architecture for structured inputs & outputs,Perceiver: General perception with iterative attention,Set transformer: A framework for attention-based permutation-invariant neural networks]。在Set Transformequery中，学习查询用于将输入维度投影到更小的输出维度，用于计算考虑或解码输出预测。类似地，percepver网络家族使用小的潜在数组来编码来自输入数组的信息。Goyal等人提出了对transformer架query修改，其中学习查询(共享工作区)作为令牌之间的通信通道，避免了二次的、成对的通信。与QnA不同，上述方法在整个输入序列上使用交叉注query在QnA中，学习到的查询在重叠的窗口之间共享。信息在本地聚合，利用强大的局部性先验，这种局部性先验已经通过卷积的大量使用得到了很好的建立。

Method

• Query和Attend是一个上下文感知的本地特征处理层。QnA的关键设计选择是一种类似卷积的操作，其中聚合核根据处理的局部区域的上下文而变化。QnA的核心是注意机制，其中重叠窗口被有效地处理以保持移位不变性。回想一下，从自注意力中的输入特征推导出三个主要实体:Query、key和Values。定义注意力权重的Query-Key点积在计算上可能会很昂贵。

• 为了克服这个限制，不再从窗口本身提取Query，而是学习它们(参见下图c)。这个过程在概念上类似于卷积核，因为学习的query决定如何聚合令牌值，重点注意力网络预定义的特征子空间。学习Query保持了注意力注机制的表达能力，并促进了一种仅使用简单快速操作的新型高效QnA实现。最后，本文的层可以扩展到执行其他功能(例如，下采样和上采样)，这在现有方法中是非常重要的[Scaling local self-attention for parameter efficient visual backbones]。

  • 

  • QnA Overview。局部层通过考虑重叠窗口(左)对图像进行操作，其中输出是通过聚合每个窗口内的信息来计算的:

    • (a)卷积通过学习应用于每个窗口的加权过滤器来进行聚合。

    • (b)独立自注意(SASA)通过自注意(self-attention)将窗口令牌组合在一起，这是一种消耗时间和内存的操作。

    • ©我们使用学习的Query，而不是所有的窗口元素彼此共享。这允许线性空间复杂性，同时保持注意力机制的表现力。

• 在详细解释QnA之前，我们将简要讨论卷积和自注意力的优点和局限性。我们设H和W为输入特征映射的高度和宽度，D为嵌入维数。否则，在本节中，我们将使用大写符号来表示矩阵或张量实体，使用小写符号来表示标量或向量。

Convolution

• 卷积层通过考虑输入特征 $X\in {\mathbb{R}}^{H\times W\times D}$ 的每个元素(如像素)的局部邻域来聚合信息。具体来说，给定一个核 $W\in {\mathbb{R}}^{k\times k\times D\times D}$，在位置(i, j)处的卷积输出为:

  • $$z_{i,j}=\sum _{(n,m)\in {\mathcal{N}}_k(i,j)}{x}_{n,m}\cdot W_{[k/2]+i-n,[k/2]+j-m},(1)$$

  • 位置(i, j)的k × k空间邻域是

  • $${\mathcal{N}}_k(i,j)=\{(n,m)|-k/2<(i-n),(j-m)\le k/2\}$$

  • (参见上图a)。为了简化表示法，我们从式(1)中省略k，将其重新写成矩阵表示法:

  • $$z_{i,j}=X_{{\mathcal{N}}_{i,j}}\cdot W,(2)$$

  • 为简洁起见，我们假设所有跨步操作的跨步为1，并应用填充来保持空间一致性。

• 卷积参数的数量在内核大小上是二次的，抑制了大型内核的使用，因此限制了捕获全局交互的能力。此外，跨不同位置重用卷积过滤器不允许基于内容的自适应过滤。然而，卷积的局部性和移位不变性有利于视觉任务。因此，卷积在计算机视觉网络中被广泛采用，深度学习框架支持公式(1)的硬件加速实现。

Self-Attention

• 视觉Transformer 网络通过自注意力机制将大小为N的序列混合，处理一个d维向量序列 $X\in {\mathbb{R}}^{N\times D}$。这些向量通常编码某种形式的图像补丁，其中N = H × W和H, W是每个空间维度上的补丁数量。具体来说，输入向量首先通过三个线性投影矩阵 $W_K,W_V,W_Q\in {\mathbb{R}}^{D\times D}$ 投影成键 $K=X{W}_K$，值 $V=X{W}_V$ 和query $Q=X{W}_Q$。那么，自注意操作的输出定义为:

  • $$SA(X)=Attention(Q,K)\cdot V=softmax(\frac{Q{K}^T}{\sqrt{D}})\cdot V,(3)$$

  • 其中Attention (Q, K)是一个大小为N × N的注意得分矩阵，它是使用应用于每行的Softmax计算出来的。

• 与卷积不同，自注意层有一个全局接受场，可以处理整个输入序列，而不影响学习参数的数量。此外，自注意力层的每个输出都是V值的输入依赖线性组合，而在卷积中，聚合在空间维度上是相同的。然而，自注意层的运行时复杂度和内存使用效率低下，使得它不太适合处理高分辨率输入。此外，有研究表明，香草Transformer不能很好地参与局部任务[Convit,ViT,Do vision transformers see like convolutional neural networks?,Training data-efficient image transformers & distillation through attention]，这是下游任务所需的先决条件。这些模型往往只有在经过长时间和data-hungry的训练过程后才会变得更加局部化。

Query-and-Attend

• 为了设计一个高性能的层，我们将自注意力机制调整为类似卷积的聚合操作。这背后的动机是，正如已经证明的[Coatnet,ViT]，自注意层比卷积层具有更好的容量，然而，卷积的归纳偏置具有更好的可转移性和泛化能力。具体来说，局部性和移位不变性先验(用于早期阶段)在图像域施加了强大的指导。

• 首先重新审视独立-自我注意方法(SASA)，其中注意力是在小的重叠k × k窗口中计算的，很像卷积。定义SASA的输出 $z_{i,j}$ 为:

  • $$z_{i,j}=Attention(q_{i,j},K_{{\mathcal{N}}_{i,j}})\cdot V_{{\mathcal{N}}_{i,j}},(4)$$

  • 换句话说，为了在本地聚合令牌，在每个本地窗口的令牌之间应用自注意力，并从窗口中心提取单个query(参见上图b)，其中 $q_{i,j}=X_{i,j}{W}_Q$。

• 虽然SASA具有表达性和局部性，但通过类似于输入自适应卷积的操作，它需要大量的内存使用。具体来说，据我们所知，所有公开可用的实现都使用展开操作，从输入张量中提取补丁。如果简单地实现此操作，则内存需求将增加 $k^2$ 。Vaswani等利用光环扩展的局部注意改进了SASA的内存需求。然而，这种实现需要比QnA多x3-x10的内存，同时速度要慢x5-x8，这取决于k(见上图)。这一限制使得SASA层无法处理高分辨率图像、使用更大的内核或使用相当大的批处理。

QnA - Single Query

• 为了缓解SASA的计算限制，我们通过引入学习Query来重新定义式(4)中的key-query点积。正如我们稍后将看到的，这种修改利用了权重共享原则(就像卷积一样)，并使QnA层的有效实现成为可能。首先用在训练过程中学习到的单个d维向量 $\hat{q}$ 替换方程(4)中的query $q_{i,j}$。更具体地说，我们定义位置(i, j)处QnA层的输出为:

• $$z_{i,j}=Attention(\hat{q},K_{{\mathcal{N}}_{i,j}})\cdot V_{{\mathcal{N}}_{i,j}},(5)$$

• 通过上述修改，我们将query-key点积解释为键在d维query方向上的标量投影。因此，根据标记值与query向量的相对方向聚合标记值。直观地说，key现在可以被提取，这样相关特征的键将与 $\hat{q}$ 紧密对齐。这意味着网络可query化查询方向来检测上下文相关的特征。

QnA - Multiple Queries

• 事实证明，在我们的范式下，性能可以进一步提高，计算开销最小，额外的内存可以忽略不计。原始的方法是在考虑多头注意时增加通道或注意头。虽然这增强了表现力，但额外的磁头会导致更大的内存占用和计算开销。为了提高层的表达性，我们可以使用l个不同的query $\hat{Q}\in {\mathbb{R}}^{L\times D}$ 而不是一个。然而，简单地在公式(5)中插入 $\hat{Q}$ 会导致L × D输出，这扩大了L的内存使用(也称为交叉注意)。相反，我们将Query学习到的注意图加权求和为单个注意图(针对每个注意头)，并使用它来汇总值。因此我们的QnA输出变成:

  • $$z_{i,j}=(\sum _{i\in [L]})W_i\ast Attention(\hat{Q},K_{N_{i,j}}))\cdot V_{N_{i,j}},(6)$$

  • 其中 $W\in {\mathbb{R}}^{L\times k^2}$ 是一个习得的权重矩阵，∗是元素的乘法运算。在这种情况下使用的整体额外空间是 $O(L\times k^2)$，相对较小，与原始解相反，它需要$O(L\times D)$额外空间。

QnA Variants

• 本文的层自然地适应了香草自注意力层所做的改进[Attention is all you need]。具体来说，我们在所有模型中都使用了相对位置嵌入和多头注意力(进一步的细节可以在附录F中找到)。

• 上采样和下采样使用QnA下采样可以很容易地实现使用跨行窗口。为了将令牌扩展到5倍，我们可以使用带有 $L=s^2$ 学习query的QnA层。query个查询的结果分配为上采样输出中的一个条目，我们有效地构建了一个大小为s × s的空间动态上采样核。我们在附录f中更正式地定义了上采样操作。我们表明，QnA可以用于有效地执行上采样函数，并提高了性能，表明它可以合并到其他视觉任务中，如图像合成。

Implementation & Complexity Analysis

• 跨窗口的共享学习Query允许我们使用现有深度学习框架(例如Jax)中可用的高效操作来实现QnA。特别是，query-key点积可以在整个输入序列上计算一次，从而避免了额外的空间分配。然后，我们可以使用基于窗口的操作来有效地计算重叠窗口上的softmax操作，导致线性的时间和空间复杂性(见上图)。我们方法的完整实现细节见附录D，以及Jax/ flex中的代码片段。

The QnA-ViT Architecture

• QnA- vit架构由视觉Transformer块(用于全局上下文)和QnA块(用于局部上下文)组成。QnA块与ViT块具有相似的结构，只是我们将多头自注意层替换为QnA层。我们提出了一系列遵循ResNet设计的架构。具体来说，我们使用一个4阶段的分层架构。基本尺寸D根据模型尺寸不同而不同。下面我们指出了我们在每个阶段使用的层数(T代表viti -block, Q代表qna -block):

  • T i n y : D , T , Q = { 64 , [ 0 , 0 , 4 , 2 ] , [ 3 , 4 , 3 , 0 ] } S m a l l : D , T , Q = { 64 , [ 0 , 0 , 12 , 2 ] , [ 3 , 4 , 7 , 0 ] } B a s e : D , T , Q = { 96 , [ 0 , 0 , 12 , 2 ] , [ 3 , 4 , 7 , 0 ] } Tiny:D,T,Q=\{64,[0,0,4,2],[3,4,3,0]\}\\ Small:D,T,Q=\{64,[0,0,12,2],[3,4,7,0]\}\\ Base:D,T,Q=\{96,[0,0,12,2],[3,4,7,0]\}\\ Tiny:D,T,Q={64,[0,0,4,2],[3,4,3,0]}Small:D,T,Q={64,[0,0,12,2],[3,4,7,0]}Base:D,T,Q={96,[0,0,12,2],[3,4,7,0]}

Experiments

Image Recognition & ImageNet-1K Results

• 设置:我们使用ImageNet1K基准来评估我们的方法，并遵循DEiT的训练配方，只是我们省略了EMA和重复增强。完整的训练详情请参阅附录B。

• 结果:不同模型之间的总结比较见下表。从表中可以看出，大多数基于transformer的视觉模型在top-1精度方面优于基于CNN的模型，即使CNN模型是使用强训练过程训练的。例如，使用标准ImageNet训练的ResNet50达到了76.6%的top-1准确率。然而，如[Resnet strikes back: An improved training procedure in timm]所述，通过更好的训练，其准确率可以飙升至80.4%。的确，这是一个非常令人印象深刻的改进，但与transformer模型相比还差得远。特别是，本文的模型(小版本)在ResNet上改进了1.3%，参数和FLOPs减少了40%。

  • 

  • ImageNet-1K预训练结果。所有模型在输入大小为224 × 224的条件下进行预训练和测试。标有↑384的模型后来也在 $384^2$ 分辨率上进行了微调和测试，如下[Fixing the train-test resolution discrepancy]。精度、参数计数和浮点操作在相应的出版物中报告。吞吐量使用timm库计算，在单个NVIDIA V100 GPU上使用16GB内存。对于 $Qn{A}_{7\times 7}$，使用的窗口大小是7×7而不是3×3。本文的模型达到了与最先进的模型相当的结果，参数更少，计算复杂度更低。

• 过研究上表中报告的视觉transformer，可以明显看出吞吐量差距。一种特别强的ViT是Focal-ViT ;在它的小版本中，它比ResNet101提高了0.7%，而后者的吞吐量提高了x1.4倍。尽管如此，本文的模型在速度和准确性的权衡方面脱颖而出。将QnA-Tiny与Focal-Tiny进行比较，我们的准确度仅降低0.5%，而吞吐量、参数计数和Flops次数改善了2倍。本文甚至可以通过训练具有更大接受场的QnA来缩小这个差距。例如，将QnA的接受域设置为7x7而不是3x3，可以实现82.0%的准确率，对模型速度和大小的影响可以忽略不计。

• 最后，我们注意到大多数视觉transformer都达到了类似的Top-1精度。更具体地说，Tiny模型(在参数和flop数量方面)实现了大致相同的Top-1精度，为81.2-82.0%。在较大的模型中，准确度差异甚至更不显著(例如，基本变体的准确度仅相差0.1%)，而且这种准确度差异很容易因许多因素而向任何一方倾斜，甚至通过选择不同的种子。尽管如此，我们的模型更快，同时使用更少的资源。

• 更好的准确性和效率权衡背后的原因:由于几个原因，QnA-ViT实现了更好的准确性和效率权衡。首先，QnA速度很快，这对提高吞吐量至关重要。此外，大多数视觉transformer的参数计数是由于线性投影矩阵。本文的方法通过省略query投影来减少线性投影的数量(即，$W_q$ 矩阵被query习查询取代)。此外，前馈网络比自注意力需要×2更多的参数。本文的模型使用比现有模型更小的嵌入维数而不牺牲精度。即NesTTiny使用192个嵌入维数，而SwinTiny和Focal-Tiny使用96个嵌入维数。另一方面，本文的方法实现了相似的特征表示能力，但维数较低，为64。

• 最后，其他参数高效方法通过对更大的输入图像进行训练来实现低参数计数。这被证明可以提高图像分类的准确性。然而，这是以低吞吐量和更多的flop为代价的。例如，在456 × 456分辨率的图像上进行训练和测试的effentnet - b5，仅使用30M个参数，准确率达到83.6%。尽管如此，网络的吞吐量是170张/秒，使用9.9 GFLOPs。与我们的基本模型相比，QnA以两倍的吞吐量实现了相似的准确性。此外，值得注意的是，这些模型是通过神经架构搜索优化的，这是一种更好的架构设计的自动化方法。我们认为，采用类似目的的方法[Nvit]甚至可以进一步优化本文的模型参数计数。

Ablation & Design Choices

• query数量:使query个查询允许我们捕获不同的特征子空间。我们将SASA作为基准，它从窗口元素中query自注意力查询。由于它占用大量内存，我们不能考虑类似于QnA-ViT的SASA变体。相反，我们考虑一种轻量级变体，它将本地注意力注与SASA结合在一起。所有SASA层都使用3x3的窗口大小。降采样的执行类似于QnA- vit，除了我们用SASA代替QnA。

• 最后，局部-自注意力层使用7x7的窗口大小，没有重叠(参见附录E.1)。结果总结在下表中。可以看到，我们取得了与SASA相当的结果。此外，两个query的性能优query个查询，但这种改进很快就会饱和。因此，我们query使用两个查询，因为它具有效率和表达性。

  • 

  • Multiple queries effect.。我们比较了SASA与QnA在不同query量下的性能。可以看到，使query个查询提高了QnA，达到了相当的性能，使用的内存减少了一个数量级。

• Number of heads:大多数视觉transformer采用大 head 尺寸(如≥32)。然而，我们发现QnA层有更多的头。我们用不同的训练设置训练了基于QnA和自注意力层的各种模型来验证这一点。我们的实验发现，head 维度d = 8最适合QnA层。与之前的工作[Going deeper with image transformers]类似，在混合模型中，同时使用了自注意层和QnA层，我们发现自注意层仍然需要较大的头部维度(即d = 32)。

• 此外，本文发现对于小型网络，使用更多的QnA磁头会更好(提高1%)。此外，当在没有强增强的情况下训练模型时，这种性能差距更明显(详见附录E.2)。直观地说，由于QnA层是局部的，因此它从局部模式识别中获益更多，而不像全局上下文那样需要表达性表示。

• 需要多少个QnA层?为了验证QnA的表达能力，考虑了十几个不同的模型。每个模型由四个阶段组成。在每个阶段，考虑使用自注意力层和QnA层。摘要报告见下表(完整报告见附录E.3)。在本文的实验中，我们得出结论，QnA层在早期阶段是有效的，可以在不影响模型性能的情况下取代全局注意力。QnA快速，提高了模型的效率。最后，QnA层是一个非常有效的下采样层。例如，我们考虑了两种主要由 transformer 块组成的基线架构，

  • (1)一个模型使用简单的2x2步进卷积来约简特征映射(在[Swin transformer]中采用)，

  • (2)另一个模型基于NesT[Aggregating nested transformers]中使用的下采样，即3x3卷积，然后是层规范和maxpooling层。

• 两种模型的准确率相近，均为81.2%。另一方面，当仅仅用QnA层替换下采样层时，看到在不增加参数计数和FLOPs的情况下，性能提高了0.7%。注意，要实现良好的性能，仍然需要全局的自我注意力。然而，它可以通过局部操作来减少，例如QnA。

• 深度模型:为了扩大我们的模型，我们选择在网络的第三阶段增加层数(如以前的作品[Deep residual learning for image recognition]所示)。这种设计选择主要是出于效率的考虑，其中空间和特征维度在第三阶段是可控的。特别是，本文将第三阶段的层数从7层增加到19层，并考虑四种配置，其中每种配置随所使用的QnA层数而变化。模型的精度如下表所示。从表中可以看出，减少全局注意力的数量可以保持模型的准确性。这表明，虽然注意力注可以捕获全局信息，但在一定程度上是有益的，并且出于效率考虑，架构设计可以施加注意力关注。

  • 

  • 消融研究和设计选择。在前两列中，本文指定了每个阶段使用的全局注意力层和QnA层的数量。更多细节请参见4.2节，更多配置请参见提供的材料。

Object Detection

• 设置:为了评估预训练网络的表示质量，我们使用了DETR框架，这是一个基于transformer的端到端对象检测框架。我们使用三个主干进行评估;ResNet50，以及QnA- vit的两个变体，即QnA- vit - tiny和QnA- vit - tiny -7x7, QnA- vit - tiny对所有QnA层使用7x7的接受器(而不是3x3)。完整的训练细节在补充材料中提供。

• 重新审视DETR设计:DETR与基于cnn的框架效果相当。然而，当在较小的物体上测试时，它的平均精度不太理想。DETR模型使用一个普通的transformer编码器来处理从骨干网络中提取的输入特征。如前所述，全局注意力受到地方性问题的影响。为了展示在现有基于transformer的网络中整合QnA的潜力，我们提出了der -QnA架构，其中两个transformer块被四个QnA块取代。

• 结果:我们在下表中报告了结果。可以看到，与ResNet50骨干网相比，QnA-Tiny训练的DETR获得了+2.2的更好AP。使用更大的接受野(7×7)进一步将AP提高0.4。然而，很大的改进是由于在大型对象上有更好的性能(+0.7)。最后，当将QnA合并到DETR编码器中时，我们获得了额外的+0.6AP(相对于使用DETR模型来说是+1.0AP)。更特别的是，将QnA与DETR结合在一起，在小物体上实现了令人印象深刻的+2.2 AP改进，这表明QnA的局域性的好处。

  • 

  • 基于COCO数据集的DETR目标检测。将qna - viti - tiny与DETR结合在ResNet50骨干网上显着改善(高达3.2)。接受域为7x7的QnA提高了大对象( $A{P}_L$ )的平均精度，将QnA纳入DETR网络可以提高较小对象(表明局域性)的性能。

QnA as an upsampling layer

• 我们建议QnA可以适用于除分类和检测之外的其他任务。为了证明这一点，我们使用$L_1$ 重建损失在CelebA数据集上训练了一个自编码器网络。我们考虑两个简单的基于卷积的基线。特别是，一个基线使用双线性上采样来升级特征映射，另一个基线使用转置卷积层。定性和定量结果分别如下图所示。图中显示，基于qna的自编码器实现了更好的定性和定量结果，并且引入了更少的工件(参见附录G了解更多细节)。

  • 

  • 定性自编码器结果。我们使用卷积层(a-b)和© QnA层训练一个简单的自编码器。我们展示了来自CelebA测试集的重建图像。QnA显示较好的重建。

  • 

  • 与上图中相同的卷积基线相比，报告了定量Auto-Encoder结果。我们将我们的方法(灰色)与双线性上采样(绿色)和基于转置卷积的上采样(粉红色)进行比较。在 $L_1$ 损耗(左上)、pSNR(右上)、SSIM(左下)和MSE(右下)指标上显示了跨时代(水平轴)的一致改善。

Limitations & Conclusion

• 我们提出了QnA，一种新颖的局部注意力层，具有线性复杂度，也是移位不变的。通过严格的实验，证明QnA可以作为通用层，在不影响精度部分的情况下提高视觉transformer的效率。此外，我们在目标检测设置中评估了我们的方法，并对现有的基于自注意的方法进行了改进。

• 我们的层也可以用作上采样层，我们认为这对于在其他任务中合并transformer是必不可少的，例如图像生成。最后，由于QnA是基于注意力的，它需要额外的中间内存，而卷积可以无缝地运行，不需要额外的分配。尽管如此，QnA比卷积具有更强的表达能力。

• 此外，全局自注意力块在捕获全局上下文方面更强大。因此，我们相信我们的层减轻了注意力注和卷积之间的差距，未来的工作应该结合所有三个层来实现最佳性能的网络。

A. Attention Visualization

• 在QnA中，每个窗口的聚合核是由学习到的query和窗口键之间的注意力得分得出的。为了可视化整个图像的注意力，我们选择将在所有相关窗口中获得的每个空间分配的分数相加。您可以在下图中找到可视化。如下图所示，注意力是内容感知的，这表明窗口聚合核是空间自适应的。

  • 

  • 不同头部的QnA注意可视化。为了可视化特定位置的注意力得分，我们将在所有窗口中获得的该位置的注意力得分相加。注意图被上采样并覆盖在图像上，以获得更好的可视化效果。

• 对于每个窗口中学习到的内核，请参见下图。

  • 

  • QnA聚合核可视化。注意核在可视化中被平铺而不是重叠，造成均匀的网格效果。较亮的区域表明注意力得分较高。(放大时观看效果最好)。

B. Full training details

B.1. Image Classification

• 我们使用ImageNet-1K基准测试来评估我们的方法，该基准测试包含来自1000个类别的1.28M张训练图像和50,000张验证图像。我们遵循DEiT的训练配方，只是省略了EMA和重复增强。特别是，我们使用AdmaW优化器训练了300个epoch的所有模型。

• 我们采用线性缩放学习率 $lr=5e-4\cdot \frac{batchsize}{256}$，预热时间根据模型大小和权重衰减 $wd=5e-2$ 而变化。对于增强，我们将RandAugment、mixup和CutMix与标签平滑和颜色抖动结合使用。最后，随机深度逐渐增大。请注意，这个训练配方(与以前的论文略有不同)正在成为在ImageNet基准上训练Vision Transformer的标准。最后，我们将所有QnA层中的query归一化为单位向量，以获得更好的训练稳定性。

B.2. Object Detection

• 我们在COCO 2017检测数据集上训练DETR模型，该数据集包含118K个训练图像和5k个验证集大小。我们利用DETR的训练设置，其中输入被调整大小，使得短边在480到800之间，而长边最多为1333。检测transformer设置初始学习率为1e−4，骨干网设置初始学习率为1e−5。由于计算限制，我们使用75 epoch的短训练调度程序，批大小为32。50次后，学习率按0.1缩放。我们使用[Scenic: A JAX library for computer vision research and beyond.]中的实现来训练DETR。

C. QnA-ViT architecture

• QnA- vit架构由transformer模块和QnA模块组成。首先，我们将图像分成4×4小块，并将它们投影成输入令牌。视觉transformer模块由多头注意层(MSA)和反向瓶颈前馈网络(FFN)组成，扩展4。block- 1的输出为:

  • $$\begin{gathered} z_l^{\prime }=MSA(LayerNorm(z_{l-1}))+z_{l-1} \\ {z}_l=FFN(LayerNorm(z_l^{\prime }))+z_l^{\prime } \end{gathered}$$

• QnA块具有类似的结构，只是我们用QnA层代替了MSA层。down采样是使用步长设置为2的QnA块执行的(为了启用跳过连接，我们使用步长相似的1×1-convolution)。我们采用预归一化来稳定训练。最后，我们在分类头之前使用全局平均池化，在池化操作之前使用LayerNorm。完整的体系结构细节可以在下表中找到。

  • 

  • QnA-ViT架构细节。QnA用于对两个连续阶段之间的特征映射进行下采样。在阶段3中，我们首先使用全局自注意力块。

D. Implementation & complexity - extended

• 在本节中，我们将提供有关有效实施QnA的详细信息。为了简化讨论，我们只考虑没有位置嵌入的单个query。让我们首先通过扩展注意层内部的softmax操作来检查QnA层的输出:

  • $$\begin{gathered} z_{i,j}=Attention(\hat{q},K_{{\mathcal{N}}_{i,j}})\cdot K_{{\mathcal{N}}_{i,j}} \\ =\frac{{\sum }_{N_{i,j}}{e}^{\hat{q}{K}_{n,m}}{v}_{n,m}}{{\sum }_{N_{i,j}}{e}^{\hat{q}{K}_{n,m}}} \end{gathered}$$

  • 回想一下， $N_{i,j}$ 是k×k-window在位置(i, j)。虽然看起来我们需要计算每个窗口的query-key 点积，但请注意，由于我们在每个窗口上使用相同的query，因此我们可以为整个输入计算一次 $\hat{q}{K}^T$。此外，我们可以利用矩阵乘法结合性，并通过首先计算 $\hat{q}{W}_k^T$ 来提高计算复杂度(这种融合实现通过避免分配关键实体来减少内存)。一旦我们计算了query-key点积，我们就可以使用许多深度学习框架支持的sumreduce操作有效地聚合点积。更具体地说，让 $Su{m}_k(...)$ 是对每个k × k窗口的值求和的函数，则:

  • $$z_{i,j}=\frac{{\sum }_{N_{i,j}}{e}^{\hat{q}{K}_{n,m}}{v}_{n,m}}{{\sum }_{N_{i,j}}{e}^{\hat{q}{K}_{n,m}}}=\frac{Su{m}_k(e^{\hat{q}{K}^T}\ast V)}{Sum(e^{\hat{q}{K}^T})}$$

  • 其中*是元素的乘法。我们的方法的伪代码可以在算法1中找到。我们进一步提供了QnA模块的代码片段，用Jax/ flex实现(参见下图)。

  • 

  • QnA模块在Jax和flex中实现的代码片段。使用预先训练的网络权重的全面实施将向公众开放。

  • 

• 复杂度分析:在单个query变体中，提取值和计算key-query点积需要 $2HW{D}^2$ 计算和 $HW+HWD$ 额外空间。此外，使用上述方法计算softmax需要额外的 $O(k^{2}HWD)$ 计算(用于求和和除法)和 $O(HWD)$ 空间(与k(即窗口大小)无关)。query多查询变量(L = 2)，见下图的经验比较。

E. Design choices - full report

E.1. Number of queries

• 为了验证使用多个query的有效性，我们训练了一个由局部自注意力块和QnA块组成的轻量级QnA- vit网络。我们将所有局部注意力注层的窗口大小设置为7x7，并对QnA层使用3x3的接受场。执行的所有降采样都是基于qna的。SASA基线使用了类似的架构，其中我们用SASA层替换了QnA层。每个阶段使用的QnA/SASA块数为[2,2,2,0]，局注意力关注块数为[1,1,5,2]。

E.2. Number of heads

• 使用更多的注意力头对QnA是有益的。更具体地说，我们进行了两个实验，一个是我们只使用QnA块，另一个是我们同时使用QnA块和(全局)自我注意力块。在第一个实验中，我们使用标准的ImageNet训练预处理，这意味着我们使用随机裁剪和随机水平翻转。在第二个实验中，我们使用DeiT训练预处理。我们在下表中显示了完整的报告。

  • 

  • The affect of number of heads on QnA。我们使用Inception预处理训练两个网络，即随机裁剪和水平翻转。我们根据ResNet-18和ResNet-50网络设置QnA块的数量(每个阶段的块数量在第一列中说明)。可以看出，在网络中传播时，使用固定的头维比增加头维要好。较浅的网络受益于有许多头脑，而较深的网络受益于更多的头脑。因此，我们建议增加深度网络的头维数，以获得更好的内存利用率。

  • 首先，从上表中，我们注意到，训练较少epoch的浅qnnet需要许多注意头。此外，最好在各个阶段之间保持固定的维度头部——这可以通过将两个连续阶段之间的头部数量加倍来实现。对于更深层次的网络，使用更多头的优势变得不那么重要了。这是因为网络可以通过利用其附加层捕获更多的特征子空间。最后，当同时使用QnA和ViT块时，仍然最好为QnA层使用更多的磁头，而对于ViT块，最好通过使用更少的磁头来使用高维表示(见下表)。

  • 

  • The number of heads affect on QnA and ViT Blocks.QnA仍然受益于更多的头部，而ViT块需要更高的维度表示，特别是对于更长的训练。

E.3. How many QnA layers do you need?

• 为了理解使用QnA层的好处，我们考虑了十几个结合了普通ViT和QnA块的网络体系结构。对于ViT块，我们尝试在早期阶段使用全局注意力，但发现最好使用局部注意力注，并将窗口大小限制为最多14x14。我们将架构选择分为三组:

  • 我们考虑在早期阶段改变QnA块的数量。

  • 我们在第三阶段改变QnA块的数量。

  • 我们使用较小的窗口大小的本地自注意力块。即，在所有阶段的所有ViT块中都是7x7窗口。

• 最后，经过DeiT预处理，对所有网络进行300次epoch的训练。完整的报告见下表。

  • 

  • 你真正需要多少QnA ?-完整报告。第一行表示每个阶段QnA和局部自注意(SA)块的数量。在前两个子表中，窗口大小为14 × 14，最后一个阶段设置了7 × 7的窗口大小。在最后一个子表中，我们将所有阶段的窗口大小减小为7 × 7。

• 从上表中可以看出，局部自我注意在早期阶段不是很有益，可以通过仅使用QnA块来省略。此外，在更深的阶段使用更多的全局注意力块会更好，但是通过使用相当数量的QnA块可以减少网络的延迟。最后，当使用较小的窗口大小时，局部自我注意变得不那么有效。特别是，由于QnA是移位不变的，它可以减轻缺乏跨窗口相互作用，这反映在我们使用更多QnA块时获得的改进增益中。

F. QnA Variants - extended version

• 正如文中所讨论的，我们在我们的层中加入了多头注意和位置嵌入。

• 位置嵌入自注意力是一种排列不变操作，这意味着它不假设输入标记之间存在任何空间关系。这个属性在图像处理中是不可取的，因为相对上下文是必不可少的。位置编码可以注入到自注意机制中来解决这个问题。根据最近的文献，我们使用了相对位置嵌入。这在注意力方案中引入了空间偏差，使Eq.(3)(来自主要文本)现在为:

  • $$Attention(Q,K)=softmax(\frac{Q{K}^T}{\sqrt{D}}+B),(7)$$

  • 其中， $B\in {\mathbb{R}}^{k\times k}$ 是一个习得的相对位置编码。注意，对于QnA层中的每个query学习不同的偏差，这增加了 $O(L\times k^2)$ 额外空间。

• 多头注意:与最初的自注意层一样，我们使用多个头部是为了允许QnA层同时捕获不同的特征。事实上，正如我们将在附录E.2中展示的那样，使用更多的注意力头对QnA是有益的。

• 在多头注意中，所有queryQ、键K和值V实体被分割成h个子张量，这些子张量对应于较低维空间RD/h中的向量。更具体地说，设 $Q^{(i)}、K{(i)}、V^{(i)} $ 分别query个查询、键和值的第i个子向量。头i的自我注意力变成:

  • $$\begin{gathered} hea{d}_i=Attention(Q^{(i)},K^{(i)})V^{(i)} \\ =softmax(Q^{(i)}{K}^{(i)T}/\sqrt{d_h}+B)V^{(i)},(8) \end{gathered}$$

  • 式中 $d_h=D/h$，多头注意力(MHA)的输出为:

  • $$MHA(Q,K,V)=Concat(hea{d}_1,...,hea{d}_h)W_O,(9)$$

  • 其中，$W_O\in {\mathbb{R}}^{D\times D}$ 为最终投影矩阵。

• 因子s的上采样可以使用 $s^2$ 学习的query来定义，即: $\hat{Q}\in {\mathbb{R}}^{s^2\times D}$。准确地说，每个窗口 $N_{i,j}$ 的输出表示为:

  • $$z_{i,j}=Attention(\hat{Q},K_{{\mathcal{N}}_{i,j}})\cdot V_{{\mathcal{N}}_{i,j}}$$

  • 注意，由Eq.(10)给出的窗口输出现在是一个大小为 $s^2\times D$ 的矩阵，而总输出Z是一个大小为 $H\times W\times s^2\times D$ 的张量。为了形成上采样输出 $Z^s$，我们需要重塑张量并排列其轴:

  • $$\begin{gathered} Z_{(1)}^s=Reshape(Z,[H,W,s,s，D]) \\ {Z}_{(2)}^s=Permute(Z_{(1)}^s,[0,2,1,3,4]) \\ {Z}^s=Reshape(Z_{(2)}^s,[H\ast s,W\ast s,D]) \end{gathered}$$

G. QnA as an upsampling layer

• 在本文中，我们展示了如何将QnA用作上采样层。特别是，我们训练了一个简单的自编码器网络，由五个下采样层和五个上采样层组成。我们使用L1重构损失作为目标函数来训练自编码器。我们考虑了三种不同的编码器层:

  • ConvS2-IN:步长2的3x3卷积，然后是实例规范化层。

  • ConV-in-max:步长1的3x3卷积，步长2的实例规范化层和maxpooling。

  • LN-QnA:层规范化，然后是一个3x3的单queryQnA层(没有跳过连接)。

• 还有三个不同的解码器层:

  • bilinear - conv-in: x2双线性上采样，然后是3x3卷积和实例归一化。

  • ConvTransposed-IN:一个2d的转置卷积，然后是实例归一化。

  • LN-UQnA:层归一化，然后是3x3上采样QnA层(没有跳过连接)

• 对于本文的基线网络，我们发现最好在编码器路径中使用ConvIn-Max，并选择bilineear-convin或ConvTranspose-IN作为解码器路径。对于基于QnA的自编码器，我们使用QnA层进行下采样和上采样。所有的网络都是在CelebA数据集上训练的，其中所有的图像都是中排的，并调整大小为 $256^2$ 的分辨率。使用Adam优化器对所有网络进行10个epoch的训练(根据最佳测试损失选择学习率)。
  范化，然后是一个3x3的单queryQnA层(没有跳过连接)。

• 还有三个不同的解码器层:

  • bilinear - conv-in: x2双线性上采样，然后是3x3卷积和实例归一化。

  • ConvTransposed-IN:一个2d的转置卷积，然后是实例归一化。

  • LN-UQnA:层归一化，然后是3x3上采样QnA层(没有跳过连接)

• 对于本文的基线网络，我们发现最好在编码器路径中使用ConvIn-Max，并选择bilineear-convin或ConvTranspose-IN作为解码器路径。对于基于QnA的自编码器，我们使用QnA层进行下采样和上采样。所有的网络都是在CelebA数据集上训练的，其中所有的图像都是中排的，并调整大小为 $256^2$ 的分辨率。使用Adam优化器对所有网络进行10个epoch的训练(根据最佳测试损失选择学习率)。