将Transform骨干网络（T2T-ViT,LV-ViT）加速超31%，使用多准则Tokens融合方法实现高效ViT

Abstract

视觉Transformer（ViT）已成为计算机视觉领域的主要骨干网络。为了实现更高效的ViT，最近的研究通过剪枝或融合冗余的tokens来减少自注意力层的二次计算成本。然而，这些研究面临着由于信息丢失而导致的速度与精度的权衡。在本文中，我们认为在融合tokens时需要考虑tokens之间的多样化关系，以最大限度地减少信息丢失。为此，我们提出了一种多准则token融合（MCTF）方法，基于多准则（即相似性、信息量和融合后tokens的大小）逐步融合tokens。此外，我们利用了一步提前注意力（one-step-ahead attention），这是捕捉tokens信息量的改进方法。通过使用token减少一致性训练配备MCTF的模型，我们在图像分类（ImageNet1K）中达到了最佳的速度与精度的权衡。实验结果证明，MCTF在有无训练的情况下，一致超越了之前的减少方法。具体来说，带有MCTF的DeiT-T和DeiT-S分别减少了约44%的FLOPs，同时分别提升了0.5%和0.3%的性能。我们还展示了MCTF在各种视觉Transformer中的适用性（如T2T-ViT、LV-ViT），实现了至少31%的加速且无性能下降。

源代码：https://github.com/mlvlab/MCTF

Introduction

视觉Transformer（ViT）最初是为自然语言处理任务开发的自注意力机制，现已被提议用于解决视觉任务。随着ViT的出现，Transformers成为了广泛应用于各种视觉任务（如分类、目标检测、分割等）的主流架构。与传统方法（如卷积神经网络，CNNs）相比，仅由自注意力和MLP构建的ViTs提供了极大的灵活性和令人印象深刻的性能。然而，尽管具有这些优势，随着tokens数量的增加，自注意力的二次计算复杂性成为了Transformers的主要瓶颈。随着对大规模基础模型（如CLIP）的兴趣日益增加，这一限制变得更加显著。为此，一些研究提出了高效的自注意力机制，包括在预定义窗口内的局部自注意力。
最近，越来越多的研究致力于通过优化ViTs中的token减少方法，而不改变其架构。早期的研究主要集中在剪枝无信息量的tokens以减少token数量。另一类研究则尝试融合tokens而不是丢弃它们，以尽量减少信息丢失。然而，大多数token融合方法仍然常常出现性能下降的情况。我们注意到，token融合方法通常只考虑单一准则（如tokens的相似性或信息量），导致次优的token匹配。例如，基于相似性的token融合容易将前景tokens组合在一起，而基于信息量的融合则经常合并明显不同的tokens，导致表示的塌陷。此外，如果过多的tokens被融合为一个token，则信息丢失是不可避免的。
为了解决这些问题，我们引入了多准则token融合（MCTF），它基于多准则优化视觉Transformers的token融合。与之前只考虑单一准则的工作不同，MCTF通过以下多准则来测量tokens之间的关系：（1）相似性以融合冗余tokens，（2）信息量以减少无信息量的tokens，（3）tokens的大小以防止由于大尺寸tokens导致的信息丢失增加。为了应对连续层注意力图之间的不一致性，我们采用了一步提前注意力，该方法明确估计了下一层tokens的信息量。最后，通过引入用于微调模型的token减少一致性，我们的模型在现有工作中表现出更优的性能，如图1所示。令人惊讶的是，我们的MCTF甚至比“完整”的基础模型（红色虚线）在计算复杂度降低的情况下表现更好。具体而言，在DeiT-T和DeiT-S中，它分别带来了0.5%和0.3%的性能提升，同时减少了约44%的FLOPs。我们在T2T-ViT和LV-ViT中观察到了类似的加速（31%），且无性能下降。

3.Method

我们首先回顾了自注意力和token缩减方法（第3.1节）。然后，我们介绍了利用一步预测注意力的多准则token融合（第3.2节）及其所需的步骤（第3.3节）。最后，我们在第3.4节中引入了一种结合token缩减一致性的训练策略。

