UniFormer：无缝集成 Transformer，更高效的时空表征学习框架 | 论文解读

对图像和视频上的表征学习而言，有两大痛点：局部冗余——视觉数据在局部空间、时间、时空邻域具有相似性，这种局部性质容易引入大量低效的计算；全局依赖——要实现准确的识别，需要动态地将不同区域中的目标关联，建模长时依赖。而现在的两大主流模型 CNN 和 ViT，往往只能解决部分问题。CNN 通过卷积有效降低了局部冗余，但受限的感受野难以建模全局依赖；而 ViT 凭借自注意力可捕获长距离依赖，但在浅层编码局部特征十分低效，导致过高冗余。

为能同时解决上述两大痛点，上海人工智能实验室联合商汤科技共同提出一种新的 UniFormer（Unified Transformer）框架， 它能够将卷积与自注意力的优点通过 Transformer 进行无缝集成。与经典的 Transformer 模块不同，UniFormer 模块的相关性聚合在浅层与深层分别配备了局部全局token，能够同时解决冗余与依赖问题，实现了高效的特征学习。包括 ICLR2022 接受的 video backbone，以及为下游密集预测任务设计的拓展版本，UniFormer 整套框架在各种任务上都能表现出比现有 SOTA 模型更好的性能。

论文标题：UniFormer: Unified Transformer for Efficient Spatiotemporal Representation Learning

论文链接：https://arxiv.org/abs/2201.04676

扩展版本：UniFormer: Unifying Convolution and Self-attention for Visual Recognition

代码链接：https://github.com/Sense-X/UniFormer

*UniFormer 性能提升

1、图像分类：在 Token Labeling 的加持下，仅靠 ImageNet-1K 训练，39GFLOPs 的 UniFormer-L-384 在 ImageNet 上实现了 86.3% 的 top-1 精度。

2、视频分类：仅用 ImageNet-1K 预训练，UniFormer-B 在 Kinetics-400 和 Kinetics-600 上分别取得了 82.9% 和 84.8% 的 top-1 精度（比使用 JFT-300M 预训练，相近性能的 ViViT 的 GFLOPs 少 16 倍）。在 Something-Something V1 和 V2 上分别取得 60.9%和 71.2% 的 top-1 精度，为同期模型的 SOTA。 

3、密集预测：仅用 ImageNet-1K 预训练，COCO 目标检测任务上取得了 53.8 的 box AP 与 46.4 的 mask AP；ADE20K 语义分割任务上取得了 50.8 的 mIoU；COCO 姿态估计任务上取得了 77.4 的 AP。后文将会介绍为下游任务设计的训练和测试时模型适配。

图1：图像分类与视频分类任务性能比较（上方为 ImageNet 上 224x224 与 384x384 分辨率输入）

设计灵感

对图像和视频上的 representation learning 而言，目前存在两大痛点：

• Local Redundancy: 视觉数据在局部空间/时间/时空邻域具有相似性，这种局部性质容易引入大量低效的计算。

• Global Dependency: 要实现准确的识别，需要动态地将不同区域中的目标关联，建模长时依赖。

现有的两大主流模型 CNN 和 ViT，往往只关注解决以上部分问题。Convolution 只在局部小邻域聚合上下文，天然地避免了冗余的全局计算，但受限的感受也难以建模全局依赖。而 self-attention 通过比较全局相似度，自然将长距离目标关联，但通过如下可视化可以发现，ViT 在浅层编码局部特征十分低效。

图2：DeiT 可视化后，可以发现即便经过三层 self-attention，输出特征仍保留了较多的局部细节。任选一个 token 作为 query，可视化 attention 矩阵可以发现，被关注的 token 集中在 3x3 邻域中（红色越深表示关注越多）。

图3：TimeSformer 可视化后，同样可以发现即便是经过三层 self-attention，输出的每一帧特征仍保留了较多的局部细节。任选一个 token 作为 query，可视化 spatial attention 和 temporal attention 矩阵都可以发现，被关注的 token 都只在局部邻域中（红色越深表示关注越多）。

