Press "Enter" to skip to content

ICCV2021 MIT-IBM沃森开源CrossViT:Transformer走向多分支、多尺度

今日分享 ICCV 2021 论文 『CrossViT: Cross-Attention Multi-Scale Vision Transformer for Image Classification』 , MIT-IBM 沃森人工智能实验室开源《CrossViT》,Transformer 开始走向多分支、多尺度(附目前多尺度ViT异同点对比)。

 

详细信息如下:

 

 

论文链接:https://arxiv.org/abs/2103.14899

 

项目链接:https://github.com/IBM/CrossViT

 

导言:

 

与卷积神经网络相比,最近出现的视觉Transformer(ViT)在图像分类方面取得了很好的结果。受此启发,在本文中,作者研究了 如何学习Transformer模型中的多尺度特征表示来进行图像分类 。 为此,作者提出了一种双分支Transformer来组合不同大小的图像patch,以产生更强的图像特征。本文的方法用两个不同计算复杂度的独立分支来处理小patch的token和大patch的token,然后这些token通过attention机制进行多次的交互以更好的融合信息。

 

此外,为了减少计算量,作者开发了一个简单而有效的基于cross-attention的token融合模块。在每一个分支中,它使用单个token(即 [CLS] token)作为query,与其他分支交换信息。 本文提出cross-attention的计算复杂度和显存消耗与输入特征大小呈线性关系 。 实验结果表明,本文提出的CrossViT的性能优于其他基于Transformer和CNN的模型。例如,在ImageNet-1K数据集上,CrossViT比DeiT的准确率高了2%,但是FLOPs和模型参数增加的非常有限。

 

       01       

 

Motivation

 

Transformer使NLP任务中序列到序列建模的能力取得了很大的飞跃。Transformer在NLP中的巨大成功激发了其在计算机视觉领域的应用。在ViT之前的一些工作主要将Transformer中的Self-Attention和CNN进行结合。虽然这些结合CNN和Self-Attention方法达到了比较不错的性能,但与纯粹的基于Self-Attention的Transformer相比,它们在计算方面的可拓展性非常有限。

 

ViT使用一系列embedding式的图像patch作为标准Transformer的输入,这是第一个与CNN模型性能相当的无卷积Transformer网络。然而,ViT需要非常大数据集,如ImageNet21K和JFT300M来进行预训练。之后的DeiT表明,数据增强和模型正则化可以在较少的数据下训练高性能的ViT模型。在此之后,ViT就逐渐成为了CV任务中的主流模型之一。

 

在这项工作中,作者研究了 如何学习Transformer模型中的多尺度特征表示来进行图像识别 。 多尺度的特征已经在很多工作中证明了对于CV任务是有效的,但多尺度特征对视觉Transformer的潜在好处仍有待验证。受到多分支CNN架构的启发, 作者提出了一个双分支Transformer来组合不同大小的图像patch,以产生更强的视觉特征用于图像分类 。

 

本文的方法用两个具有不同计算复杂度的独立分支来分别处理大patch的token和小patch的token,这些token多次融合以相互补充信息。 本文的重点是研究并设计适合视觉Transformer的多尺度特征融合方法 。 在本文中,作者通过一个有效的交叉注意模块来实现这一点,其中每个Transformer分支创建一个non-patch token(即 [CLS] token)作为代理,并通过attention机制与另一个分支交换信息。( 由于这里只使用[CLS] token进行信息交互,所以这一步attention的计算复杂度是线性的,而非二次的。 )

 

 

在ImageNet-1K数据集上,CrossViT比DeiT的准确率高了2%,但是FLOPs和模型参数增加的非常有限(如上图所示)。

 

       02       

 

方法

 

2.1.Vision Transformer的概述

 

视觉Transformer(ViT)首先通过将图像按照一定的patch大小划分,然后将图像转换为一个patch的序列,并将每个patch线性投影成embedding。为了执行分类任务,还要在序列中添加了一个额外的分类token(CLS)。此外,由于Transformer编码器中的自注意是与位置无关的,而视觉应用高度需要位置信息,因此ViT向每个token添加了位置embedding,包括CLS token。

 

然后,将所有token输入到堆叠的Transformer编码器,最后使用CLS token进行分类。Transformer编码器由一系列块组成,其中每个块包含带有FNN和多头自注意(MSA)。FFN在隐藏层上包含具有扩展比为r的两层多层感知器,在第一个线性层后应用一个GELU非线性激活函数。在每个块之前应用Layer normalization(LN),在每个块之后应用残差连接。ViT的输入和第k个块的处理可以表示为:

 

 

其中,和分别为CLS和patch token,为位置embedding。N和C分别是patch token的数量和embedding的维度。

 

