Backbone创新 | 中科大联合百度提出全新Transformer Backbone

集智书童公众号

发布于 2023-09-04 10:54:00

2260

发布于 2023-09-04 10:54:00

当前，Vision Transformer（ViT）在许多视觉任务中占据主导地位。通过Token稀疏化或降维（在空间或通道上）来解决其Token多Head自注意力（MHSA）的二次复杂度缺点已经得到广泛应用。然而，MHSA中存在的冗余通常被忽视，同样忽视了前馈网络（FFN）的冗余。为此，作者提出了Hallucinate注意力图和FFN压缩来填补这些空白。具体而言，作者观察到原始ViT中存在类似的注意力图，并提议使用更简单的操作从剩余的一半注意力图中生成Hallucinate注意力图，这被称为Hallucinate-MHSA（hMHSA）。至于FFN，作者对其隐藏到输出投影矩阵进行分解，并利用重参化技术加强其能力，使之成为紧凑的FFN（cFFN）。通过作者提出的模块，在基于ViT的各种Backbone结构（包括简单的DeiT，混合的NextViT和分层的PVT结构）中，可以实现10％-20％的浮点操作（FLOPs）和参数（Params）的减少，同时性能也非常有竞争力。

1、简介

最近，Vision Transformer（ViT）在各种计算机视觉任务上取得了显著的性能提升，例如图像分类、目标检测和语义分割。同时，广泛认为，由于Token-Level多Head自注意力（MHSA）模块带来的二次计算复杂度，Vision Transformer的计算成本较高。这种模型复杂性使得视觉Transformer在实际应用中的推理速度是一个关Key 因素时，相对于卷积神经网络（CNN）处于劣势，特别是与高效的CNN变体相比。

因此，计算机视觉界对于平衡Vision Transformer的效果和复杂性给予了很大关注。以视觉Transformer的复杂性为根本原因，先前的研究主要集中在设计更简单的多Head自注意力（MHSA）。一般来说，现有的关于缓解MHSA复杂性的方法要么选择稀疏化Vision Transformer的输入Token，要么在计算注意力图时减少空间和通道维度。

例如，在[DynamicViT, Ia-red2, AdaViT]中提出了舍弃不重要Token的方法，PVT设计了分层结构和空间降维注意力（SRA），Swin-Transformer利用移位窗口自注意力（SW-SA）避免了直接的全局自注意力，而ScalableViT中在实现MHSA时使用了通道维度压缩的方法。

这些方法确实有效地降低了MHSA的复杂性，然而，它们中没有一个提出了以下问题：

作者是否真的需要通过昂贵的Query和Key信号之间的相关计算来获得每个注意力图呢？

另外一个Value得一提的是，除了广受关注的MHSA之外，现有研究往往忽视了最初在传统的Vision Transformer模块中提出的前馈网络（FFN）。同时，在最近有意设计的基于Transformer的Backbone网络（例如PVT和NextViT）大大降低了复杂性的情况下，作者想知道“FFN的计算成本是否真的微不足道，并且还能采取哪些进一步措施来平衡FFN的效率和效果？”

为了回答这些问题，作者首先调查了传统MHSA模块的冗余性。作者可视化了3种基于Transformer的Backbone网络DeiT-S、PVTv2和NextViT中每个模块的注意力图。以DeiT-S的第2个模块为例，作者在图1a中展示了注意力图。

此外，作者经验性地发现这些网络中注意力图之间的平均余弦相似度超过50%，这将在表4中展示。因此，通过复杂的Query-Key相关计算生成的注意力图是冗余的或相关的。然后，如图1b所示，作者通过参数数量和浮点运算数（FLOPs）分析了多个Backbone网络的模型复杂性。

很明显，FFN始终占据这些模型中相当大的计算成本，但通常被忽视。这与先前的研究结果[10, 11]一致。以上观察结果表明，仍然存在改进的设计空间，激励作者记住注意力图之间的相似性，以开发更好的MHSA设计，并寻求更精心设计的FFN以降低成本。