无论是 spatial attention 抑或是 temporal attention，在 ViT 的浅层，都仅会倾向于关注 query token 的邻近 token。但是 attention 矩阵是通过全局 token 相似度计算得到的，这无疑带来了大量不必要的计算。相较而言，convolution 在提取这些浅层特征时，无论是在效果上还是计算量上都具有显著的优势。那么为何不针对网络不同层特征的差异，设计不同的特征学习算子，从而将 convolution 和 self-attention 有机地结合物尽其用呢？

本论文中提出的UniFormer (Unified transFormer)，旨在以 Transformer 的风格，有机地统一 convolution 和 self-attention，发挥二者的优势，同时解决 local redundancy 和 global dependency 两大问题，从而实现高效的特征学习。

模型架构

图4：模型整体框架，标红维度仅对 video 输入作用，对 image 输入都可视作 1

模型整体框架如上图所示，借鉴了 CNN 的层次化设计，每层包含多个 Transformer 风格的 UniFormer block。

1、UniFormer block

每个 UniFormer block 主要由三部分组成，动态位置编码 DPE、多头关系聚合器 MHRA 及 Transformer 必备的前馈层 FFN，其中最关键的为多头关系聚合器：

2、MHRA

与多头注意力相似，将关系聚合器设计为多头风格，每个头单独处理一组 channel 的信息。每组的 channel 先通过线性变换生成上下文 token  ，然后在 token affinity   的作用下，对上下文进行有机聚合。

基于前面的可视化观察，可以认为在网络的浅层，token affinity 应该仅关注局部邻域上下文，这与 convolution 的设计不谋而合。因此，将局部关系聚合  设计为可学的参数矩阵：

3、local MHRA

其中  为 anchor token，  为局部邻域  任一 token，  为可学参数矩阵，(i − j) 为二者相对位置，表明 token affinity 的值只与相对位置有关。这样 local UniFormer block 实际上与 MobileNet block 的设计风格相似，都是 PWConv-DWConv-PWConv（见原论文解析），不同的是研究引入了额外的位置编码以及前馈层，这种特别的结合形式有效地增强了 token 的特征表达。

在网络的深层，需要对整个特征空间建立长时关系，这与 self-attention 的思想一致，因此，通过比较全局上下文相似度建立 token affinity：

4、global MHRA

其中  为不同的线性变换。先前的 video transformer，往往采用时空分离的注意力机制 ，以减少 video 输入带来的过量点积运算，但这种分离的操作无疑割裂了 token 的时空关联。相反，UniFormer 在网络的浅层采用 local MHRA 节省了冗余计算量，使得网络在深层可以轻松使用联合时空注意力，从而可以得到更具辨别性的 video 特征表达。再者，与以往 ViT 中使用绝对位置编码不同，这里采用卷积风格的动态位置编码，使得网络可以克服 permutation-invariance 的同时，对不同长度的输入更友好。

流行的 ViT 往往采用绝对或者相对位置编码 ，但绝对位置编码在面对更大分辨率的输入时，需要进行线性插值以及额外的参数微调，而相对位置编码对 self-attention 的形式进行了修改。为了适配不同分辨率输入的需要，采用了最近流行的卷积位置编码设计动态位置编码:

5、DPE

其中 DWConv 为零填充的的深度可分离卷积。一方面，卷积对任何输入形式都很友好，也很容易拓展到空间维度统一编码时空位置信息。另一方面，深度可分离卷积十分轻量，额外的零填充可以帮助每个 token 确定自己的绝对位置。

整体框架

1、图像分类

表1：模型细节

