本文介绍的是中科院深圳先进技术研究院、商汤和上海 AI Lab 的研究者合作完成的 UniFormer，包括 ICLR 2022 接收的视频 backbone，以及为下游密集预测任务设计的拓展版本，在各种任务上都能取得相较流行 SOTA 模型更好的性能。代码、模型、日志以及训练脚本都已开源。。
- 论文 V1 地址：https://arxiv.org/pdf/2201.04676.pdf
- 论文 V2 地址：https://arxiv.org/pdf/2201.09450.pdf
- 项目地址：https://github.com/Sense-X/UniFormer
- 图像分类：在 Token Labeling [1] 的加持下，仅靠 ImageNet-1K 训练，39G 的 UniFormer-L-384 在 ImageNet 上实现 86.3% 的 top-1 精度；
- 视频分类：仅用 ImageNet-1K 预训练，UniFormer-B 在 Kinetics-400 和 Kinetics-600 上分别取得了 82.9% 和 84.8% 的 top-1 精度（比使用 JFT-300M 预训练，相近性能的 ViViT [2] 的 GFLOPs 少 16 倍）。在 Something-Something V1 和 V2 上分别取得 60.9% 和 71.2% 的 top-1 精度，为同期模型的 SOTA；
- 密集预测：仅用 ImageNet-1K 预训练，COCO 目标检测任务上取得 53.8 的 box AP 与 46.4 的 mask AP；ADE20K 予以分割任务上取得 50.8 的 mIoU；COCO 姿态估计任务上取得 77.4 的 AP。后面会介绍为下游任务设计的，训练和测试时模型适配。
图像分类与视频分类任务性能比较（上方为 ImageNet 上 224x224 与 384x384 分辨率输入）
研究动机
对图像和视频上的表征学习而言，有两大痛点：
- 局部冗余（local redundancy）：视觉数据在局部空间 / 时间 / 时空邻域具有相似性，这种局部性质容易引入大量低效的计算；
- 全局依赖（global dependency）：要实现准确的识别，需要动态地将不同区域中的目标关联，建模长时依赖。
现有的两大主流模型 CNN 和 ViT，往往只关注解决问题之一。卷积只在局部小邻域聚合上下文，天然地避免了冗余的全局计算，但受限的感受野难以建模全局依赖。而自注意力通过比较全局相似度，自然将长距离目标关联，但如下可视化可以发现，ViT 在浅层编码局部特征十分低效。
DeiT 可视化，可以看到即便是经过了三层的自注意力，输出特征仍保留了较多的局部细节。我们任选一个 token 作为查询，可视化注意力矩阵可以发现，被关注的 token 集中在 3x3 邻域中（红色越深关注越多）。
TimeSformer 可视化，同样可以看到即便是经过了三层的自注意力，输出的每一帧特征仍保留了较多的局部细节。我们任选一个 token 作为查询，可视化空间注意力和时序注意力矩阵都可以发现，被关注的 token 都只在局部邻域中（红色越深关注越多）。
无论是空间注意力抑或时序注意力，在 ViT 的浅层，都仅会倾向于关注查询 token 的邻近 token。要知道注意力矩阵是通过全局 token 相似度计算得到的，这无疑带来了大量不必要的计算。相较而言，卷积在提取这些浅层特征时，无论是在效果上还是计算量上都具有显著的优势。那么为何不针对网络不同层特征的差异，设计不同的特征学习算子，将卷积和自注意力有机地结合物尽其用呢？
本文的 UniFormer (Unified transFormer)，旨在以 Transformer 的风格，有机地统一卷积和自注意力，发挥二者的优势，同时解决局部冗余和全局依赖两大问题，实现高效的特征学习。
方法
模型整体框架，标红维度仅对视频输入作用，对图像输入都可视作 1。
模型整体框架如上所示，借鉴了 CNN 的层次化设计，每层包含多个 Transformer 风格的 UniFormer block。
UniFormer block
每个 UniFormer block 主要由三部分组成，动态位置编码 DPE、多头关系聚合器 MHRA 以及 Transformer 必备的前馈层 FFN，其中最关键的为多头关系聚合器：
MHRA
与多头注意力相似，我们将关系聚合器设计为多头风格，每个头单独处理一组 channel 的信息。每组的 channel 先通过线性变换生成上下文 token 
，然后在 token affinity 
的作用下，对上下文进行有机聚合。
基于前面的可视化观察，研究者认为在网络的浅层，token affinity 应该仅关注局部邻域上下文，这与卷积的设计不谋而合。因此，他们将局部关系聚合
设计为可学的参数矩阵：
local MHRA
其中
为 anchor token，
为局部邻域 
任一 token，
为可学参数矩阵， 
为二者相对位置，表明 token affinity 的值只与相对位置有关。这样一来，local UniFormer block 实际与 MobileNet block [3] 的风格相似，都是 PWConv-DWConv-PWConv（见原论文解析），不同的是研究者引入额外的位置编码以前前馈层，这种特别的结合形式有效地增强了 token 的特征表达。
在网络的深层，研究者需要对整个特征空间建立长时关系，这与自注意力的思想一致，因此通过比较全局上下文相似度建立 token affinity：
global MHRA
其中
为不同的线性变换。先前的视频 transformer 往往采用时空分离的注意力机制 [4]，以减少视频输入带来的过量点积运算，但这种分离的操作无疑割裂了 token 的时空关联。相反，UniFormer 在网络的浅层采用 local MHRA，节省了冗余计算量，使得网络在深层可以轻松使用联合时空注意力，从而可以得到更具辨别性的视频特征表达。再者，与以往 ViT 中使用绝对位置编码不同，这里采用卷积风格的动态位置编码，使得网络可以克服排列不变形（permutation-invariance）的同时，对不同长度的输入更友好。
流行的 ViT 往往采用绝对或者相对位置编码 [5]，但绝对位置编码在面对更大分辨率的输入时，需要进行线性插值以及额外的参数微调，而相对位置编码对自注意力的形式进行了修改。为了适配不同分辨率输入的需要，研究者采用了最近流行的卷积位置编码 [6] 设计动态位置编码:
DPE
其中 DWConv 为零填充的的深度可分离卷积。一方面，卷积对任何输入形式都很友好，也很容易拓展到空间维度统一编码时空位置信息。另一方面，深度可分离卷积十分轻量，额外的零填充可以帮助每个 token 确定自己的绝对位置。
框架
图像分类
模型细节
研究者设计了三种不同规模的模型，每个模型包含 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 模型可比。