在这项工作中，作者提出了使用Hallucinate注意力图和FFN压缩的视觉Transformer设计。受到GhostNNet的启发，该方法直接从CNN的Conv输出中获得特征图，作者设计了一种Hallucinate的MHSA（hMHSA），通过使用更简单的操作将一半的注意力图从另一半中生成，而不是通过昂贵的Q-K相关计算获得所有的注意力图。通过这种方式，模型可以有效地减少MHSA的复杂性。令人惊讶的是，这些通过简单方式生成的注意力图甚至更少冗余，具体结果将在4.2.2节中展示。

此外，受到FFN中通常存在冗余的事实的启发，作者提出了紧凑型FFN（cFFN）来节省FLOPs。作者不是直接降低扩展比例（即压缩FFN隐藏维度），这必然会降低模型的能力和性能，而是首先使用矩阵分解来减轻FFN中隐藏层到输出层的投影矩阵的计算量，以实现冗余抑制，然后利用现成的重参化技术来使其更紧凑并增强其能力。

最后，作者将hMHSA和cFFN应用于各种Backbone网络，包括直接结构（DeiT）、混合结构（NextViT）和分层结构（PVT）。

实验结果表明，在具备竞争性能的同时，可以实现10%-20%的FLOPs和参数节省。本文的贡献如下：

作者利用每个MHSA中注意力图之间的相关性/冗余，提出了Hallucinate MHSA（hMHSA）。hMHSA通过使用简单的操作从另一半生成一半的注意力图，以实现更好的效率和效果平衡。
另外，作者设计了紧凑型FFN（cFFN），其中隐藏到输出投影矩阵进行了分解以减少冗余，并利用现成的重参化技术来压缩FFN。
作者将提出的hMHSA和cFFN应用于各种Backbone网络，并通过实验证明了其有效性，并在具备竞争性能的同时实现了10%-20%的FLOPs和参数减少。

2、本文方法

作者提出了hMHSA和cFFN用于Vision Transformer的构建。一般而言，它们可以直接用来替代许多基于ViT的框架中的MHSA和FFN模块，以实现更好的效益和效率平衡。

2.1 hMHSA模块

在介绍作者提出的hMHSA之前，让作者先回顾一下原始的MHSA。对于输入

X \in R^{N×C}

，其中

为序列长度，

为通道维度，MHSA模块首先生成3个可训练的线性转换矩阵，得到Query (Q)、Key (K)和Value(V)

\in R^{N×C}

。然后，对于多Head计算，它们被Reshape为

R^{h×N×C_h}

，其中h表示Head的数量，

C_h = C/h

表示每个Head的维度。然后，原始MHSA的计算过程可以总结如下：

其中，

Scale(·)

是缩放操作，用于避免点积的结果过大，

Proj(·)

是一个可训练的投影操作，用于融合来自不同Head 的信息。

首先，

的每个Head

Q_i

与对应的K的Head

K_i

进行缩放点积计算，得到

A^{i,:,:}_{pre} \in R^{N×N}

。为了方便起见，作者使用

A_{pre} \in R^{h×N×N}

表示

Softmax(·)

之前的注意力图，而

Softmax(·)

之后的注意力图表示为

A_{post} \in R^{h×N×N}

。然后，

A_{post}

作为

的检索器，并且

Proj(·)

将检索结果的所有Head连接起来，并在其上应用一个全连接层，得到输出

。

受到Head 维度

A_{pre}

中相似性和冗余性发现的启发，作者设计了 hMHSA 模块，以更简单的操作从另一半中Hallucinate出一半的注意力图，以达到更高的效率。

如图 2 所示，普通 MHSA 模块和作者的 hMHSA 模块之间的差异有两个：

生成的

\hat Q

和

\hat K

；

Softmax(·)

之前插入的IHH和CHH模块。

为了演示清楚，作者还制定了hMHSA模块的计算流程如下：

其中

IHH(·)

代表Intra-Head Hallucination 操作，

CHH(·)

代表 Cross-Head Hallucination 操作。为了方便起见，作者将

Softmax(·)

之前的真实（由

\hat Q

和

\hat K

计算）注意力图表示为

A^r_{pre}

，将更简单的

