通过重叠头自注意力增强视觉 Transformer 的学习能力 !

科技   2024-11-14 09:01   上海  

点击下方卡片,关注「集智书童」公众号

点击加入👉「集智书童」交流群




想要了解更多:

前沿AI视觉感知全栈知识👉「分类、检测、分割、关键点、车道线检测、3D视觉(分割、检测)、多模态、目标跟踪、NerF

行业技术方案👉AI安防、AI医疗、AI自动驾驶
AI模型部署落地实战👉「CUDA、TensorRT、NCNN、OpenVINO、MNN、ONNXRuntime以及地平线框架」

欢迎扫描上方二维码,加入「集智书童-知识星球」,日常分享论文、学习笔记、问题解决方案、部署方案以及全栈式答疑,期待交流!

免责声明
凡本公众号注明“来源:XXX(非集智书童)”的作品,均转载自其它媒体,版权归原作者所有,如有侵权请联系我们删除,谢谢。

近年来,视觉 Transformer 在大多数视觉任务上取得了显著的进步,实现了最先进的性能。这一成功关键在于引入了多头自注意力(MHSA)模块,该模块使得每个头可以通过独立应用注意力机制来学习不同的表示。

在本文中,作者通过在MHSA中重叠头的方式,实证地证明了视觉 Transformer 可以进一步改进。

作者引入了多重叠头自注意力(MOHSA),其中头与其相邻的两个头进行 Query 、 Key和Value 的重叠,而对于第一个和最后一个头,采用零填充进行扩展,因为它们只有一个相邻的头。

作者提出了各种重叠比的模式,以充分研究作者方法的最优性能。

所提出的做法在四个基准数据集上的五个 Transformer 模型上进行了评估,结果表明性能有了显著的提升。源代码将在发表时公开。

I Introduction

自Transformer从语言领域转向视觉领域以来,视觉Transformer已逐渐在许多视觉任务中占据主导地位,如图像识别[2][3][4]和目标检测[5][6]。由于Transformer中的全局注意力机制,当训练数据丰富且训练周期足够长时,视觉Transformer在视觉任务上可以获得更好的性能。尽管视觉Transformer的收敛率通常比卷积对应物慢,但全局信息交流是Transformer模型成功的关键要素之一。

多头自注意力(MHSA)[1]是Transformer模型的核心模块,使得自注意力机制可以在多个头中实现,从而学习不同的表示。自注意力机制最初是为了处理语言数据以捕捉语言中的词间长程关系。当自注意力应用于图像分类和识别任务时,每个图像patch token都可以与所有其他patch token进行交互以获得图像的全局视图,这对视觉Transformer的优秀性能具有重要意义。

MHSA表示 Query 、 Key和Value 被分成不同的头,并在每个头中独立计算自注意力。具有多头自注意力的Transformer模型通常比具有单头自注意力的模型性能更好。尽管MHSA提高了Transformer模型的性能,但不同头之间的关系和交互却鲜有关注。

Transformer模型成功的主要原因是不同 Token 之间有效信息交换的有效性,从而使每个 Token 都能获得上下文信息的全局视角。由于MHSA的优越性能,大多数Transformer模型默认使用MHSA。然而,每个头中的 Query 、 Key和Value 是分开的,没有重叠,当在各个头中计算注意力时,没有信息交换。换句话说,在计算当前头的注意力时,其他头中的信息并不存在。尽管 Token 在注意力之后会通过线性投影进行处理,但信息交换仅限于每个 Token 。

在本文中,作者提出信息交换在视觉 Transformer (Vision Transformers)的注意力计算过程中的作用可以提高其性能。这一目标可以通过在每个头中重叠 Query 、 Key和Value 与相邻头中的 Query 、 Key和Value 来实现。为此,作者提出了一种名为多重重叠头自注意力(MOHSA)的方法,通过重叠头来改进多头自注意力(Multi-Head Self-Attention)机制,使得每个头中的 在计算注意力时,由相邻头的 重叠,如图1所示。通过重叠头,其他头中的信息也可以参与当前头的注意力的计算。由于在将不同头中的 Token 拼接在一起后,重叠会使 Token 的维度增加,因此线性投影会降低 Token 的维度至原始大小。头间信息交流可以为视觉 Transformer 提供更好的性能。此外,作者还设计了许多重叠比例的范式来研究作者提出方法的最佳性能。

