UniFormer: Unifying Convolution and Self-attention for Visual Recognition

paper链接: https://arxiv.org/abs/2201.09450

一、引言

虽然CNN可以在小邻域内通过卷积有效地减少局部冗余，但有限的感受野使其难以捕获全局依赖。另外，VIT可以通过自注意力有效地捕获长期依赖关系，而所有令牌之间的盲目相似性比较导致高冗余。为了解决这些问题，本文提出了一种新的统一transFormer(UniFormer)，它可以在一个简洁的transFormer格式中无缝地集成卷积和自注意的优点。与典型的transFormer块不同，UniFormer块中的关系聚合器在浅层和深层分别配备了局部和全局令牌亲和性，允许处理冗余和依赖，以实现高效和有效的表示学习。UniFormer块堆叠成一个新的强大的骨干，并将其用于从图像到视频领域，从分类到密集预测的各种视觉任务。
本文以图像和视频领域中著名的vit(即DeiT和TimeSformer)为例，进行可视化，它们的注意力分布在浅层。如下图所示，两种vit确实捕捉到了浅层的详细视觉特征，而空间和时间注意是多余的。可以很容易地看到，给定一个锚标记，空间注意力主要集中在局部区域的标记上(主要是3×3)，而从该图中的其他标记中学到的很少。同样，时间注意力主要集中在相邻帧中的标记，而忽略了远处帧中的其他标记。然而，这种局部焦点是通过在空间和时间上对所有标记进行全局比较获得的。显然，这种冗余注意方式带来了大量不必要的计算负担，从而恶化了ViTs中的计算精度平衡。

在此基础上，提出了一种新的统一transFormer(UniFormer)。它以简洁的transFormer格式灵活地统一了卷积和自注意，可以同时解决局部冗余和全局依赖，从而实现有效和高效的视觉识别。具体来说，UniFormer模块由三个关键模块组成，即动态位置嵌入(DPE)、多头关系聚合器(MHRA)和前馈网络(FFN)。关系聚合器的独特设计是UniFormer与以前的CNN和VIT的关键区别。在浅层中，关系聚合器用一个小的可学习参数矩阵捕获局部令牌相关性，它继承了卷积风格，可以通过局部区域的上下文聚合在很大程度上减少计算冗余。
在深层，关系聚合器通过令牌相似度比较学习全局令牌亲和力，它继承了自注意风格，可以自适应地从遥远的区域或框架构建远程依赖。通过逐步分层地堆叠局部和全局的UniFormer块，可以灵活地整合它们的协作能力，以促进表示学习。最后，为视觉识别提供了一个通用而强大的主干，并通过简单而精细的适应成功地解决了各种下游视觉任务。

二、实现细节

下图显示了简洁的统一transFormer(UniFormer)。为了简单描述，以T帧的视频为例，图像输入可以看作是一个单帧的视频。
因此，红色突出显示的维度仅用于视频输入，而对于图像输入，所有维度都等于1。可以看到，UniFormer是一个基本的transFormer格式。

具体来说，UniFormer块由三个关键模块组成:动态位置嵌入(DPE)、多头关系聚合器(MHRA)和前馈网络(FFN):

考虑输入令牌张量$X_{in}\in R^{C\times T\times H\times W}$(对于图像输入T = 1)，首先引入DPE，将位置信息动态地集成到所有令牌中(Eq. 1)。它对任意输入分辨率都很友好，并且很好地利用了令牌顺序，从而更好地进行视觉识别。然后，使用MHRA通过关系学习利用其上下文标记来增强每个标记(Eq。
2).通过在表层和深层灵活设计令牌亲和力，MHRA可以巧妙地统一卷积和自注意力，减少局部冗余，学习全局依赖。最后，像传统的ViTs一样加入FFN，它由两个线性层和一个非线性函数，即GELU (Eq. 3)组成。通道数先按4的比例扩展，然后恢复，这样每个令牌都会单独增强。
如前所述，传统的CNN和VIT侧重于解决局部冗余或全局依赖，导致精度不理想或/和不必要的计算。为了克服这些困难，引入了一种通用的关系聚合器(RA)，它统一了卷积和用于标记关系学习的自注意。通过在浅层和深层分别设计局部和全局token亲和性，实现了高效的表示学习。具体来说，MHRA以多头风格利用令牌关系:

给定输入张量$X\in R^{C\times T\times H\times W}$(对于图像输入T = 1)，首先将其重塑为一个长度为L=T ×H×W的符号序列$X\in R^{L\times C}$。Rn(·)表示第N个头中的RA, $U\in R^{C\times C}$是一个可学习的参数矩阵，可以对N个头进行积分。每个RA由令牌上下文编码和令牌关联学习组成。应用线性变换将原始标记编码为上下文标记$V_n(X)\in R^{L\times \frac{C}{N}}$。
随后，RA可以在令牌亲缘性$A_n\in R^{L\times L}$的指导下总结上下文。

