大道至简 | 设计 ViT 到底怎么配置Self-Attention才是最合理的?(一)

本文涉及的产品
函数计算FC,每月15万CU 3个月
简介: 大道至简 | 设计 ViT 到底怎么配置Self-Attention才是最合理的?(一)

Transformer已成为深度学习中的主要架构之一,尤其是作为计算机视觉中卷积神经网络 (CNN) 的强大替代品。然而,由于Self-Attention在长序列表示上的二次复杂性,特别是对于高分辨率密集预测任务,先前工作中的Transformer训练和推理可能非常昂贵。为此,我们提出了一种新颖的少注意力视觉Transformer(LIT),基于Transformers中早期的Self-Attention仍然专注于局部模式,并在最近的分层视觉Transformers中带来较小的好处。

具体来说,提出了一个分层 Transformer,使用纯多层感知器(MLP)在早期阶段对丰富的局部模式进行编码,同时应用Self-Attention模块在更深层捕获更长的依赖关系。此外,进一步提出了一种 Learned Deformable Token Merging Module,以非均匀方式自适应地融合信息Patch

所提出的 LIT 在图像识别任务(包括图像分类、对象检测和实例分割)上取得了可观的性能,可作为许多视觉任务的强大支柱。

开源地址:https://github.com/zhuang-group/LIT


1简介


Transformers在自然语言处理(NLP)和最近在计算机视觉(CV)领域取得了长足的进步。受CNN中金字塔设计的启发,最近的分层视觉转换器(HVT)将Transformer Block划分为多个阶段,并随着网络的深入逐渐缩小特征图。然而,早期的高分辨率特征图导致了很长的token序列,由于Self-Attention的二次复杂度,带来了巨大的计算成本和内存消耗。例如,大小为56×56×96 的特征图在一个Multi-Head Self-Attention(MSA)中需要2.0G FLOPs,而ResNet-18的整个模型只需要1.8G FLOPs。如此巨大的计算成本使得将Transformer应用到广泛的计算机视觉任务中变得非常困难。

HVT的早期阶段,已经做出了一些努力来降低计算成本。例如,一些工作减少了MSA层中Self-Attention Head的数量或进一步减少了Transformer Block的数量。另一行工作建议通过启发式近似来权衡MSA的准确性和效率,例如空间缩减注意力(SRA)和基于Shift Window Multi-Head Self-Attention(SW-MSA)。也有研究当特征图的分辨率相当大时使用卷积层。然而,早期采用Self-Attention层对最终性能有多大贡献仍不清楚。

在本文中提出了一种Less attention Vision Transformer(LIT)来解决上述HVT问题。具体来说,建议在早期Stage使用MLP层来捕获局部信息,同时引入具有足够数量的Head来处理的MSA层后期的长期依赖关系。

作者的动机来自2个方面。首先,先前的CNNTransformer相关研究表明,模型浅层会专注于局部信息,而深层倾向于捕获高级语义或全局关系,这就产生了在早期Stage是否有必要使用Self-Attention的问题。其次,从理论的角度来看,一个具有足够Head的Self-Attention应用于图像可以表达任何卷积层。

然而,MSA层中较少的Head理论上阻碍了逼近具有大kernel-size的卷积层的能力,其中极端情况与1×1卷积一样具有表现力,可以被视为独立应用于每个像素的标准FC层。虽然最近的HVT在早期阶段采用很少的Head来提供金字塔表示,但作者认为这不是最优的,因为这样的配置会引入高计算和内存成本,但带来的收益却很小。

需要强调的是,通过在前期Stage利用MLP Block可以避免Self-Attention在高分辨率特征图上产生的巨大计算成本和内存占用。此外,在后期Stage应用Self-Attention来捕获远程依赖关系。综合实验结果表明,经过如此简单的架构设计后可以带来模型性能和效率之间的最佳平衡点。

此外,最近的HVT要么采用标准卷积层,要么采用线性投影层来合并邻近的Token,旨在控制特征图的规模。然而,考虑到并非每个像素对输出单元的贡献都相同,这种方法阻碍了Vision Transformer对几何变换建模的表示能力。

为此,作者提出了一个受Deformable Convolutions启发的Deformable Token Merging(DTM) 模块,在该模块中,学习了一个偏移网格以自适应地增加空间采样位置,进而合并来自特征图中子窗口的相邻Patch。通过这种方式可以获得更多信息的下采样Token用于后续处理。

主要贡献:

  • 首先,确定了早期MSA层在最近的HVT中的微小贡献,并在早期Stage提出了一个简单的HVT结构,其中包含纯MLP Block;
  • 其次,提出了一个Deformable Token Merging模块,以自适应地合并更多信息Patch以提供分层表示,并具有增强建模能力;
  • 最后,进行了广泛的实验,以表明所提出的LIT在计算复杂度和内存消耗方面的有效性。