3.1 预备知识

在Transformer中，tokens

通过自注意力机制进行处理，定义如下：

其中

，而

是可学习的权重矩阵。尽管自注意力具有出色的表达能力，但由于其二次时间复杂度

，其计算量会随着tokens数量𝑁的增加而迅速增加。为了解决这一问题，一些研究通过简单地修剪不具有信息量的tokens来减少token数量。然而，这些方法通常会导致信息损失，从而显著降低性能。因此，另一部分研究选择将冗余或不重要的tokens融合为一个新的token

，其中𝑋是原始tokens集合，𝛿表示融合函数，例如最大池化或平均池化。在本文中，我们也采用了‘token融合’而非‘token修剪’，并通过多准则最小化token缩减造成的信息损失。

3.2 多准则token融合

给定一组输入tokens

，MCTF（多准则token融合）的目标是将tokens融合为输出tokens

，其中𝑟为融合的token数量。为尽量减少信息损失，我们首先基于多准则评估tokens之间的关系，然后通过双向二部软匹配来分组和融合这些tokens。多准则吸引函数：我们首先定义一个基于多准则的吸引函数𝑊，如下所示：
 

其中 

是根据第𝑘个准则计算的吸引函数，而是用于调整第𝑘个准则影响的温度参数。两个tokens之间的吸引分数越高，越有可能被融合。在本文中，我们考虑了以下三个准则：相似性、信息量和大小。相似性：第一个准则是tokens的相似性，用于减少冗余信息。与之前的工作需要tokens的接近性类似，我们利用tokens集合之间的余弦相似度：
 

基于相似性的token融合有效地消除了冗余tokens，但它往往会过度地结合信息量大的tokens，如图2b所示，导致信息损失。
信息量：为了最小化信息损失，我们引入了信息量以避免融合信息丰富的tokens。为了量化信息量，我们在自注意力层中测量平均注意力分数

，表示每个token对其他tokens的影响：

当时，其他tokens没有影响。根据信息量分数，我们定义了一个基于信息量的吸引函数：
 

其中和 分别是和 的信息分数。当两个tokens都不重要时

，权重变高（

），使得这两个tokens更容易被融合。图2c显示了结合相似性和信息量权重后，前景对象中的tokens融合较少。大小：最后一个准则是tokens的大小，表示被融合tokens的数量。尽管tokens并未被丢弃，而是通过融合函数（如平均池化或最大池化）合并，但随着组成tokens数量的增加，保留所有信息变得困难。因此，更倾向于在较小tokens之间进行融合。为此，我们最初将tokens𝑋的大小设置为1，并跟踪每个token的组成（融合）tokens数量，并定义一个基于大小的吸引函数：
 

在图2d中，tokens基于多准则（相似性、信息量和大小）被融合。我们观察到相似tokens之间的融合发生，且前景tokens或大tokens的融合得到了适当的抑制。
双向二部软匹配：给定基于多准则的吸引函数W，我们的MCTF执行了一种称为二部软匹配的松弛双向二部匹配。二部匹配的一个优势在于它减轻了tokens之间相似度计算的二次成本，即

,其中

.此外，通过放松一对一的对应约束，可以通过高效算法获得解。在此松弛匹配问题中，首先将token集合𝑋分割为源和目标

,如图3的步骤1所示。给定一组二进制决策变量，即和之间的边矩阵

，二部软匹配可表述为：

其中满足
 

其中
 

表示token和之间存在边的变量，而

此优化问题可以通过两个简单的步骤解决：1）为每个𝑖找到最大化的最佳边，2）选择具有最大吸引分数的前𝑟条边。然后，基于软匹配结果，我们将tokens分组为
 

其中表示与匹配的tokens集合。最终的融合结果 可表示为
 

其中
 

是一个简单的融合操作（例如，最大池化），如图3的步骤2-3所示。

3.3 一步预测注意力的MCTF