(一)、Local MHRA

尽管浅层的VIT比较了所有令牌之间的相似性，但它们最终学习了局部表示。
这种冗余的自注意设计在浅层中带来了较大的计算成本。基于这一观察，建议在一个小的邻域中学习令牌亲和力，这恰好与卷积滤波器的设计有相似的见解。因此，建议将局部亲和度表示为浅层中的可学习参数矩阵。

其中$a_n\in R^{t\times h\times w}$是可学习参数，$X_j$指${\Omega }_i^{t\times h\times w}$中的任意相邻的令牌。(i−j)表示标记i和j之间的相对位置。请注意，相邻标记之间的视觉内容在浅层中变化微妙，因为标记的感受野很小。在这种情况下，没有必要在这些层中使令牌关联成为动态的。因此，使用一个可学习的参数矩阵来描述局部令牌相关性，它仅仅依赖于令牌之间的相对位置。
与卷积块的比较：局部MHRA可以被解释为MobileNet块通用扩展。首先，方程4中的线性变换V(·)等价于逐点卷积(PWConv)，其中每个头对应一个输出特征通道$V_n(X)$。
此外，局部令牌关联可以实例化为在每个输出通道(或头部)$V_n(X)$上操作的参数矩阵，因此关系聚合器$R_n(X)=A_n^{local}{V}_n(X)$可以解释为深度卷积(DWConv)。最后，线性矩阵U，它连接并融合了所有的头部，也可以被视为一个点卷积。因此，这种局部MHRA可以在MobileNet块中以pwconv-dwconv - pwconv的方式重新表述。此外，与MobileNet块不同的是，局部UniFormer块被设计为通用的transformer格式，即除了MHRA之外，它还包含动态位置编码(dynamic position encoding, DPE)和前馈网络(forward network, FFN)。这种独特的集成可以有效地增强令牌表示，这在之前的卷积块中没有探讨过。

(二)、全局MHR

在深层，重要的是在更广泛的token空间中开发远程关系，这自然与自注意力的设计有相似的见解。因此，通过比较全局视图中所有令牌的内容相似度来设计令牌的亲和性:

其中$X_j$可以是全局中大小为T ×H×W(对于图像输入T = 1)的任何令牌，而$Q_n(\cdot )$和$K_n(\cdot )$是两个不同的线性变换。
与变压器块比较:全局MHRA Aglobaln (Eq. 7)可以实例化为一个时空自注意，其中Qn(·)，Kn(·)和Vn(·)在ViT中成为Query, Key和V value。因此，它可以有效地学习长期依赖关系。
全局UniFormer块与以前的ViT块不同。首先，大多数视频transformer在视频域中划分时空注意，以减少令牌相似度比较中的点积计算。
但这种操作不可避免地恶化了符号之间的时空关系。相反，全局联合编码时空符号关系，生成更具鉴别性的视频表示，便于识别。由于局部UniFormer块在很大程度上节省了浅层token比较的计算，因此整个模型可以达到较好的计算精度平衡。其次，UniFormer中采用动态位置嵌入(dynamic position embedding, DPE)代替了绝对位置嵌入。它采用卷积样式，可以克服置换不变性，并且对不同长度的可视标记友好。

(三)、动态位置嵌入

位置信息是描述视觉表征的重要线索。此前，大多数vit通过绝对或相对位置嵌入对这些信息进行编码。但是，绝对位置嵌入需要对不同的输入大小进行微调，而相对位置嵌入由于自注意的修改，效果不佳。
为了提高灵活性，最近提出了卷积位置嵌入。特别是，条件位置编码(CPE)可以通过卷积运算符隐式编码位置信息，从而解锁Transformer以处理任意输入大小并提高识别性能。由于其即插即用的特性，灵活地采用它作为在UniFormer中的动态位置嵌入(DPE):

DWConv指的是零填充的深度卷积。选择这样的设计作为我DPE基于以下原因。首先，深度卷积对任意输入形状都是友好的，例如，可以直接使用它的时空版本来编码视频中的3D位置信息。其次，深度卷积是轻量级的，这是计算精度平衡的重要因素。最后，添加额外的零填充，因为它可以帮助令牌通过逐步查询它们的邻居来了解它们的绝对位置。

三、框架设计

四、实验

(一)、图像分类

(二)、视频分类

(三)、目标检测

(三)、语义分割