作者的主要贡献如下。

  1. 作者提出了多重交叠自注意力(MOHSA)方法,并证明在计算注意力时,将当前头的 Query 、 Key和Value 与相邻头的 Query 、 Key和Value 交叠,可以提高视觉 Transformer 模型的性能。作者在各种数据集上利用多种视觉 Transformer 模型来证明作者提出的方法的有效性。

  2. 针对覆盖度维数的各种变体,本文提出了用于全面调查MOHSA(多头自注意力机制)最优性能的方法。所提出的MOHSA可以集成到大多数Vision Transformer模型中,以提高它们的性能,而无需增加过多开销。

  3. 迄今为止,作者的工作是首次研究在计算注意力时,不同head之间的重叠方法。

II Related Work

Vision Transformers

视觉Transformer[2]利用嵌入的图像块作为 Token 输入,以及一个类 Token 作为 Token 输入,用于Transformer编码器来识别图像。Swin-Transformer[3]通过限制每个窗口中的注意力计算来实现计算成本的降低,并通过移位窗口来扩展 Token 的视图。Shuffle Transformer[10]提出空间混洗来在窗口之间交换信息。

MSG-Transformer[11]利用MSG Token 来表示局部窗口,并使用MSG Token 来实现窗口之间的信息通信。PVT[12]设计了一个具有层次结构的视觉Transformer,该结构不包含卷积,通过逐渐缩小特征图来实现。CaiT[13]和DeepViT[14]研究了具有更深层的视觉Transformer。

移动ViT系列[15][16][17]和EdgeViTs[18]探索了移动 Level 应用中的视觉Transformer。一些工作[19][20][21][22][23]将卷积引入视觉Transformer,以充分利用卷积和Transformer的优势。PiT[24]探索了使用单个像素而不是块作为视觉Transformer Token 的可能性。

Attention Heads Interaction

DeepViT [14] 通过提出一种可学习的矩阵,将所有 Head 的注意力图混合在一起,并在与值相乘之前进行重新注意力。重新注意力类似于最初用于语言任务的说话人注意力[25],也被CaiT[13]所使用。说话人注意力[25]在Multi-Head Attention的 Head 之前和之后使用可学习的线性投影,在softmax函数之前和之后交换注意力模块中 Head 的信息。

上述方法通过线性投影混合注意力图,仅在计算注意力图后才能实现头与头之间的信息交换。然而,混合后的注意力图被应用于没有头间信息交换的值。相比之下,作者提出了一种方法,通过将重叠以实现头间在注意力计算过程中的通信。尽管这种方法引入了略微增加的计算和参数,但它在各种数据集上显著提升了视觉Transformer模型的性能。

III Approach

Multi-Head Self-Attention

[1] Transformer 具有注意力机制来计算 Token 之间的长程关系。注意力计算包括 Query (queries)、键(keys)、值(values)以及 Query 和键的维数 。 Query 和键采用点积计算权重,这些权重用于计算最终结果。注意力计算的矩阵格式如图1所示。在计算点积[1]后,为防止计算 softmax 函数后出现的大值输入,使用了 [1]。


多头自注意力机制[1]在Transformer模型中得到应用,以提高性能,使不同头可以学习不同的表示,这比单头注意力[1]更好。每个头的注意力在方程(2)中得到说明。在每个头中分别被分为,并且针对学习不同方面的独立地在每个头中实现注意力计算。



最后,MHSA可以表示为公式(3)。所有 Head 的结果都被连接起来,是投影矩阵。

视觉Transformer[2]是将Transformer模型应用于视觉任务,主要涉及图像分类和识别。对于视觉Transformer,仅使用编码器进行特征提取。图像被分成相同大小的patch,这些patch被嵌入为视觉Transformer的token。在MHSA和FFN之前,常常使用Layer Normalization[26]。

