跳过不重要，关注关键点！ ToSA ，优化 Transformer 层的标记处理，为密集预测任务削减计算成本！

AIWalker

发布于 2024-06-27 13:02:16

710

发布于 2024-06-27 13:02:16

文章被收录于专栏：AIWalkerAIWalker

在本文中，作者提出了一种新颖的标记选择注意力方法，即ToSA，它可以识别需要关注以及可以跳过 Transformer 层的标记。更具体地说，一个标记选择器解析当前的注意力图并为下一层预测注意力图，这些注意力图随后用于选择应参与注意力操作的重要标记。剩余的标记简单地绕过下一层，并与被关注的标记连接起来，重新形成一个完整的标记集合。通过这种方式，作者减少了二次计算和内存成本，因为较少的标记参与自注意力，同时保持了整个网络中所有图像块的特征，这使得它可用于密集预测任务。作者的实验表明，通过应用ToSA，作者可以在保持ImageNet分类基准准确性的同时显著减少计算成本。此外，作者在NYU Depth V2的单目深度估计密集预测任务上进行评估，结果表明，使用带有ToSA的相当轻量级主干网络，作者可以达到类似的深度预测准确性。

1 Introduction

视觉 Transformer （ViTs）[5]是许多最新计算机视觉进展的核心，自注意力在生成关键视觉特征方面发挥着关键作用。然而，自注意力的操作与输入大小成二次方计算和内存成本。这使得在高分辨率图像上运行视觉 Transformer 以及在资源受限的设备上运行变得昂贵且具有挑战性。

研究者们探索了将注意力应用得更高效的方法，以使视觉 Transformer 更加高效。例如，Swin [12] 在图像的局部窗口内应用自注意力。其他研究者提出，先应用几层卷积层，然后仅在输入图像的显著下采样版本上应用自注意力以减少计算成本。还有一些论文设计了生成注意力的替代方法，例如移除softmax [10]，使用ReLU进行归一化[3]，在特征通道维度上应用注意力[13]，加性注意力[17]等。

然而，计算机视觉领域的许多最新进展仍然主要依赖原始的ViTs，这得益于它们的几个优点：设计简单，易于实现，存在强大的预训练方法和预训练预训练权重，它们能够随大量数据扩展[4]，并且由于标准自注意力的通用设计，它们在不同问题领域可能具有更一致的性能。

图1：（a）两个连续的标准 Transformer 层。（b）一个标准 Transformer 层后面跟着作者的ToSA层。作者提出的方法操作在连续的 Transformer 层对上，并修改第二个层使其成为标记选择性的。换句话说，只有一部分标记在第二个层中参与自注意力，其余的则绕过该层；作者的标记选择器预测标记的重要性得分以供选择。请注意，作者在整个层中保留所有标记，以方便密集预测任务，例如，分割，深度估计。

因此，研究者们试图在保持标准自注意力的同时减少ViTs的计算成本。一个常见的方法是随着网络减少标记[2, 7, 21, 23]。这有效地减少了与被关注标记数量成比例的计算和内存使用。然而，这些方法仅适用于分类任务。由于在推理过程中某些标记被丢弃或合并，这些网络不能用于需要所有图像像素/块独特特征的密集预测任务。在本文中，作者提出了一个新颖的标记选择性注意力方法，ToSA，它可以通过选择仅子集的标记来进行自注意力，使任何视觉 Transformer 更加高效。更具体地说，ToSA应用于两个连续的 Transformer 层，后者将变得标记选择性。鉴于第一层的多头注意力图，作者的标记选择器预测下一层的注意力图，并相应地产生标记的重要性得分。最重要的标记将进入下一层的自注意力，作者称之为_ToSA Transformer 层_；ToSA Transformer 层替换原来的第二个标准 Transformer 层。其余的标记简单地绕过ToSA层，并与被关注的标记重新合并。通过这种方式，作者减少了与被关注标记数量成二次方的计算和内存成本，保持了有利的标准自注意力操作，并在推理过程中保留了完整的标记集，以便网络可以作为密集预测任务的编码器。图1提供了ToSA的高层次说明。

总结来说，作者的主要贡献如下：

作者提出ToSA，这是一种新颖的方法，通过在某些层中减少参与注意力计算的标记数量，提高了多层视觉 Transformer 的效率。尽管在ToSA中并非所有标记都被关注，但它们在整个层中都被保留。这使得模型能用于密集预测任务。
为了选择要输入自注意力的标记子集，作者设计了一个标记选择器，它根据前一个 Transformer 层预测标记的重要性得分。这些得分用于确定需要在下一层关注的标记，其余的将简单地跳过下一层。
作者在标准的ImageNet分类基准[16]上评估了ToSA，表明ToSA在保持准确性的同时显著降低了计算成本。此外，作者将ToSA应用于视觉 Transformer Backbone 网络，并将其作为单目深度估计任务的编码器。作者展示了作者更高效的 Backbone 网络保持了深度预测的准确性。

2 Proposed Approach: ToSA

在本节中，作者介绍了作者提出的标记选择注意力方法，即ToSA。ToSA仅选择一组标记参与 Transformer 层的自注意力，这大大减少了计算和内存成本，同时在网络的整个过程中保持了完整的图像特征集，这使得它能够作为密集预测任务的编码器变得可行。

Self-Attention with Selected Tokens

标准自注意力机制：考虑一个输入，其中是标记的数量（即图像块的特征）而是特征维度；为了简洁起见，这里省略了批处理维度。在一个标准的视觉 Transformer 层中，首先通过线性层生成每个注意力头的 Query 、键和值如下所示：