值得注意的是,ViT与CNN的一个非常不同的设计是CLS token。在CNN中,最终的embedding通常是通过对所有空间位置的特征进行平均,而ViT使用与每个Transformer编码器上的patch token相互作用的CLS token作为最终的embedding。因此,我们可以认为CLS token是一个总结所有patch token的代理,因此作者提出的模块是基于CLS token设计的双路径多尺度的ViT。

 

2.2.多尺度Vision Transformer

 

patch大小会影响ViT的准确性和复杂性;使用细粒度的patch大小,ViT性能更好,但会导致更高的FLOPs和显存消耗。例如,patch大小为16的ViT比patch大小为32的ViT好6%,但前者需要更多4×的FLOPs。受此启发,作者提出的方法是利用来自更细粒度的patch大小的优点,同时平衡复杂性。更具体地说,作者首先引入了一个双分支ViT,其中每个分支处理不同patch大小的特征,然后采用一个简单而有效的模块来融合分支之间的信息。

 

 

上图展示了交叉注意多尺度视觉Transformer( Cross-Attention Multi-Scale Vision Transformer,CrossViT)的网络结构。模型主要由K个多尺度Transformer编码器组成,其中每个编码器由两个分支组成:

 

 

L-Branch :大分支利用粗粒度的patch大小(),这个分支有更多的Transformer编码器和更大的embedding维度。

 

S-Branch :小分支对细粒度的patch大小( )进行操作,这个分支具有更少的编码器和更小的embedding维度。

 

 

两个分支的输出特征在Cross-Attention中融合L次,利用末端的两个分支对CLS token进行预测。对于两个分支的每个token,作者还在多尺度Transformer编码器之前添加了一个可学习的位置embedding,以学习位置信息。

 

2.3.多尺度特征融合

 

 

有效的特征融合是学习多尺度特征表示的关键。在本文中,作者探索了四种策略,如上图所示。

 

设是分支i上的token序列(包括patch和CLS token),其中i可以是l或者s,分别代表大分支和小分支。和分别代表分支i的CLS和patch token。

 

All-Attention Fusion

 

 

All-Attention Fusion是简单地concat来自两个分支的所有token,而不考虑每个token的属性,然后通过自注意模块融合信息,如上图所示。但是这种方法的时间复杂度也是跟输入特征大小呈二次关系,因为所有的token都是通过自注意模块计算的。All-Attention Fusion方案的输出可以表示为:

 

 

其中,和是用来对齐特征通道维度的投影函数。

 

Class Token Fusion

 

 

CLS token可以看作是一个分支的抽象全局特征表示,因为在ViT中,它可以作为用来预测结果的最终embedding。因此,可以直接对两个分支的CLS token求和,如上图所示。这种方法计算上非常有效,因为只需要处理一个token。这个融合模块的输出可以表示为:

 

 

Pairwise Fusion

 

 

上图显示了这两个分支是如何成对融合的。基于patch token位于图像的空间位置,启发式融合方法是根据它们的空间位置将它们合并。然而,这两个分支处理不同大小的patch,因此具有不同数量的patch token。因此,作者首先执行插值操作来对齐空间大小,然后以成对的方式融合两个分支的patch token。另一方面,将这两个分支的CLS token分别融合。L-Branch和S-Branch的成对融合的输出可以表示为:

 

 

Cross-Attention Fusion

 

 

上图显示了本文提出的Cross-Attention Fusion的基本思想,其中融合涉及到一个分支的CLS token和另一个分支的patch token。此外,为了更有效地融合多尺度特征,作者首先利用每个分支的CLS token作为代理,在另一个分支的patch token之间交换信息,然后将其投影到自己的分支。

 

由于CLS token已经学习了自己分支中所有patch token之间的抽象信息,因此与另一个分支中的patch token的交互有助于融合不同尺度的信息。与其他分支token融合后,CLS token在下一层Transformer编码器上再次与自己的patch token交互,在这一步中,它又能够将来自另一个分支的学习信息传递给自己的patch token,以丰富每个patch token的表示。

 

 

大分支(L-Branch)的 cross-attention模块结构如上图所示。具体来说,对于分支l,它首先从s分支中收集patch token,并将自己的CLS token与l分支的patch token进行concat,表示如下:

 

 

其中,是用于特征对齐的投影矩阵。然后,该模块在和之间执行交叉注意(CA),其中CLS token是唯一的query,因为patch token的信息已经被融合到了CLS token中。在数学上,CA可以表示为:

 

 