IHH(·)

和

CHH(·)

生成的Hallucinate注意力图表示为

A^h_{pre}

。与 MHSA 不同，hMHSA 生成

\hat Q,\hat K ∈ R^{N× C/2}

, 并将它们Reshape为

R^{h/ 2 ×N×C_h}

以获得

A^r_{pre} ∈ R^{h/2×N×N}

。这里，为了方便起见，作者假设Head数

是偶数。同时，V 的生成方式与普通 MHSA 模块相同。

IHH(·)

和

CHH(·)

它们都是使用传统的卷积网络模块来实现的。在

IHH(·)

阶段，将

A^r_{pre} ∈ R^{h/2×N×N}

中的每个实

N×N

注意映射Reshape为

N×h×W

，其中

和

是输入

的原始空间维度，将

A^r_{pre}

变为

A^r_{pre}∈R ^{N×h/2×h×W}

。然后，在Head 维度上应用3×3深度卷积。

图2的中间部分说明了IHH在一张注意力图中的过程。在执行

IHH(·)

之后，应用

CHH(·)

对不同注意力Head之间的交互进行建模。受DeepViT的启发，

CHH(·)

用于执行变换，以沿Head 维度聚合多Head注意力图。该操作进一步对Head之间的多样性进行建模，并最终产生

A^h_{pre}

，并且

CHH(·)

可以通过1×1卷积层方便地实现（见图2的最右边部分）。在生成

A^h_{pre}

之后，作者将其与

A^r_{pre}

沿着Head 维度连接以获得

A_{pre}

，并且hMHSA的其余部分与MHSA相同。

请注意，作者的运算

IHH(·)

被设计为深度卷积，它也在GhostNet中使用。然而，它们本质上是不同的。GhostNet在通道维度上采用深度卷积，主要考虑其轻量级特性。相反，作者的操作是在Head 维度上进行的，目的是进一步建模相邻特征点之间的隶属关系信息。

元素

\hat A^r_{pre}(n,i,h,w)

表示第

个Head的1D形式（N）的第n个Token

P_{1D}

和2D形式（H×W）的第（h，w）个Token

P_{2D}

之间的相关性。

因此，深度3×3卷积可以学习局部感受野中不同标记之间的亲和力，即

P_{2D}

及其8个邻居都将有助于

P_{1D}

和

P_{2D}

之间的注意力得分。它不仅为MHSA提供了更强的建模能力，而且节省了计算成本。

2.2、cFFN Module

首先，作者还简要回顾了Transformer块中的普通FFN模块，其中包括一些符号滥用。对于输入的

X∈R^{N×C}

，当FFN隐藏层的展开比为m时，FFN的工作原理如下：

其中

M_1∈R^{C×mC}

是将

变换到隐藏维度的矩阵，

Act(·)

是非线性激活，

M_2∈R^{mC×C}

是隐藏到输出投影矩阵。

如MobileNetv2所示，低维流形的非线性变换会导致信息丢失，其表现为流形中的某些点相互塌陷。这意味着模块可能会受到一些辨别点减少的影响，从而导致建模能力较差。

因此，为了节省计算成本，作者声明不直接降低扩展率。相反，作者提出了cFFN模块来保持

M_1

和扩展比，以避免这种信息丢失，同时对

M_2

进行因子分解以进行FFN压缩：

其中

U∈R^{mC×k}

，

V∈R^{k×C}

表示

M_2

的因子分解的两个小矩阵，

是维数。特别地，当

等于

mC/(m+1)

时，cFFN的FLOP与标准的FFN相同。

通常，作者设

k=tmC/(m+1),t∈(0,1)

，作者称

为compact ratio。值得注意的是，因子分解可能会损害性能，为了进行压缩，作者在实际实现中利用了重参化技术。如图3（a）所示，在两个因子分解矩阵的训练阶段，使用Conv 1×1的r个分支，然后是Norm（BN）。

在推理阶段（图3（b）），给定Conv和BN的线性，

个分支被合并为方程4的形式。作者要强调作者的贡献，即确定和解决长期被忽视和处理不力的FFN复杂性问题。利用重参化技术作为一种现成的方法，作者能够在不牺牲性能的情况下压缩FFN模块。