其中, , 是线性变换矩阵。

然后，基于 Query 和键计算注意力图，然后与值相乘：

所有注意力头的输出被连接起来并输入到一个线性层，以生成该 Transformer 层的最终输出：

其中是一个线性投影矩阵。

标记选择注意力（ToSA）：在标准自注意力机制中，所有标记都参与注意力计算，导致O()的计算和内存成本。减少计算成本的一种方法是将参与自注意力的标记数量减少，即减少。

为此，作者需要确定应参与自注意力的那一组重要标记，以及应跳过注意力的剩余标记集。作者提出了一种新型的标记选择器，它在标准 Transformer 层和选择性注意力 Transformer 层（ToSA Transformer 层）之间操作。它消耗来自前一个标准层的未归一化注意力图（即，）并预测下一层的相应头的注意力图。这些预测用于确定重要标记。更具体地说，作者对每个头的预测注意力图进行列求和，为每个标记生成一个重要性分数。根据重要性分数，作者选择下一个层中执行自注意力的前K个标记，这作者称之为_注意力标记，而K是基于规定的注意力比率确定的。例如，表示下一个层中将有的标记参与自注意力。剩余的标记简单绕过下一个 Transformer 层，作者称之为_跳过标记。

在注意力标记通过应用标准自注意力的 Transformer 层之后，将经过注意的标记和跳过的标记结合起来，重新形成每个头的完整标记集，即：

其中是从的自注意力输出，而连接操作沿着标记维度进行。最后，将所有头的输出连接起来并通过一个线性层处理，以产生这个ToSA层的最终输出。图2为标准层和ToSA层的一对过程进行了说明。

Token Selector: Architecture and Training

图2给出了 Token 选择器架构的高级概述。Token 选择器接收来自前一个标准 Transformer 层的多 Head 未归一化注意力图（）作为输入。它包括两个带有中间ReLU激活的一维卷积层，后面是一个对数-softmax层，该层预测下一个层的完整多 Head 注意力图，；作者选择对数-softmax作为输出层，因为它生成归一化的注意力图，并且在训练中数值上更有利。

作者基于由预训练模型计算出的真实注意力图来训练 Token 选择器的注意力预测部分。更具体地说，给定一个预训练的 Backbone 网络（例如，在ImageNet分类[16]上训练的DeiT [19]），作者在连续的标准 Transformer 层和之间插入一个 Token 选择器。在训练期间， Token 选择器使用第一个 Transformer 层的图，其预测的注意力图与预训练模型的第二个 Transformer 层计算出的实际注意力图进行比较，这为训练 Token 选择器提供了监督，即：

其中表示KL散度损失。在这个过程中，预训练的 Transformer 层被冻结，仅训练 Token 选择器。请注意，作者仅使用用于预训练视觉 Transformer 模型的数据来训练作者的 Token 选择器，并不使用额外的训练数据。

Full Vision Transformer Model with ToSA

在给定一个预训练的视觉转换模型的情况下，作者可以将ToSA应用于任意一对层，其中第二层将被替换为一个ToSA层，并在它们之间训练一个 Token 选择器。以一个12层的标准视觉转换模型为例，作者可以将ToSA应用于连续的层对，并将第2层、第4层、第6层、第8层和第10层标准转换层替换为ToSA转换层；图3中提供了示意图。

一旦为将应用ToSA的层对训练好 Token 选择器，作者就冻结 Token 选择器，修改层对中的第二层以使其具有 Token 选择性，并在模型最初预训练的相同训练集上（例如，ImageNet）对整个模型进行微调。

3 Experiments

作者在ImageNet分类基准上评估ToSA。作者进一步将其用于单目深度估计，并展示其有效性。请注意，在这篇短文中，作者呈现了代表性的结果来展示作者提出方法的效能。作者计划在全文中包含更全面的结果。### 实验设置

数据集和评估：作者在ImageNet [16] 上执行标准的图像分类评估。作者进一步使用分类网络作为在NYU Depth v2 [18] 和KITTI [8] 上单目深度估计的编码器，作者在那里使用标准的深度准确度指标，如绝对相对误差（Abs Rel）和；这些指标的数学定义见[6]。

网络和实现：在这项工作中，作者使用DeiT-Tiny [19] 作为基础模型，因为它是标准的视觉 Transformer 之一，且用于训练的计算资源更容易获得，例如，作者使用2个Nvidia A100 GPU。对于ImageNet评估，作者使用ImageNet训练数据在ToSA上训练标记选择器并对DeiT-Tiny进行微调。由于DeiT-Tiny有12层，作者将ToSA应用于第2层、第4层、第6层、第8层和第10层，在每个ToSA层使用80%的标记进行自注意力。对于深度估计，作者使用NeWCRFs [22] 作为基础模型，并将编码器替换为DeiT-Tiny或带有ToSA的DeiT-Tiny。

Image Classification

表1总结了在ImageNet-1K分类基准上的评估结果。通过应用ToSA，作者在保持分类准确率的同时，显著减少了DeiT-Tiny的计算量，接近25%。由于保留了完整的标记集，作者的ToSA网络可以立即用作密集预测任务的编码器。另一方面，尽管现有工作如ToMe [1] 和 a-STAR [23] 也可以显著减少计算量，但它们依赖于减少标记的数量，即通过网络后只保留少量的标记。这对分类任务有效，但使得将这些模型用作密集任务的编码器变得具有挑战性。此外，这些现有方法会导致准确率下降。