如图1所示，我们通过引入一步预测注意力来进一步提升多准则token融合的性能。相较于在自注意力层之后执行token融合，使用一步预测注意力的MCTF（SAMP-MCTF）能够显著减少融合过程中信息的丢失。一步预测注意力：为了避免重叠融合问题，我们在自注意力中引入一步预测（或称前馈）机制。不同于通过直接计算注意力分数来判断tokens之间的关系，我们在SAMP-MCTF中首先通过多准则吸引函数计算tokens的初步关联，然后通过一步预测注意力机制来引导实际融合。这种机制不仅保留了token之间的多准则信息，还通过前馈预测有效地引导了token融合。

3.4 结合token缩减一致性的训练策略

为了训练SAMP-MCTF，我们设计了一种结合token缩减一致性的训练策略。具体来说，我们不仅最小化最终的分类或回归损失，还通过引入一致性损失来保证融合前后token的语义一致性。这种策略有助于提高SAMP-MCTF在下游任务中的性能，并增强其对信息丢失的鲁棒性。

Experiment

为了验证所提出方法的有效性，我们将MCTF与之前的token减缩方法进行了对比。为了进行比较，我们选择了token剪枝方法（如A-ViT、IA-RED2、DynamicViT、EvoViT、ATS）和token融合方法（如SPViT、EViT、ToMe、BAT）在DeiT中的表现，并报告了每种方法的效率（FLOPs (G)）和性能（Top-1 Acc (%)）。此外，为了验证MCTF在其他视觉Transformer（如T2T-ViT、LV-ViT）上的有效性，我们报告了MCTF的结果并将其与现有工作中的官方数据进行对比。我们在表1和表2中使用下标r表示每层减少的token数量。表格中的灰色表示基准模型，绿色表示性能的提升，红色表示相较于基准模型的性能下降。

token减缩方法的比较

表1总结了与现有token减缩方法的比较结果。我们展示了MCTF在DeiT中以最低的FLOPs实现了最佳性能，超越了所有之前的工作。此外，值得注意的是，MCTF是唯一在DeiT-T和DeiT-S中实现最低FLOPs且无性能下降的工作。通过对DeiT-T进行30个周期的微调，MCTF在降低近一半FLOPs的同时，准确率相较于基准模型显著提高了+0.5%。类似地，我们观察到DeiT-S的性能提升了+0.3%，同时FLOPs减少了-2.0 G。我们认为，多准则融合加上一步预测注意力有助于模型最大限度地减少信息损失；通过token减缩在类别token上的一致性损失进一步提高了模型的泛化能力。

MCTF在其他视觉Transformer中的应用

为了验证MCTF在不同ViTs中的适用性，我们展示了MCTF在其他Transformer架构中的表现（见表2）。根据先前的工作，我们将MCTF应用于LV-ViT。同时，我们还展示了MCTF与T2T-ViT的结果。如表中所示，我们的实验结果非常有前景。MCTF在这些架构中至少实现了31%的速度提升，而没有性能下降。此外，与LV-ViT结合的MCTF在FLOPs和准确率方面超过了所有其他Transformer和token减缩方法。特别值得注意的是，除MCTF外，所有token减缩方法在LV-ViT中均带来了性能下降。这些结果表明，MCTF是一种适用于多种视觉Transformer的高效token减缩方法。

无需训练的token减缩

与ToMe类似，MCTF可以在预训练的ViTs中应用，而无需额外的训练，因为MCTF不需要任何可学习的参数。我们在此将两种减缩方法应用于预训练的DeiT而不进行微调，并在表3中提供了结果。无论每层减少的token数量r是多少，MCTF始终优于ToMe。特别是在最稀疏的设置r=20中，性能差距显著（在DeiT-T中+7.0%，在DeiT-S中+3.8%）。值得注意的是，在没有任何额外训练的情况下，我们的MCTF在预训练的DeiT-S上仍表现出79.2%的竞争性能，相比于需要训练的减缩方法（例如表1中的A-ViT为78.6%，IA-RED2为79.1%，DynamicViT和SPViT分别为79.3%）。

4.2 关于MCTF的消融研究

