TPAMI2022(Submission). UniFormer: Unifying Convolution and Self-attention for Visual Revognition

这篇文章是ICLR的接受论文，目前扩展后投往TPAMI但还未接受。

论文地址

代码地址

由于视觉数据中较大的局部冗余和复杂的全局依赖，学习图像和视频确定性的表示是一项具有挑战性的任务。卷积网络能够有效通过小邻域聚集有效降低局部冗余，但感受野尺寸使其难以捕获长期依赖，视觉transformer能够有效捕获长期依赖，但是盲目地在所有token之间建立相似性比较会带来较高的冗余。文章提出了一种最新的Unified transFormer(UniFormer)，无缝集成卷积和自注意力的优点。文章在浅层和深层分别使用全局token affinity，因此能够通过有效的表示学习解决局部冗余和全局依赖。

主要问题

视频数据中面临两项挑战，局部冗余，即局部区域内的视觉内容很相似；复杂的全局依赖，即不同区域的目标之间具有动态关系。主流方法如卷积和ViTs分别使用卷积和自注意力作为核心操作，但仅能解决其中一个问题。卷积通过聚集一个小邻域之中的信息能够避免多余的计算，但有限的感受野不能构建全局依赖。自注意力通过比较视觉token之间的相似性能够学习到图片和视频之间的全局依赖性，但通常不能有效地编码浅层的局部特征。文章可是花了注意力图，从中可以看出给定一个anchor token，空间注意力主要集中于其局部邻域内，很少学习到图片其余部分的区域，而时间注意力也仅仅集中于近邻帧，远距离的长视野则没有关注。因此在所有token之间进行比较会浪费较多计算量，显著额外计算和精度之间的平衡关系。

文章贡献/创新点

文章提出了最新的Unified transFormer(UniFormer)，灵活地以transformer的格式了统一了卷积和自注意力，能够有效解决局部冗余和全局复杂依赖的问题。UniFormer包含三个组件：动态位置编码(DPE)、多头关系聚集(MHRA)和前馈神经网络(Feed-Forward)，关系聚集设计是区别于卷积和视觉注意力的主要部分。

在浅层中，关系聚集使用小的可学习矩阵来捕获局部token affinity，采用了卷积的模式来通过局部区域的情景聚集来显著降低计算冗余。深层中关系聚集使用token相似性比较来学习全局token affinity，采用自注意力的模式来自适应地建立不同区域和帧之间的长期依赖关系。通过分阶段分层式地堆叠局部和全局UniFormer块，能够灵活地结合给两者的能力来实现强大的表示学习。最终文章得到了一种通用且有效的backbone架构，能够适用于各种视觉识别任务。

UniFormer-L在ImageNet-1K上获得了86.3的top-1精度，UniFormer-B经过ImageNet-1K预训练后再Kinetics-400/Kinetics-600上获得了82.9/84.8的top-1精度，再Something-Something V1&V2上获得了60.9和71.2的top-1精度，在COCO检测任务上获得了53.8的box AP和46.4的mask AP，在ADE20K语义分割任务上获得了50.8mIoU，在COCO位置估计任务上获得了77.4AP。

文章方法

从结构图中可以看出，UniFormer一共包含三个关键组件：动态位置编码、多头关系聚集和前馈神经网络。

$$\mathbf X=\mathrm{DPE}(\mathbf X_{in})+\mathbf X_{in}\\ \mathbf Y=\mathrm{MHRA}(\mathrm{Norm}(\mathbf X))+\mathbf X\\ \mathbf Z=\mathrm{FFN}(\mathrm{Norm}(\mathbf Y))+\mathbf Y$$

动态位置编码DPE

对于输入token张量：$\mathbf X_{in}\in\mathbb R^{C\times T\times H\times W}$首先使用DPE来动态地集成所有token的位置编码信息，DPE能够适用于任何输入尺寸的图片。

$$\mathrm{DPE}(\mathbf X_{in})=\mathrm{DWConv}(\mathbf X_{in})$$

其中$\mathrm{DWConv}$是指0填充的按深度卷积，即每个通道单独进行卷积，最后拼接到一起。文章说有三个优点：

1. 能够适用于任意输入形状的图片。
2. 轻量，能够满足计算和精度之间的平衡。
3. 额外的零填充能够帮助token通过阶段性地查询周围邻域来感知绝对位置信息。

多头关系聚集MHRA

MHRA采用多头注意力地形式来利用token之间的关系：

$$\mathrm{R}_n(\mathbf X)=\mathrm A_n\mathrm V_n(\mathbf X)\\ \mathrm{MHRA}(\mathbf X)=\mathrm{Concat}(\mathrm R_1(\mathbf X);\mathrm R_2(\mathbf X);\cdots;\mathrm R_N(\mathbf X))\mathbf U$$