2.3、计算复杂度分析

一个标准的ViT块由一个MHSA模块和一个FFN模块组成。在MHSA模块中，输入特征形状为N×C，转换X到Q、K和V的3个Conv 1×1贡献了

3NC^2

的FLOPs，

QK^T

和

A_{post}V

都贡献了

N^2C

，

Proj(·)

贡献了

NC^2

。因此，其复杂度总体上等于：

对于FFN模块，假设展开比为

，则计算复杂度为：

在作者的hMHSA模块中，Conv 1×1从

映射到

\hat Q

和

\hat K

的复杂度分别等于

NC^2/2

。从

到

的映射和

Proj(·)

都贡献了

NC^2

。当CHH（DWConv 3×3）和IHH（Conv 1×1）贡献

9N^2h/2

和

N^2h^2/4

的FLOP时。因此，总的复杂性是：

作者还可以计算cFFN的复杂度，它由3个映射组成，其通道分别从

到

，

到

和

到

，如下：

从方程5和7中，作者可以看到，与标准MHSA相比，作者的hMHSA的FLOPs减少等于

NC^2+(2C−h^2−18h)N^2/4

。一般来说，C>>h因此可以确保自注意力模块的FLOPs减少。同时，对于FFN模块，作者的cFFN将其FLOP降低到

(1+t)/2

的比率，当

t<1

时，FLOP的降低也不是微不足道的。

3、实验

3.1、SOTA对比

表1给出了ImageNet分类结果，证明了作者提出的方法的有效性。具体而言，作者注意到，尽管与官方实现相比，PVTv2-b1的参数和FLOP分别减少了13.2%和15.2%，但它在作者的方法中表现出了有竞争力的性能。这些发现证实了作者的方法在处理层次结构和空间减少注意力模块方面的稳健性。

此外，配备作者方法的NextViTS实现了相当的准确性，Params和FLOP分别降低了13.9%和11.5%。这一结果突出了作者的方法在从高效模型中消除冗余方面的有效性。此外，DeiT-T和DeiT-S在Acc@1通过引入作者提出的模块，分别提高了0.7%和0.3%，证明了作者的方法在直结构中的有效性。

值得注意的是，DeiT-T和DeiT-S的FLOP和Params节省也相当可观。因此，作者的解决方案表明，作者提出的方法能够在效率-功效权衡方面改进各种基于ViT的Backbone。

3.2、消融实验

1、hMHSA的研究

首先，作者旨在寻找一种合适的策略，从Q-K生成的Head 生成Hallucinate Head 。作者对不同的Hallucinate Head生成方法进行了比较，包括直接复制

A^r_{pre}

、仅应用

IHH(·)

和仅应用

CHH(·)

。

从表2中可以看出，简单的复制会对性能造成0.5%的损害。简单地应用

CHH(·)

或

IHH(·)

。因此，这证明了作者提出的

IHH(·)

和

CHH(·)

的优越性。每个Head 中已经具有高度代表性的Token之间的交互以及不同Head 之间的相关性都可以帮助hMHSA获得竞争性能，因为完全QK计算的注意力图再次验证了作者的注意力图Hallucinate策略是有用的。

然后，作者探索了同时应用IHH和CHH的顺序，即在

CHH(·)

之前应用

IHH(·)

，反之亦然。表3中的结果表明，在

IHH(·)

之后应用

CHH(·)

。这在直觉上是合理的，因为注意力图间信号是不可比较的，所以最好在注意力图内对局部感知场中的注意力关系进行建模，而不是在融合跨注意力图信息后应用3×3 Conv。

2、MHSA v.s.hMHSA

在这个实验中，作者试图实证地测量hMHSA的冗余减少。分块相似性计算的启发，作者提出了称为“CCS”的度量，即“贡献余弦相似性”的缩写，以测量Head 相似性。如在注意力机制中所知，注意力图的行与V的列相乘以对结果做出贡献。因此，作者根据两个不同Head 对应行之间的相似性来量化它们之间的相似程度。首先，注意力Head 之间的相似性计算为：