团队设计了三种不同规模的模型，每个模型包含 4 层，前两层使用 local MHRA，后两层使用 global MHRA。对于 local MHRA，卷积核大小为 5x5，归一化使用 BN（使用 LN 性能较差）。对于 global MHRA，每个 head 的 channel 数为 64，归一化使用 LN。动态位置编码卷积核大小为 3x3，FFN 的拓展倍数为 4。对于特征下采样，团队采用非重叠卷积，其中第一次下采样卷积核大小为 4x4、步长为 4x4，其余三次下采样卷积核大小为 2x2、步长为 2x2。在每次下采样卷积之后，额外增加 LN 归一化。网络最后接平均池化层与线性分类层，输出最终预测。当使用 Token Labeling 时，额外加入一个线性分类层以及辅助损失函数。

对于 UniFormer-S，团队设计了增强版本，每层 block 数量为 [3, 5, 9, 3]，下采样使用重叠卷积，FLOPs 控制为 4.2G，保证与其他 SOTA 模型可比。

2、视频分类

对于 video 使用的 3D backbone，通过加载 ImageNet-1K 预训练的  UniFormer-S 和 UniFormer-B，并进行卷积核展开。具体来说，动态位置编码和 local MHRA 分别展开为 3x3x3 和 5x5x5 卷积。对于下采样层，只在第一次下采样同时压缩时间和空间维度，而在其余三次下采样仅压缩空间维度，即第一次下采样卷积核大小为 3x4x4、步长为 2x4x4，其余卷积核大小为 1x2x2、步长为 1x2x2。这样就可以在减少计算量的同时，还能保证模型的高性能。对于 global MHRA，直接继承相应参数，将时空 token 序列化进行统一处理。

3、密集预测

对于下游密集预测任务，直接使用原始网络作为 backbone 并不合适。因为这些任务往往输入大分辨率图像，比如目标检测中输入 1333x800 的图像，使用 global MHRA 会带来过多的计算量。以 UniFormer-S 为例，统计了不同分辨率输入时，不同操作所需的计算量。

图5：左边图片为模型整体计算量与第三/四层中 MatMul 运算所需计算量的结果；右边图片为第三层采用不同风格的 MHRA 所需的 MatMul 计算量的结果

从上图中可以看到，第三层中 MHRA 所需的 MatMul 运算随着分辨率的增加而急剧上升，在输入分辨率为 1008x1008 时，甚至占了总运算量 50% 以上，而第四层仅为第三层的 1/28。因此，仅对第三层的 MHRA 进行改进。

受先前工作的启发，将 global MHRA 应用在限制的窗口内，这样会把原本  的复杂度降至  ，其中 p 为窗口大小。然而直接应用纯 window 化操作，不可避免地会带来性能下降，为此团队将window 和 global 操作结合。每个 hybrid 分组中包含 4 个 block ，前 3 个为 window block，最后1 个为 global block。UniFormer-S 和 UniFormer-B 分别包含 2 个和 5 个分组。

图6：不同任务在训练和测试时使用特定改进 

如上图所示，在五种任务上，训练和测试采用特定改进。对目标检测任务，由于训练和测试时输入分辨率都很大（如 1333x800），在第三层都采用 hybrid block。对姿态估计任务， 输入分辨率相对较小（如 384x288），在第三层采用原本的 gloabl block。而对于语义分割任务，由于往往在测试时使用几倍于训练输入的分辨率（如 2048x512 vs. 512x512），因此在训练时，对第三层采用 global block，而在测试时采用 hybrid block，但需要注意测试时 hybrid block 中 window size 需要与训练时 global block 的感受野一致（如 32x32），感受野不一致会导致急剧性能下降。这种设计可以保证训练高效的同时，提高测试的性能。

实验结果

1、图像分类

团队用ImageNet-1K 进行图像分类实验，采用了 DeiT 的代码与训练策略，UniFormer-S/B/L 使用 droppath 比例为 0.1/0.3/0.4，对大模型额外加入 Layer Scale 防止网络训练崩溃 。对于 Token Labeling，沿用了其代码训练框架与超参。

表2 ：ImageNet-1K