视觉Transformer[2]可以应用于视觉任务,并在视觉任务上取得显著性能,因为注意力机制的有效信息交换。自注意力机制计算每个token与所有其他token的点积,从而在它们之间有效交换信息。在图像分类和识别等视觉任务中,每个patch token具有图像的全局视图,这对于每个patch获取上下文信息非常重要。

然而,在计算注意力时,每个head都是独立地实现注意力机制的,因此在计算过程中,一个head并没有其他head的注意力信息。

尽管投影矩阵是在所有头都拼接在一起之后实现的,但只有每个 Token 的信息会发生变化。在计算每个头中的注意力时,并没有信息交换。因此,作者提出将相邻头的信息进行重叠,以增强每个头中计算注意力时的信息交换。

Multi-Overlapped-Head Self-Attention

根据上述分析,作者提出了一种简单而有效的策略,通过改进多头自注意力机制来增强视觉Transformer的性能。为了在头之间交换信息,作者利用软除法而不是硬除法将Q、K和V划分到不同的头中。

这个过程可以用公式(4)分别表示为Q、K和V。在公式(4)中,作者使用"part"来表示相邻头之间的部分重叠。


在公式(4)中,的原硬除法结果,这在公式(2)中得到展示。对于软除法,也包含其相邻 Head 的部分信息,从而使得与两个相邻头重叠,这在图2中得到展示。

左图表示对不同头的原始硬除法实现,而右图则代表作者提出的将对不同头的软除法实现。对于只有相邻头一个的第一头和最后头,采用零填充方法构建第一个和最后一个相邻头。


在计算了每个头中重叠的,,和后,结果被连接在一起,如图5-6所示。重叠的头在连接后会略微增加 Token 的维度。因此,投影矩阵会将连接后的维度()投影到原始 Token 维度(),以便 Token 可以以相同的维度进入下一层。在方程(3)中的投影矩阵会将连接后非重叠的头()投影到原始 Token 维度()。由于重叠头的维度略大于非重叠头的维度,头数不变,作者提出的MOHSA中的投影矩阵比原始MHSA中的投影矩阵具有更多的参数。

Overlapping Ratios

重叠比例对于所提出的方法的效用至关重要。在实验中,作者设计了几种重叠比例的范式。在本工作中,作者利用重叠维度来展示重叠比例,如图3所示。从图3中,每个 Head 的重叠维度是其一侧相邻 Head 的重叠维度。在图3中,蓝色部分是 Head 的原始部分,红色部分是与相邻 Head 重叠的部分。对于只有单侧相邻 Head 的第一个和最后一个头,作者利用零填充来处理缺失的相邻 Head 。

除了所有层的固定重叠尺寸外,作者还通过根据层深度改变重叠尺寸来实现一些重叠比的变化。在本论文中,一个层包括注意力模块和FFN模块。重叠比的变化在表1中得到了展示。在表1中,“inc (x层)”表示重叠尺寸每x层增加1,而“dec (x层)”表示重叠尺寸每x层减少1,这是根据层深度进行的“inc (x层)”的反向过程。此外,“0索引”表示重叠尺寸从0开始为“inc”,结束为0为“dec”,而“1索引”表示重叠尺寸从1开始为“inc”,结束为1为“dec”。例如,inc (2层)(在实验中用“inc-0 (2)”表示)在总共12层中,重叠尺寸(0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5)从第1层到第12层。dec (1层)(在实验中用“dec-1 (1)”表示)在总共12层中,重叠尺寸(12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1)从第1层到第12层。

IV Experiments

在实验中,作者选择了三个代表性模型来研究所提出的算法:基础的ViT [2],更深层的Transformer CaiT [13],以及基于窗口的分层Transformer Swin-Transformer [3]。更具体地说,作者选择了ViT-Tiny [2],ViT-Small [2],CaiT-xxs12 [13],CaiT-xxs24 [13]和Swin-Tiny [3]来探索所提出算法的有效性。