其中，

A^l_n

为第

个块中的第

个Head，

<·,·>

为两个向量之间的内积。然后，作者将第n个块的Head 相似度

S_n

定义为：

最后，作者对整体B块相似度进行平均，得到贡献余弦相似度：

因此，该指标直接展示了不同Head 之间的相似性。

如表4所示，作者计算了三个模型中注意力图的贡献余弦相似性。结果表明，在所有模型中，作者的hMHSA Value都低于官方的ViT块，这表明作者的hMHSA在同一块中的Head 之间的内部相似性较小。

此外，作者在图4中可视化了由作者的hMHSA模块生成的DeiT-S的注意力图。可以看出，尽管12个Head 中有一半是由另一半生成的，但与标准MHSA相比，实际上它们在这些地图中的相似性较低（图1a）。例如，只有一对图非常相似，而其他图则表现出不同的注意力模式。这表明作者的hMHSA有能力灵活地生成不太相似的注意力图。

3、cFFN的研究

首先，作者探索了FFN压实比简单地减少m的隐藏层尺寸的优越性。为了公平比较，作者将cFFN的紧凑比设置为t=2/3，并设置m，使得香草FFN与cFFN共享相同的Params和FLOP。

从表5中可以看出，简单地减少m会使PVTv2-b1的性能严重降低0.5%。同时，通过作者的因子分解，69.1%Acc@1这比直接减少m高出0.3%。它认为分解隐藏到输出的映射矩阵是更好的设计选择。

此外，作者在cFFN中进一步研究了重参化技术。为了更好地了解重参化技术的效果，作者在训练时间上消融了可重参化的分支的数量r，特别是1、2和3。该表显示，r=2比r=1获得更好的性能，而r=3导致精度下降。此外，有了更多的重参分支，训练成本也更高。因此，作者选择r=2作为cFFN。

然后，作者证明了对m2进行因式分解是一个更好的选择。作者分别对m1和m2进行了因式分解和参数重构。从表6中可以看到压缩M2的性能比M1高0.2%。实验结果表明，在作者的实现中所做的压缩m2是优越的。

最后，控制紧凑比t可以帮助实现效率和性能之间的平衡。作者进行了实验，以在作者的实施中根据经验寻求适当的压实比t，结果如表7所示。作者可以观察到，在Acc@1.当t等于3/4时，Params和FLOP的节省变得非常有限。因此，作者选择2/3的紧凑比t作为作者的最终选择。

3.3、下游任务

为了验证所提出方法的可转移性，作者展示了下游任务的分类性能，结果如表8所示。在本实验中，作者遵循惯例，分别使用384×384和448×448对CIFAR-100和斯坦福犬上的所有模型进行训练和测试。根据ImageNet预训练的权重对模型进行初始化，并在这些数据集上进行微调。从这些结果中，作者可以发现，作者的三个模型的变体与其原始实现是可比较的。

4、局限性与总结

在本文中，作者研究了MHSA中注意力图的相似性（或冗余性），并提出通过使用更简单的操作（称为hMHSA）从另一半产生Hallucinate来利用这种性质。作者的hMHSA机制发现了降低MHSA成本的新方向。此外，作者还注意到了基于ViT架构中经常被忽视的FFN部分，并提出了作者的cFFN策略来对其进行去冗余。

实验结果表明，hMHSA和cFFN可以被纳入各种基于ViT的Backbone中，包括直接、分层和混合卷积以及MHSA结构。这些模型在FLOP和Params方面的复杂性可以进一步降低，同时，作者的架构的整体性能具有相当的竞争力。

然而，出于概念验证的目的，作者只关注用于模型复杂性测量的FLOP和Params。目前没有关注模型吞吐量，因为吞吐量在很大程度上取决于硬件加速技术，不同的实现可能会导致吞吐量的巨大变化。作者将以下内容作为作者未来的工作：要么是具有现代硬件加速库（例如，CUDNN或TensorRT）意识的hMHSA模块设计，要么是作者当前hMHSA模件的定制硬件加速友好实现。