由于在query中只使用CLS token,所以在交叉注意(CA)中生成attention map的计算复杂度和显存消耗是线性的,而不是像在full-attention中那样是二次的,这使得整个过程更高效。此外,就像Self-Attention一样,作者也在CA中使用多个head,并将其表示为MCA。在交叉注意后,没有使用前馈网络FFN。具体来说,CA的输出的计算如下:

 

 

与其他三种简单的启发式方法相比,交叉注意获得了最好的精度,同时对多尺度特征融合有效。

 

2.4.  CrossViT

 

 

基于上面提出的结构,作者提出了多个不同计算量和参数量的CrossViT模型,具体结构如上表所示。

 

       03       

 

实验

 

3.1. Main Results

 

Comparisons with DeiT

 

 

上表比较了CrossViT和相似计算量下的DeiT进行的比较,可以看出,相比于DeiT,CrossViT具有明显的性能优势。

 

Comparisons with SOTA Transformers

 

 

上表展示了CrossViT和SOTA模型的对比结果,与ViT-B相比,CrossViT-18†的准确率高了4.9%(77.9% vs 82.8%)。

 

Comparisons with CNN-based Models

 

 

上表比较了CrossViT和CNN模型的性能。可以看出,除了EfficientNet,CrossViT相比于大多数CNN结构还是具有性能上的优势的。

 

Transfer Learning

 

 

为了验证本文方法的泛化性能,作者在不同的数据集上做了实验,可以看出CrossViT在其他数据集上同样具有性能的优势。

 

3.2. Ablation Studies

 

Comparison of Different Fusion Schemes

 

 

上表显示了不同fusion方法的性能。可以看出Cross-Attention相比于其他方法具有性能上的优势。

 

Effect of Patch Sizes

 

 

与Baseline相比,上表中的A展示了Patch Size的影响,可以看出使用(12,16)的patch大小能够在更少的FLOPs下能够获得了更好的准确性。

 

Channel Width and Depth in S-branch

 

 

尽管cross-attention设计的比较轻量级,作者在实验中还是测试了使用更复杂的s分支的性能,如上表(B和C)所示。这两种模型都增加了FLOPs和参数,但没有提高精度。

 

Depth of Cross-Attention and Number of Multi-Scale Transformer Encoders

 

 

为了增加两个分支的融合频率,可以堆叠更多的交叉注意模块(L),也可以堆叠更多的多尺度Transformer编码器(K)。如上表D和E所示,可以看出过频繁融合并不能提供任何性能的改进。

 

       04       

 

总结

 

在本文中,作者提出了一种用于学习多尺度特征的双分支视觉Transformer——CrossViT,以提高图像分类的识别精度。为了有效地结合不同尺度的图像patch token,作者进一步开发了一种基于交叉注意的融合方法,能够在线性时间复杂度下,进行两个分支的有效信息交换。通过大量的实验,作者证明了CrossViT的性能优于目前的CNN和Transformer模型。

 

       05       

 

思考

 

在Vision Transformer领域中,除了这篇CrossViT,近期还有另一篇论文CrossFormer( CrossFormer: A Versatile Vision Transformer Based on Cross-scale Attention )同样研究了Vision Transformer中的多尺度特征。下面我们来比较一下这两篇论文的异同点。

 

CrossFormer

 

 

 

CrossViT

 

 

两篇文章的模型结构如上图所示,相同的是,这两篇文章 都用了视觉特征的多尺度信息 ,来获得更加丰富和鲁棒的视觉特征,从而提升视觉任务的性能。

 

不同之处主要存在以下几点:

 

 

CrossFormer是一个多阶段结构,不同的阶段具有不同的分辨率(这一点和CNN很像,空间维度不断缩小,通道维度不断增加);而CrossViT与Transformer更相似,在所有的Block中都保持了相同的分辨率和通道维度。

 

CrossFormer的多尺度特征体现在每个Stage内,它的多尺度是由于CrossFormer中每个stage的CEL模块在进行patch embedding的时候进行了多尺度embedding造成的;CrossViT的多尺度特征体现在Transformer Block之前的patch embedding,它在patch embedding的时候生成了两个尺度的特征,然后CrossViT中每个Block在比较小的计算量下,进行了这两个尺度特征的信息交互。

 

CrossFormer的多尺度信息会更加丰富一些(第1个stage有4个尺度的特征信息;第2/3/4个stage都有2个尺度的特征信息);CrossViT在整个计算过程中存在2个尺度的视觉信息(即文中的S-Branch和L-Branch的尺度)

 

 

作者介绍

 

研究领域:FightingCV公众号运营者,研究方向为多模态内容理解,专注于解决视觉模态和语言模态相结合的任务,促进Vision-Language模型的实地应用。

 

Be First to Comment

发表回复

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