结果如上所示，其中带*为加入 Token Labaleing 进行训练，UniFormer-B dagger 为前述 UniFormer-B 的增强版本。可以看到 UniFormer 在不同规模下都取得了 SOTA 性能，UniFormer-L-384 仅需 39G 的 FLOPs，即可取得 86.3% 的 top-1 精度。

2、视频分类

团队在 Kinetics-400/600 以及 Something-Something V1/V2 上进行了视频分类实验，沿用了 MViT 的代码和训练策略。对 Kinetics 采用 dense 采样方式，加载了 ImageNet-1K 的预训练模型加速训练，droppath 比例保持与前述一致。对 Sth-Sth 采用 uniform 采样方式，加载了 Kinetics 的预训练，droppath 设置为前述两倍，并不进行水平翻转。

表3（左）：Kinetics；表4（右）：Sth-Sth

结果如上图所示，仅使用 ImageNet-1K 预训练，结果在 Kinetics 上取得了与使用 JFT-300M 预训练的 ViViT-L、使用 ImageNet-21K 预训练的 Swin-B 相近的性能，计算量大幅减小。而在 Sth-Sth 上，UniFormer 大幅高于先前 CNN 和 ViT 的结果，取得了新的 SOTA 结果，在 Sth-Sth V1 上为 61.2%，V2 上为 71.4%。

3、目标检测与实例分割 

团队在 COCO2017 上进行了目标检测和实例分割实验，沿用了 mmdetection 的代码框架，配置了 Mask R-CNN 与 Cascade Mask R-CNN 框架，使用 Swin Transformer 的训练参数，均取得了 SOTA 性能。

表5 ：COCO 目标检测，Mask R-CNN

表6 ：COCO 目标检测，Cascade Mask R-CNN

4、语义分割 

团队在 ADE20K 上进行了语义分割实验，沿用了 mmsegmentation 的代码框架，配置了Semantic FPN 与 UperNet 两种框架，分别使用了 PVT 和 Swin 的训练参数，均取得了SOTA 性能。

表7、表9 ：ADE20K 语义分割，左：Semantic FPN；右：UperNet

5、人体姿态估计 

团队在 COCO2017 上进行了姿态估计实验，沿用了 mmpose 的代码框架，配置了简单的 Top-down 框架，使用了 HRFormer 的训练参数，取得了 SOTA 性能。

表8：COCO 姿态估计，Top-down

6、消融实验

团队进行了详尽的消融实验，首先在图像和视频分类任务上验证了 backbone 的结构设计。其次，对 video backbone 的预训练、训练以及测试策略进行了探究。最后，验证了下游密集预测任务上改进的有效性。

Model designs for image and video backbones

表10 ：结构设计

• FFN：首先将 local block 替换为 MobileNet block，其中 ReLU 替换为 GeLU，expand ration 设置为 3 保持计算量一致，保留动态编码以公平比较，可以发现 local block 在图像 和视频分类任务上都会明显优于 MobileNet block。由此可知 transformer 风格，以及其特有的 FFN 确实增强了 token 的特征表达。

• DPE：将动态位置编码去掉，在图像和视频任务上性能均下降，视频任务上更是掉了 1.7%，由此可知位置编码有助于更好的时空特征学习。

• Local MHRA size：将 local MHRA 的卷积核大小设置为 3、5、7、9，性能差异并不大，最终采用大小为 5 的卷积核，以取得最好的计算量与准确率的权衡。

• MHRA configuration：由纯 local MHRA（LLLL）出发，逐层替换使用 global MHRA。结果可以发现，仅使用 local MHRA 时，计算量很小，但性能下降明显。逐层替换 global MHRA 后，性能逐渐提升。但全部替换为 global MHRA 后，视频分类准确率急剧下降， 计算量急剧上升，这主要是因为网络缺失了提取细节特征的能力，冗余的 attention 在有限的 video 数据下导致了急剧的过拟合。

Pre-training, training and testing for video backbone