各模型的参数设置如表2所示。对于Swin-Tiny,头数和层数根据阶段有所不同。CaiT-xxs12和CaiT-xxs24有两个额外的类层。这些模型将在CIFAR-10 [7],CIFAR-100 [7],Tiny-ImageNet [8]和ImageNet-1k [9]上进行训练和评估。

实验通过100个epoch,其中warmup epoch为20个。优化器采用AdamW [27]。在CIFAR-10 [7]、CIFAR-100 [7]、Tiny-ImageNet [8]上的实验在4个NVIDIA P100 GPU上运行,总batch size为128。在ImageNet-1k [9]上的实验在4个NVIDIA V100 GPU上进行,总batch size为512。ViT和CaiT的初始学习率分别为0.0005。

对于Swin-Tiny,Tiny-ImageNet的初始学习率为0.00025,ImageNet的初始学习率为0.0005。实验中的图像均被重新缩放到224。其他设置遵循Swin-Transformer [3]中的设置。

Cifar-10

表3展示了ViT-Tiny和ViT-Small在CIFAR-10上的消融研究。在表3中,"固定1"表示所有层的叠叠维数为1,"固定一半"表示所有层的叠叠维数为头维数的一半。 "inc"表示随着层深的增加,叠叠维数增加,"dec"表示随着层深的增加,叠叠维数减少。例如, "inc-0 (3)"说明每3层叠叠维数增加1,且以0为索引。

对于ViT-Tiny,在变体"inc-0 (6)"下,最佳准确率为85.65;对于ViT-Small,在变体"inc-0 (1)"下,最佳准确率为87.30。对于ViT-Tiny,"固定一半"的准确率高于"固定1"。对于ViT-Small,"固定一半"的准确率低于"固定1"。在固定重叠维度范式下,重叠比例的更多并不能保证更高的准确率。

尽管重叠头会略微增加计算和参数,但从实验结果中作者可以看到,参数和FLOPs的增加微不足道。表4展示了CaiT-xxs12在CIFAR-10上的实验结果。最佳结果80.04显著提升了CaiT-xxs12的性能3%,且开销微不足道。

Cifar-100

表5展示了ViT-Tiny和ViT-Small在CIFAR-100上的实验结果。"inc-0 (2)"使得ViT-Tiny取得了63.01的最佳结果,而ViT-Small的最佳结果是64.97,对应的模型是"dec-1 (3)"。这两个结果都比原始模型提高了超过1%。与在CIFAR-10上的结果类似,"固定一半"对于ViT-Tiny的准确率更高,而"固定1"对于ViT-Small的表现更好。此外,"固定一半"对于ViT-Small并不具有有效性。

表7展示了CaiT在CIFAR-100上的性能。对于CaiT-xxs12,使用所有层都为1重叠维数可以提高准确率近2%。对于CaiT-xxs24,使用"inc-1 (3)"可以显著提高5%的准确率。此外,"inc-1 (3)"比使用 Head 一半维数作为重叠维数更好,因为它利用了更多的参数和FLOPs。

Tiny-ImageNet

表7展示了在Tiny-ImageNet上的实验结果。利用"inc-1 (2)",ViT-Tiny的准确率可以提高1.61%。ViT-Small则通过使用"inc-1 (1)"提高了1.27%。此外,将头维度的一半作为重叠维度比将1作为所有层的重叠维度更适合ViT-Tiny,而对于ViT-Small来说并非如此。

对于CaiT-xxs12和CaiT-xxs24,使用固定重叠维度1进行所有层设置的"固定1"可以显著提高准确性,分别提高了2.39%和3.58%。此外,使用"inc-1 (1)"的CaiT-xxs24的准确性提高了7.41%。作者方法在CaiT模型上的有效性显著,可以忽略略微增加的参数和计算。

对于Swin-Tiny,固定1范式能够将准确性提高1.23%,而"dec-1 (1)"能够显著提升性能。重叠头的有效性也在基于窗口的分层Transformer模型中得到证实。

ImageNet