2本文方法


2.1 整体架构

image.png图1

LIT的整体架构如图1所示。设为输入的RGB图像,其中H和W分别表示高度和宽度。首先将I拆分为不重叠的PatchPatch大小为4×4,因此每个Patch的初始特征维度为4×4×3=48。

接下来,利用线性嵌入层将每个Patch投影到维度,用作以下过程的初始输入。整个模型分为4个Stage。让s∈[1,2,3,4] 作为Stage的索引,在每个Stage使用个Block,其中前2个Stage仅使用MLP Block来编码局部信息,后两个Stage使用标准Transformer Block来编码处理更长的依赖关系。

在每个Stage,将输入特征图缩放为,其中和分别表示第s个Stage的块大小和隐藏维度。对于最后2个Stage,在每个Transformer Block中设置个Self-Attention heads。

2.2 LIT中的Block设计

如图1所示,LIT采用2种类型的模块:MLP模块和Transformer模块。在前期Stage应用MLP Block。具体来说,MLP Block建立在MLP之上,该MLP由2个FC层组成,中间有GELU。对于第s阶段的每个MLP,使用的扩展比。

具体来说,第1个FC层将token的维度从扩展到,另一个FC层将维度缩减回。形式上,令为第s阶段的输入,l为块的索引,MLP Block可以表示为:

image.png

其中LN表示层归一化。在最后阶段,ViT中描述的Transformer Block包含一个MSA层和一个MLP,可以表示为:

image.png

使用这种架构有2个主要优点:

  • 首先,避免了在早期Stage由长序列引入的巨大计算成本和内存占用;
  • 其次,与最近使用子窗口缩小注意力图或减少注意力图的工作不同在key和value矩阵的空间维度上,在最后2个Stage保留标准MSA层,以保持LIT处理远程依赖关系的能力。

备注

在这里通过考虑卷积层FC层MSA层之间的关系来证明在前期Stage应用纯MLP Block的合理性。

这里建议参考一下字节跳动近期新出的TRT-ViT,有对应的结论给出,同时给出了非常详细的模型设计准则,对应推文的链接如下:

建议背诵 | 字节用4大准则教你设计一个拥有CNN的速度,Transformer精度的模型!


首先,从回顾标准卷积层开始。令为输入特征图,令为卷积权重,其中K为Kernel-size,和分别为输入和输出通道维度。为简单起见,省略了偏置项,并用表示,其中(i,j)表示像素索引,。给定K×K个采样位置的卷积核,像素p的输出可以表示为:

其中是采样索引到预先指定的偏移量ΔK的双射映射。

例如,令ΔK={(−1,−1),(−1,0),...,(0,1),(1,1)}为3×3的kernel,膨胀率为1,则g(0)=(-1,-1)表示第1个采样偏移量。当K=1时,权重张量W等价于一个矩阵,使得$。在这种情况下,Eq(4)可以表示FC层,像素p的输出定义为:

最后,让是MSA层中的Head数,$是第h个Head的可学习参数。在特定的相对位置编码方案下,Cordonnier等人证明了像素p处的MSA层的输出可以表示为:

其中是Head到像素位移的双射映射。在这种情况下,Eq(6)可以看作是Kernel-size为 的卷积层近似。

从Eqs(4)-(6)观察到,虽然具有足够数量的Head的MSA层能够逼近任何卷积层,但理论上更少的Head限制了这种逼近的能力。作为极端情况,具有一个Head的MSA层是只能逼近FC层。

请注意,MSA层在实践中肯定不等同于卷积层。然而,d'Ascoli 等人观察到,早期的MSA层可以在训练时学会卷积的类似表征。考虑最近的HVT在早期Stage采用很少的Head,这种卷积行为可能会限制在小的感受野内。

图 3

在图3中,在可视化中显示PVT-S中的早期MSA层确实只关注query像素周围的一小块区域,而删除它们会导致性能轻微下降,但会显着降低模型复杂性。这证明了在前2个stage应用纯MLP Block是合理的。

2.3 Deformable Token Merging

以前关于 HVT 的工作依赖于Patch Merge来实现金字塔特征表示。然而,从规则网格合并Patch,并忽略了并非每个Patch对输出单元的贡献相同的事实。受可变形卷积的启发,提出了一个Deformable Token Merging模块来学习偏移网格,以自适应地采样更多信息块。形式上,可变形卷积被表述为:

与Eq(4)中的标准卷积运算相比,DC为每个预先指定的偏移量g(k)学习一个偏移量Δg(k)。学习Δg(k) 需要一个单独的卷积层,它也应用于输入特征图X。为了以自适应方式合并Patch,在DTM模块中采用一个DC层,可以表示为:

其中 BN 表示批量归一化,并使用了 GELU

在消融研究中,当物体的尺度和形状发生变化时,DTM 中的采样位置会进行自适应调整,从而受益于学习到的偏移量。另请注意,与Baseline中的常规网格采样相比,轻量级DTM引入了可忽略的FLOP和参数,因此使其成为最近HVT的即插即用模块。

相关实践学习
【文生图】一键部署Stable Diffusion基于函数计算
本实验教你如何在函数计算FC上从零开始部署Stable Diffusion来进行AI绘画创作,开启AIGC盲盒。函数计算提供一定的免费额度供用户使用。本实验答疑钉钉群:29290019867
建立 Serverless 思维
本课程包括: Serverless 应用引擎的概念, 为开发者带来的实际价值, 以及让您了解常见的 Serverless 架构模式
相关文章
|
机器学习/深度学习 编解码 数据可视化
【即插即用】涨点神器AFF:注意力特征融合(已经开源,附论文和源码链接)
【即插即用】涨点神器AFF:注意力特征融合(已经开源,附论文和源码链接)
3062 1
|
机器学习/深度学习 网络架构 计算机视觉
UNet详细解读(一)论文技术要点归纳
UNet详细解读(一)论文技术要点归纳
284 0
|
4月前
|
机器学习/深度学习 算法 网络架构
神经网络架构殊途同归?ICML 2024论文:模型不同,但学习内容相同
【8月更文挑战第3天】《神经语言模型的缩放定律》由OpenAI研究人员完成并在ICML 2024发表。研究揭示了模型性能与大小、数据集及计算资源间的幂律关系,表明增大任一资源均可预测地提升性能。此外,论文指出模型宽度与深度对性能影响较小,较大模型在更多数据上训练能更好泛化,且能高效利用计算资源。研究提供了训练策略建议,对于神经语言模型优化意义重大,但也存在局限性,需进一步探索。论文链接:[https://arxiv.org/abs/2001.08361]。
55 1
|
5月前
|
机器学习/深度学习 计算机视觉
YOLOv10实战:SPPF原创自研 | SPPF_attention,重新设计加入注意力机制 | NEU-DET为案列进行展开
【7月更文挑战第1天】 优点:为了利用不同的池化核尺寸提取特征的方式可以获得更多的特征信息,提高网络的识别精度; 如何优化:在此基础上加入注意力机制,能够在不同尺度上更好的、更多的获取特征信息,从而获取全局视角信息并减轻不同尺度大小所带来的影响; SPPF_attention,重新设计加入注意力机制 ,在NEU-DEU任务中mAP50从0.683提升至0.703;
690 3
|
机器学习/深度学习 自然语言处理
深度剖析Transformer核心思想 "Attention Is All You Need"
深度剖析Transformer核心思想 "Attention Is All You Need"
265 1
|
7月前
|
机器学习/深度学习 编解码 自然语言处理
【VIT】小白入门篇:从各个角度认识Vision Transformer
【VIT】小白入门篇:从各个角度认识Vision Transformer
455 0
【VIT】小白入门篇:从各个角度认识Vision Transformer
|
机器学习/深度学习 算法 大数据
Vision Transformer 必读系列之图像分类综述(三): MLP、ConvMixer 和架构分析(下)
在 Vision Transformer 大行其道碾压万物的同时,也有人在尝试非注意力的 Transformer 架构(如果没有注意力模块,那还能称为 Transformer 吗)。这是一个好的现象,总有人要去开拓新方向。相比 Attention-based 结构,MLP-based 顾名思义就是不需要注意力了,将 Transformer 内部的注意力计算模块简单替换为 MLP 全连接结构,也可以达到同样性能。典型代表是 MLP-Mixer 和后续的 ResMLP。
1171 0
Vision Transformer 必读系列之图像分类综述(三): MLP、ConvMixer 和架构分析(下)
|
机器学习/深度学习 编解码 自然语言处理
Transformer 系列| Transformer又搞事情!TransReID首次在ReID中应用,结果喜人(文末获取论文)(一)
Transformer 系列| Transformer又搞事情!TransReID首次在ReID中应用,结果喜人(文末获取论文)(一)
438 0
Transformer 系列| Transformer又搞事情!TransReID首次在ReID中应用,结果喜人(文末获取论文)(一)
|
网络架构 计算机视觉
UNet++详细解读(一)论文技术要点归纳
UNet++详细解读(一)论文技术要点归纳
364 0
|
数据可视化 计算机视觉
大道至简 | 设计 ViT 到底怎么配置Self-Attention才是最合理的?(二)
大道至简 | 设计 ViT 到底怎么配置Self-Attention才是最合理的?(二)
133 0