表11、12、13、14 ：迁移性能、卷积核展开方式、训练采样方式

• Transfer learning：表 11 比较了不同结构的预训练性能以及小数据集上的迁移学习性能，可以发现，联合的时空学习方式，不仅在预训练模型上性能更好，在小数据集上的迁移学习性能提升明显。而纯 local MHRA 以及时空分离的学习方式，迁移小数据训练未能带来提升。

• Infalting methods：表 12 中比较了是否对卷积核进行展开，可以发现，展开为 3D 卷积核，在场景相关的数据集 Kinetics-400 上性能接近，但在时序相关的数据集 Sth-Sth V1 上提升明显，尤其是在强的预训练的加持下，这表明 3D 卷积核对时空特征的学习能力更强。

• Sampling strides of dense sampling：表 13 中比较了在训练 Kinetics 时，使用不同间隔采样的结果。可以发现更稀疏的采样，在单 clip 测试上效果往往更好，但在多 clip 测试时，间隔 4 帧采样更好。 

• Sampling methods of Kinetics pre-trained model：由于加载 Kinetics 预训练模型训练 Sth-Sth，而 Sth-Sth 采用 uniform 采样，所以有必要知道预训练覆盖更多帧是否能带来提升。表 14 的结果表明，预训练的不同采样方式差别并不大，16x4 采样在大部分条件下性能都较好。

图7（左）：不同数据集预训练。 图8（右）：不同测试策略。

• Pre-trained dataset scales：图 7 比较了不同规模数据预训练的结果，可以发现对于小模型，大数据集预训练的提升非常明显，而对于大模型则相差无几。 

• Testing strategies：图 8 比较了不同的测试策略，可以发现对于使用 dense 采样方式训练的场景相关数据集 Kinetics 而言，多 clip 测试方案较好，且 1x4 综合性能最优。对于使用 uniform 采样方式训练的时序相关数据集 Sth-Sth 而言，多 crop 测试方案较好，且 3x1 综合性能最好。

Adaption designs for downstream tasks

表15、16、17：下游改进

• object detection：表 15 比较了目标检测任务上，第三层采用不同类型 block 的结果。尽管在 1x 训练时，hybrid block 的性能比纯 global block 性能略差，但经过 3x 的充分训练后，hybrid block 的性能已经能和纯 global block 持平。

• semantic segmentation：表 16 比较了语义分割任务上，第三层采用不同类型 block 的结果。可以发现更大的窗口，以及 global block 的使用都能明显提升性能，由于纯 global block 计算量较大，团队采用性能相近的 hybrid block。

• pose estimation：表 17 分别比较了姿态估计任务上，第三层采用不同类型 block 的结果。由于图像分辨率较小，zero padding 消耗了更多的计算量。

Visualizations

图9：图像分类与视频分类

图10：目标检测、语义分割、姿态估计

总结与思考

在过去一年多，研究团队在视频模型设计上尝试了 CNN（CTNet，ICLR2021）、ViT（UniFormer，ICLR2022）及 MLP（MorphMLP，arxiv）三大主流架构，发现：Transformer 风格的模块+CNN 的层次化架构+convolution 的局部建模+DeiT 强大的训练策略，保证了模型的下限不会太低。但相比 convolution 以及 linear layer 而言，self-attention 的性价比仍是最高的，同等计算量尤其是小模型下，self-attention 带来的性能提升明显，并且对大数据集预训练的迁移效果更好。

不过传统 ViT 对不同分辨率输入并不友好，并且对大分辨率输入的计算量难以承受，这在 UniFormer 中都尽可能以简洁的方式解决，DWConv 有限制地引入，也并不会导致过多的显存开销与速度延迟，最后在不同的任务上都能取得很好的性能权衡。UniFormer 提供了一个尽可能简单的框架，也希望后面的研究工作能在这个框架的基础上，去考虑视频中的运动信息、时空维度的冗余性、帧间的长时关系建模等等更复杂的问题，实现更大的突破。