给定$\mathbf X\in\mathbb R^{C\times T\times H\times W}$，首先将形状变换为序列形式：$\mathbf X\in\mathbb R^{L\times C}$，其中$L=T\times H\times W$。然后利用线性变换来讲原始的token编码的成情景token$\mathrm{V}_n(\mathbf X)\in\mathbb R^{L\times \frac CN}$，然后利用token affinity$\mathrm A_n\in\mathbb R^{L\times L}$来指导聚合情景。最后使用$\mathbf U\in\mathbb R^{C\times C}$来集合$N$个注意力头。有两种方式来学习$\mathrm A_n$

局部MHRA

为了降低计算量和冗余数据，仅在小邻域内构建affinity矩阵：

$$\mathrm A_n^{local}(\mathbf X_i,\mathbf X_j)=a_n^{i-j}$$

其中$j\in\Omega_i^{t\times h\times w}$表示$i$周围的邻域，$a_n\in\mathbb R^{t\times h\times w}$是可学习的参数，$(i-j)$表示token$i$和token$j$之间的相对位置。

此处不太理解$a_i^{i-j}$中的$i-j$是标号还是幂。

文章指出局部MHRA可以看作MobileNet网络块的扩展，线性变换$\mathrm V(\cdot)$是将$\mathbf X\in\mathbb R^{L\times C}$的通道维度缩减为$\mathbb R^{L\times\frac CN}$，等价于PWConv，即$1\times 1$卷积。局部token affinity$\mathrm A_n^{local}$是在每个通道执行，因此可以$\mathrm R_n(\mathbf X)=\mathrm A_n^{local}\mathrm V_n(\mathbf X)$可看作一种DWConv。最后$\mathbf U$执行所有注意力头，也可以看作PWConv，因此总体可以看作是PWConv-DWConv-PWConv的MobileNet块结构。

文章此处对比MHRA和MobileNet块结构的分析很有亮点。

• PWConv执行的$1\times 1$的卷积操作是对每一个空间维度像素执行线性变换，就等价于右乘矩阵做整体线性变换。但
• DWConv是对每个通道单独的卷积，卷积仅对周边小邻域执行聚集，而左乘矩阵操作$\mathrm R_n(\mathbf X)=\mathrm A_n^{local}\in R^{L\times L}\mathrm V_n(\mathbf X)$是对空间维度中的所有$L$个元素执行$L$次聚集，两者有很大不同。但此处$\mathrm A_n^{local}$仅计算每个元素核周围邻域内元素的关系，即仅$\mathrm A_n^{local}$对角线附近元素有值。因此当邻域大小$\Omega^{t\times h\times w}$和卷积核尺寸相等时，左乘该矩阵就等价于邻域内做卷积。但全局MHRA就不具备这种等价性。

文章和MobileNet不同之处在于还多了动态位置编码和前馈神经网络，因此能够增强token表示能力。

全局MHRA

全局MHRA不像局部MHRA那样直接作为网络可学习的元素，而是遵循自注意力的做法，对网络中每个元素执行关系相似性度量：

$$\mathrm A_n^{global}(\mathbf X_i,\mathbf X_j)=\frac{e^{Q_n(\mathbf X_i)^\top K_n(\mathbf X_j)}}{\sum_{j'\in\Omega_{T\times H\times W}}e^{Q_n(\mathbf X_i)^\top K_n(\mathbf X_{j'})}}$$

其中$Q_n(\cdot)$和$K_n(\cdot)$都是线性变换。

文章对比了使用的全局MHRA和以往的transformer架构的不同。首先文章没有像以往文章那样拆分空间和时间注意力，而是就直接整体编码，计算量会非常高。其次，文章采用了动态位置编码来克服置换不变性。

这里的分析没能凸显全局MHRA的优点，全局MHRA就是一个普通的自注意力，所说的两个优点来自于局部MHRA和动态位置编码。

仿真实验

应用于分类任务的网络结构为：

目标检测和实例分割

文章由于使用了全局MHRA，计算量非常高，分析后发现主要是阶段3占据了超过50%的FLOPs，而阶段4只占阶段3的1/28。

因此对阶段3进行了改进。在预定义的窗口($14\times 14$)内执行自注意力操作，将计算复杂度从$\mathcal O(w^2h^w)$降低为$\mathcal O(whp^2)$。为了避免过于降低模型性能，文章在阶段3继承了窗口和全局UniFormer块，每个混合组包含3个窗口块和1个全局块，UniFormer-Small/Base两种backbone的阶段3分别包含2/5个混合组。

总体可以看出UniFormer有提升，但相比已有方法如Focal-T、CSWin-S等不是非常显著。文章最有价值的工作就是指出邻域内的affinity矩阵和MobileNet块中的DWConv等价，但文章在邻域内计算$\mathrm A_n^{local}$，导致只有对角线附近元素有值，其余都为0，可能会增加较多的空间复杂度。另外UniFormer在应用于目标检测时还是需要采用以往工作的窗口自注意力来降低计算复杂度，因此模型设计的通用性不太强。