Press "Enter" to skip to content

想要更好的Transformer架构?ViTAS搜索机制助你炼出高性能模型

 

Transformer通过将图像切分成片元序列并利用自注意力机制对图像建模,在众多视觉任务上取得了令人瞩目的成果。但目前架构主要沿袭了NLP领域的原始Transformer架构,视觉领域的 Transformer 还没有被充分的探索。本文提出的 ViTAS搜索机制 ,正是基于这一问题,更高效地利用并设计视觉Transformer模型。

 

为了更高效地设计视觉 Transformer 模型架构, 来自悉尼大学、商汤等机构的研究人员提出了视觉 Transformer 架构自动搜索器ViTAS ,不同token嵌入、序列大小、注意力头数量和宽高的架构都可以从构建的超级 Transformer 中导出。同时为了适应不同架构的变化,还在超级 Transformer 中引入了私有类别token和自注意力图。实验表明搜索出的 Transformer 架构可以达到甚至超过现有纯粹基于Transformer 架构的模型效果,在 ImageNet 实现了74.7%的Top精度。

 

一、 视觉Transformer架构及其搜索空间

 

视觉 Transformer 中,输入图像首先被切分为片元并进行编码,而后送入多个序列堆叠的 Transformer 模块中。每个模块包含了多头自注意力模块中进行处理,紧随其后是两层残差连接的多层感知机。每个 Transformer 模块可以表示为:

 

 

其中LN代表LayerNorm,MHSA代表多头自注意力。在多头自主力模块中,输入首先通过线性映射获取对应的query、key和value向量,query和key可用于度量特征间的注意力,并用于聚合value。在多头注意力机制中,query、key和value向量被分为h组,而后通过全连接层映射获取最终的多头自注意力输出结果:

 

 

其中Z代表输入线性映射结果LN(X),SA代表了第i个注意力头的自注意力。虽然Transformer模型在视觉领域取得了极大的进展,但简单的模块堆叠却不一定是最优的架构。不同的深度的模块也许需要不同层级的特征抽取,因此需要在架构层级进行统一地设计和优化。

 

为了探索ViT的架构空间,研究人员将 Transformer 结构中的所有关键变量提取出来构建了庞大的架构搜索空间 ,包括自注意力头的数量、片元大小、每一层的输出维度、架构深度等。

 

 

视觉Transformer的搜索空间,图像片元首先被线性映射层(橘色)编码,而后输入到级联的Transformer模块中(灰色)进行特征抽取。每个模块中特征被一次送入多头自注意力模块和残差链接的多层感知机模块,每个都包含了两个全连接层。因此图像片元大小、线性映射输出维度、蓝色层的输出维度、自注意力头数量都可以进行搜索优化。在模型架构层面,每个模块进行独立的参数搜索,ViT的深度则在架构层面进行选择。

 

下表中呈现了架构层面对ViTAS搜索空间的定义。其中最大深度设置为16、图元大小设置为{14, 16, 32},自注意力头大小则设置为了{3, 6, 12, 16}。基础版本的超级transformer简单地共享了class token和FC1层。Class Token的共享机制是利用了预先定义了最大的并根据需要的尺寸动态选择输出尺寸[36,30]。同时多头注意力层中的FC1层共享则将输出特征均分到h个组(头)上去,不同的h将会使得每组具有不同的维度D/h。此外, 为了适应输出维度, 每一层维度可以从一组设置中进行选择。 在给定搜索空间后,每个模块都会自由选择自身参数来最优化输出。

 

 

二、视觉 Transformer 架构搜索

 

在神经网络NAS的启发下研究人员首先在提出的搜索空间优化 Transformer 架构。由于资源限制无法对搜索空间中的每一种设置进行训练,为了构建简化版本的超级Transformer,NAS中的one-shot权重共享策略被借鉴过来。Transformer定义了搜索空间A和模型权重W的超级 Transformer。 优化过程主要分为以下两步:

 

 

其中 loss 代表了训练损失、Acc代表了验证集上的精度。W* 代表了一组训练后的模型权重,α 和 α* 代表架构和最优架构。首先从架构空间 A 中均匀采样架构 α 进行训练,而后通过验证集上的精度测评搜索出最优的架构α*。

 

为了进一步提高搜索架构的性能,研究人员提出了一种用于视觉 Transformer 的新型权重共享定制范式。此外为了提升 Super-Transformer 的搜索,还进一步探索了超级 Transformer 在不同架构下稳定的训练、更好地进行性能评测。

 

首先研究人员引入了私有的 class token 和自注意力图来实现权值共享。ViT 中原始的 class token与输入特征图合并同时优化,这种方式对于单个网络的训练非常有效。但对于用于不同 patch 尺寸的不同架构训练时其中包含了不同的语义信息,会影响patch size 的参数搜索。为了解决这一问题,为每个对应的 patch size 分配了对应 class token,保证了在 patch size 变换时私有的 class token 可有效编码特征。同时为了促进参数优化,每个 class token 需要进行充分地更新。因此每个片元大小只有一个class token。私有 class token 的加入量模型的重训练Top-1精度从55.0%提升到了56.6%。

 

另一方面,ViT 中的 Transformer 模块使用多头自注意力来适应片元 token 间的全局互相关信息,由于输出特征被评价分配到了不同的注意力头上D/h,在不同头数目下直接共享 query 和 key 会影响局域性的表达。因此为每个注意力头分配了私有的 query 和 key 用于自相关定位。因此多头自注意力的公式可以别重新改写为:

 

 

其中每个头都拥有自己私有的query和key。

 

 

图中显示了ViTAS中私有的class token和自注意力图。在私有class token中,不同的片元尺寸会被分配独立的class token. 而在私有自注意力图中,value会被所有的头共享,而每个头则具有自己私有的q和k。

 

对于超级 Transformer 的训练是为了公平地评定每个架构的性能,这也决定了下一搜索阶段的优化架构。因此在训练中需要高效、便捷的策略实现稳定训练。为了克服先前架构训练中存在的问题, 研究人员提出了以下三种策略来稳定训练过程:弱增强、正则化与同一性概率搜索 ,感兴趣的小伙伴可以在论文4.3中找到更详细的说明。下表中展示了研究人员将数据增强手段逐渐移除,发现只有CutMix和Erasing方法可以提升重训练精度,搜索到更好的架构。移除所有增强策略后可以改善搜索精度,但重训练精度有所下降。

 

 

三、实验结果

 

通过 ViTAS 搜索得出的模型在 ImageNet 的 Top-1 精度上超过了大多数现有的纯 Transformer 架构,其参数量和效率都有明显的提升:

 

 

自动架构搜索出的模型在其他基准数据集上也拥有十分优异的迁移性能:

 

 

同时在ViT架构进行搜索的过程中,研究人员发现的四条重要的规律: 对于小架构和有限的计算开销,模块越少越好;虽然自注意力头越多速度越慢,但ViT精度会不断提高;ViT的第一个模块中FC层的输出维度趋小;每种架构的最后一个模块都保持了几乎所有维度来适应任务需求。

 

 

如果想要了解更多细节信息,请参看论文和同时开源的代码:

 

https://github.com/xiusu/ViTAS

 

From: arxiv  编译: T.R.

 

Illustrastion from Ico ns8

 

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注