通过在ImageNet上进行训练和测试,作者可以研究在大型数据集上作者提出的方法的 effectiveness。ImageNet上的实验结果如下表所示。几乎所有配备作者提出的MOHSA的模型在ImageNet上都有显著的改进。

ViT-Tiny可以通过使用"fixed half"方法提高1%以上,该方法将所有层的重叠维度设置为头维度的一半。ViT-Small在"dec-0 (1)"下提高了0.79%,且开销可以忽略不计。此外,在ImageNet上,将所有层的重叠维度设置为1并不是ViT-Tiny的有效方法。然而,对于ViT-Small,使用1作为所有层的重叠维度比使用所有层头维度的一半作为重叠维度要好。

作者的方法在CaiT模型上取得了更大的性能提升。CaiT-xxs12在采用"inc-1 (1)"后提高了1.17%,而CaiT-xxs24在仅使用1作为所有层的重叠维度时,其准确率显著提高了3.70%,这证明了作者的方法在显著提高模型性能的同时,保持了最小的开销。对于CaiT-xxs24,仅使用1作为所有层的重叠维度,准确率就提高了3%以上,这充分展示了作者提出的方法的有效性和效率。

对于Swin-Tiny,仅仅将1作为重叠维度在ImageNet上并不有效。尽管将重叠维度增加到头维度的一半可以提高Swin-Tiny的准确性,但这种方法与通过参数和计算更少的层深度来变化重叠维度的范式几乎得到相同的结果。

表9说明了在将重叠范式"inc-0 (1)"应用于Q、K或V时,使用ViT-Tiny和CaiT-xss12在ImageNet上的消融研究。从表9的实验结果可以看出,将作者的方法应用于V对提升性能的影响更大,而将重叠范式应用于Q、K和V可能产生最佳性能。

Analysis

在实验中,作者利用作者提出的MOHSA的不同变体在各种Vision Transformer模型和数据集上进行,以说明MOHSA的有效性。总体而言,这些模型在各种数据集上的改进显著,且开销很小。对于作者提出的这种方法的不同变体,结果表明有一些差异。对于固定范式,所有层的交集维度相同,在大多数情况下,通过增加1个交集维度即可提高性能,固定模式下的交集维度增加不能保证获得更好的结果。此外,根据层深度的不同交集维度变体的性能在大多数情况下优于固定范式在各种模型和数据集上的表现。与具有高交集维度的固定范式相比,根据层深度变化交集维度的方法可以节省参数数量和计算成本。

图4显示了在原始模型与作者的模型之间训练过程中,CaiT-xss24在验证集或测试集上的准确性。对于CIFAR-100、Tiny-ImageNet和ImageNet,分别选取"inc-1 (3)"、"inc-1 (1)"和"inc-0 (1)"作为作者的模型进行比较。在训练初期,原始模型与作者的模型之间的准确性几乎无差异,但随着训练的进行,差距显著增大。

此外,针对各种数据集,CaiT模型带来了主要的改进。这可能是因为CaiT [13]在softmax函数之前和之后都使用了talking-heads注意力 [25]。配备了作者的MOHSA,CaiT可能在不同头之间进行更有效的信息交换,从而展现出优越的性能。

V Conclusion

在这篇论文中,作者提出了一种简单而有效的模块MOHSA,通过重叠头部来改进Vision Transformers中的原始MHSA,这允许在每个头部进行注意力计算时进行信息交换。

据作者所知,这是首次提出重叠头部的工作,并在不同的 Vision Transformer 模型在各种数据集上实现了显著的性能提升。

作者希望这项工作能够激励社区进一步探索Vision Transformers的结构。

参考文献

[0]. Improving Vision Transformers by Overlapping Heads in Multi-Head Self-Attention.

扫码加入👉「集智书童」交流群

(备注:方向+学校/公司+昵称



点击下方“阅读原文”,
了解更多AI学习路上的「武功秘籍」


集智书童
书童带你领略视觉前沿之美,精选科研前沿、工业实用的知识供你我进步与学习!
 最新文章