【ICLR2022】UniFormer：无缝集成 Transformer，更高效的时空表征学习框架

2022 年 2 月 16 日 专知

来源：“上海人工智能实验室”

对图像和视频上的表征学习而言，有两大痛点：局部冗余——视觉数据在局部空间、时间、时空邻域具有相似性，这种局部性质容易引入大量低效的计算；全局依赖——要实现准确的识别，需要动态地将不同区域中的目标关联，建模长时依赖。而现在的两大主流模型 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

扩展版本：

https://www.zhuanzhi.ai/paper/8ee119122de9ee20679ec23ce65b230b

代码链接：

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 $V_{n}(X)$ ，然后在 token affinity $A_{n}$ 的作用下，对上下文进行有机聚合。

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

3、local MHRA

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

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

4、global MHRA

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