我们提供了消融研究来验证MCTF的各个组件。除非另有说明，我们所有实验均在微调后的DeiT-S上进行，MCTF的r = 16。通过调整每层减少的token数量r ∈ [1, 20]，我们提供了FLOPs-准确率图。

多准则

我们在图7a中探讨了多准则的有效性。首先，在多准则方面，我们为MCTF采用了三个准则，即相似性（sim.）、信息量（info.）和大小。相较于双准则（sim. & info.）和多准则（sim. & info. & size），单一准则的相似性和信息量显示出相对较低的性能。具体来说，当r = 16时，单一准则的性能分别为79.7%和79.4%（相似性和信息量）。然后，采用双准则（sim. & info.），MCTF达到79.8%的准确率。最后，通过尊重所有三个准则（sim. & info. & size），我们获得了80.1%的准确率，提升了+0.3%。随着r的增加，这些性能差距变得更大，证明了多准则在token融合中的重要性。

一步预测注意力与token减缩一致性

为了证明一步预测注意力与token减缩一致性的有效性，我们还提供了有无每个组件的MCTF结果（见图7b）。当去除一步预测注意力或token减缩一致性时，每个FLOP的准确率均有所下降。这种显著的下降表明这两种方法对MCTF至关重要。简言之，通过采用一步预测注意力和token减缩一致性，MCTF在广泛的FLOPs范围内有效地减轻了性能下降。

设计选择的比较

表4中展示了设计选择的消融结果。首先，我们的双向二部匹配，能够在两组中捕捉到双向关系，增强了与单向二部匹配相比的准确性。其次，对于池化操作δ，考虑大小s和注意力的加权和比其他方法（如最大池化或平均池化）更为优越。最后，我们比较了精确和近似注意力的结果。对于精确注意力，我们分别对一步预测注意力和融合后的自注意力层中的注意力进行相似性计算。否则，我们按照第3.3节所述，通过一步预测注意力来近似。正如表中所示，我们的近似注意力在效率上取得了显著改进（FLOPs减少了0.4 G），同时保持了性能。

4.3 MCTF的分析

定性结果

为了更好地理解MCTF，我们在图8中提供了MCTF的定性结果。我们在ImageNet-1K上可视化了DeiT-S最后一个块的融合token，并通过相同的边框颜色标注了融合的token。如图所示，由于token是通过多准则（如相似性、信息量、大小）进行融合的，我们在信息量较大的前景对象中保留了更多多样化的token。例如，在第三张仓鼠的图像中，背景块（包括手）被融合为一个token，而前景token融合得较少，同时保留了仓鼠的眼睛、耳朵和面部等细节。简言之，与背景相比，前景token的融合较少，适中的大小保留了主要内容的信息。

大小准则的合理性

图9展示了有无大小准则的token减缩后的大小分布直方图。具体来说，我们测量了最后一个块中最大token的大小，并提供了直方图。在有大小准则的情况下，融合的token倾向于具有较小的大小，分别显示平均大小为39.3/49.2（有和无大小准则）。如预期的那样，MCTF成功地抑制了导致信息丢失的大型token，从而提高了性能。

结论

本文的贡献总结如下：

1.提出了一种新颖的多准则token融合方法，考虑了相似性、信息量和大小等多准则，以捕捉tokens之间的复杂关系并最小化信息丢失。
2.为了测量tokens的信息量，利用了一步提前注意力，以在后续层中保留关注的tokens。
3.提出了一种新的微调方案，通过token减少一致性来提升配备MCTF的transformers的泛化性能。广泛的实验表明，MCTF在各种ViTs中实现了最佳的速度-精度权衡，超越了所有之前的token减少方法。

文章引用自：

Multi-criteria Token Fusion with One-step-ahead Attention for Efficient Vision Transformers

关注我的公众号auto_driver_ai(Ai fighting), 第一时间获取更新内容。

AiFighing是全网第一且唯一分享自动驾驶实战，以代码、项目的形式讲解自动驾驶感知方向的关键技术，从算法训练到模型部署。