UniFormer & UniFormerV2 论文笔记

Author: Sijin Yu

1. Information

• Uniformer

标题: UniFormer: Unified Transformer for Efficient Spatiotemporal Representation Learning

arXiv URL: https://arxiv.org/abs/2201.04676

code: https://github.com/Sense-X/UniFormer

期刊/会议: ICLR2022

发表时间: 2022 年 01 月

• UniformerV2

标题: UniFormerV2: Spatiotemporal Learning by Arming Image ViTs with Video UniFormer

arXiv URL: https://arxiv.org/abs/2211.09552

code: https://github.com/OpenGVLab/UniFormerV2

期刊/会议: ICCV2023

发表时间: 2022 年 11 月

2. Abstract (UniFormer)

• 在高维的视频中学习丰富的时空特征表示具有挑战性, 因为在视频的每帧之间, 存在大量的局部冗余和复杂的全局冗余.

• 关于此, 最近的主流框架有: 3D CNNs 和 ViTs.

• 3D CNNs 可以高效地整合局部内容, 以降低局部冗余, 但是因为其感受野有限, 无法降低全局冗余.

• ViTs 可以通过自注意力机制降低全局冗余, 但是因为其盲目地比较所有 tokens, 无法降低局部冗余.

• 我们提出一种新式框架 Unified transFormer (UniFormer), 集成了 CNN 和 ViT 的优点, 并且在计算消耗和正确率之间取得了很好的平衡.

• 与传统的 transfomers 不同, 通过在浅层和深层分别学习局部和全局 token affinity [⚠️注: token affinity (词单元亲和力) 度量了两个 token 之间的关联性], 我们的 relation aggregator (关系聚合器) 可以同时处理时空冗余和依赖关系.

• 我们在主要的 video benchmarks 上做了实验. 在只用 ImageNet-1k 预训练的情况下, UniFormer 在 Kinetics-400 上获得了 82.9% 的准确率, 在 Kinetics-600 上获得了 84.8% 的准确率, 同时需要的 GELOPs 比其它 sota 模型少了 10 倍 [⚠️注: GELOP 指每秒十亿次浮点运算]. UniFormer 在 Something-Something V1 上获得了 60.9% 的正确率, 在 Something-Something V2 上获得了 71.2% 的正确率.

3. Methods (UniFormer)

3.1 Overview

• 由 CNN 层 (灰色梯形) 和 transformer 层 (粉色圆角矩形) 组成.

• 一个 transformer 层 (Stage) 由若干个 UniFormer 块组成.

• 一个 UniFormer 块包含三个主要模块:

  • DPE: Dynamic Position Embedding, 动态位置包埋.

  • MHRA: Multi-Head Relation Aggregator, 多头关系聚合器.

  • FFN: Feed Forward Network, 全连接层.

  具体为:

  $$\begin{array}{}\text{(1)}& \mathbf{\text{X}}=\text{DPE}({\mathbf{\text{X}}}_{in})+{\mathbf{\text{X}}}_{in}\end{array}$$
  $$\begin{array}{}\text{(2)}& \mathbf{\text{Y}}=\text{MHRA}(\text{Norm}(\mathbf{\text{X}}))+\mathbf{\text{X}}\end{array}$$
  $$\begin{array}{}\text{(3)}& \mathbf{\text{Z}}=\text{FFN}(\text{Norm}(\mathbf{\text{Y}}))+\mathbf{\text{Y}}\end{array}$$

  其中, ${\mathbf{\text{X}}}_{in}\in {\mathbb{R}}^{3\times T\times H\times W}$$\textbf X_{in}\in \mathbb R^{3\times T\times H\times W}$.

3.2 Dynamic Position Embedding (DPE)

$\text{DWConv}(\cdot )$$\text{DWConv}(\cdot)$ 是 0 填充的简单 3D depthwise 卷积.它的输出不改变

3.3 Multi-Head Relation Aggregator (MHRA)

• 给定一层的输入张量: $\mathbf{\text{X}}\in {\mathbb{R}}^{C\times T\times H\times W}$$\textbf X \in \mathbb R^{C\times T\times H\times W}$.

• 将其 reshape 成一个 tokens 的序列: $\mathbf{\text{X}}\in {\mathbb{R}}^{L\times C}$$\textbf X\in\mathbb R^{L\times C}$. 这里 $L=T\times H\times W$$L=T\times H\times W$.

• 一层 MHRA 有 $N$$N$ 个 Relation Aggregator (RA). (即 $N$$N$ 头). 以 $R_n(\cdot )$$R_n(\cdot)$ 指代第 $n$$n$ 个.

• 在一个 $R_n(\cdot )$$R_n(\cdot)$ 里的行为:

  $$\begin{array}{}\text{(5)}& R_n(\mathbf{\text{X}})=A_n{V}_n(\mathbf{\text{X}})\in {\mathbb{R}}^{L\times \frac{C}{N}}\end{array}$$

  • $V_n(\cdot )$$V_n(\cdot)$ 是一个全连接层, 输出形状是 ${\mathbb{R}}^{L\times \frac{C}{N}}$$\mathbb R^{L\times\frac CN}$.

  • $A_n\in {\mathbb{R}}^{L\times L}$$A_n\in \mathbb R^{L\times L}$ 表示 token affinity (词单元亲和力). 它在浅层时用于提取局部亲和力, 在深层时用于提取全局亲和力. 即: MHRA 分为 Local MHRA 和 Global MHRA.

    对于 Local MHRA:

    • 对于 tokens 序列 $\mathbf{\text{X}}\in {\mathbb{R}}^{L\times C}$$\textbf X\in \mathbb R^{L\times C}$, 写成 $\mathbf{\text{X}}=[{\mathbf{\text{X}}}_1,\cdots ,{\mathbf{\text{X}}}_L{]}^T$$\textbf X=[\textbf X_1,\cdots,\textbf X_L]^T$, 其中 ${\mathbf{\text{X}}}_i\in {\mathbb{R}}^{1\times C}$$\textbf X_i\in\mathbb R^{1\times C}$ 是一个 token.

    • $A_n\in {\mathbb{R}}^{L\times L}$$A_n\in\mathbb R^{L\times L}$ 是一个矩阵, 其第 $i$$i$ 行、第 $j$$j$ 列的元素 (即 index 为 $(i,j)$$(i,j)$ 的元素) 是 tokens 序列 $\mathbf{\text{X}}$$\textbf X$ 中的 ${\mathbf{\text{X}}}_i$$\textbf X_i$ 和 ${\mathbf{\text{X}}}_j$$\textbf X_j$ 之间的时空亲和力. 这里记作: $A_n^{local}({\mathbf{\text{X}}}_i,{\mathbf{\text{X}}}_j)$$A_n^{local}(\textbf X_i,\textbf X_j)$.

    • 给定一个 token ${\mathbf{\text{X}}}_i$$\textbf X_i$, 可以确定一个关于它的长方体 ${\mathrm{\Omega }}_i^{t_i\times h_i\times w_i}$$\Omega_i^{t_i\times h_i\times w_i}$. 这个长方体度量了它计算局部亲和力的感受野. 具体地:

      $$\begin{array}{}\text{(6)}& t_i=\lfloor \frac{i}{H\times W}\rfloor \end{array}$$
      $$\begin{array}{}\text{(7)}& h_i=\lfloor \frac{i-t_i\times H\times W}{W}\rfloor \end{array}$$
      $$\begin{array}{}\text{(8)}& w_i=(i-t_i\times H\times W)\mathrm{mod}W\end{array}$$

      这里本质上是为了保证:

      给定一个目标 token ${\mathbf{\text{X}}}_i$$\textbf X_i$, 任何一个在它感受野 ${\mathrm{\Omega }}_i^{t_i,h_i,w_i}$$\Omega_i^{t_i,h_i,w_i}$ 内的邻居 token ${\mathbf{\text{X}}}_j$$\textbf X_j$ 应满足:

      $$\begin{array}{}\text{(9)}& |t_i-t_j|\le \frac{t}{2}\end{array}$$

      $$\begin{array}{}\text{(10)}& |h_i-h_j|\le \frac{h}{2}\end{array}$$

      $$\begin{array}{}\text{(11)}& |w_i-w_j|\le \frac{w}{2}\end{array}$$

    • 令 $a_n\in {\mathbb{R}}^{t\times h\times w}$$a_n\in\mathbb R^{t\times h\times w}$ 为可学习参数, 则

      对于在 ${\mathbf{\text{X}}}_i$$\textbf X_i$ 感受野 ${\mathrm{\Omega }}_i^{t_i\times h_i\times w_i}$$\Omega_i^{t_i\times h_i\times w_i}$ 内的 ${\mathbf{\text{X}}}_j$$\textbf X_j$:

      $$\begin{array}{}\text{(12)}& A_n^{local}({\mathbf{\text{X}}}_i,{\mathbf{\text{X}}}_j)=a_n[t_i-t_j,h_i-h_j,w_i-w_j]\end{array}$$

      对于不在 ${\mathbf{\text{X}}}_i$$\textbf X_i$ 感受野 ${\mathrm{\Omega }}_i^{t_i\times h_i\times w_i}$$\Omega_i^{t_i\times h_i\times w_i}$ 内的 ${\mathbf{\text{X}}}_j$$\textbf X_j$:

      $$\begin{array}{}\text{(13)}& A_n^{local}({\mathbf{\text{X}}}_i,{\mathbf{\text{X}}}_j)=0\end{array}$$

    对于 Global MHRA:

    • 矩阵 $A_n\in {\mathbb{R}}^{L\times L}$$A_n\in\mathbb R^{L\times L}$ 的 index 为 $(i,j)$$(i,j)$ 的元素由 $A_n^{global}({\mathbf{\text{X}}}_i,{\mathbf{\text{X}}}_j)$$A_n^{global}(\textbf X_i, \textbf X_j)$ 计算.

    • 对于 Global 情况, 感受野变为 ${\mathrm{\Omega }}^{T\times H\times W}$$\Omega^{T\times H\times W}$.

    • $Q_n(\cdot )$$Q_n(\cdot)$ 和 $K_n(\cdot )$$K_n(\cdot)$ 是两个不同的全连接层, 类似自注意力中的存在. 它们将 ${\mathbb{R}}^C$$\mathbb R^{C}$ 映射为 ${\mathbb{R}}^{hidden}$$\mathbb R^{hidden}$. $hidden$$hidden$ 是超参数.

    • $A_n^{global}({\mathbf{\text{X}}}_i,{\mathbf{\text{X}}}_j)$$A_n^{global}(\textbf X_i, \textbf X_j)$ 的计算规则为

      $$\begin{array}{}\text{(14)}& A_n^{global}({\mathbf{\text{X}}}_i,{\mathbf{\text{X}}}_j)=\frac{\exp (Q_n({\mathbf{\text{X}}}_i{)}^T{K}_n({\mathbf{\text{X}}}_j))}{\sum _{j^{\prime }\in {\mathrm{\Omega }}^{T\times H\times W}}\exp (Q_n({\mathbf{\text{X}}}_i{)}^T{K}_n({\mathbf{\text{X}}}_{j^{\prime }}))}\end{array}$$

• 最后, MHRA 层的行为是:

  $$\begin{array}{}\text{(15)}& \text{MHRA}(\mathbf{\text{X}})=\text{Concat}(R_1(\mathbf{\text{X}});R_2(\mathbf{\text{X}});\cdots ;R_N(\mathbf{\text{X}}))U\end{array}$$

  其中, $U\in {\mathbb{R}}^{C\times C}$$U\in \mathbb R^{C\times C}$ 是可学习的融合矩阵.

4. Experiment & Result (UniFormer)

4.1 和 SOTA 的对比

下图: 在 Kinetics-400 和 Kinetics-600 上的对比, 以低计算代价达到 SOTA 性能.

下图: 在 Something-Something V1 和 Something-Something V2 上达到 SOTA.

下图: 在计算代价和准确率之间的平衡. (横轴: 计算代价, 纵轴: 准确率)

4.2 消融实验

下图: 消融实验.

• Unified 表示是否使用 local MHRA, Joint 表示是否使用 global MHRA, DPE 表示是否使用 DPE. ✖️ 表示使用 MobileNet 块.

• Type 中的 L/G 表示四层 Stage 分别使用 Local MHRA 还是 Global MHRA.

下图: Local MHRA 中 长方体感受野的敏感性.

下图: 迁移学习的能力. 在 K400 上训练, 在 SSV1 上测试.

下图: 采样方法的敏感性. (帧数 $\times$$\times$ 时间步)

下图: Multi-clip/crop 实验.

[⚠️注: Multi-clip 指将视频按不同帧切分为多个视频, 对每个视频做分类, 然后将预测结果组合起来获得最终预测结果. Multi-clop 指将每一帧裁剪为不同的区域, 分别得到不同的多个视频, 然后将各视频的预测结果组合起来获得最终预测结果]

5. Conclusion (UniFormer)

主要创新点:

• 将 CNN 和 ViT 融合.

• 设计了 token 亲和力 $A_n$$A_n$, 与自注意力有异曲同工之妙. 解决了自注意力机制盲目比较所有 token 的问题.

主要成果:

• 能同时处理局部冗余和全局冗余.

• 计算代价相对小.

6. Abstract (UniFormerV2)

• 学习具有区分性的时空特征表示是视频理解的关键问题.

• ViTs 通过自注意力展现了他们在学习长视频依赖性上的强大能力. 但不幸的是, 因为它们盲目地比较所有 tokens, 在处理局部视频冗余问题上展现出局限性.

• UniFormer 通过在 transformer 架构中整合 CNNs 和自注意力为关系聚合器 (relation aggregator), 成功地缓解了这一问题.

• 但是 UniFormer 在视频微调之前, 需要复杂的图像预训练步骤. 这限制了它的广泛应用.

• 与之相反, 大量开源的 ViTs 已经在丰富的图像监督下预训练.

• 基于这些观察, 我们通过结合预训练的 ViTs 和 UniFormer, 提出一种搭建强大视频网络族的通用范式. 我们称之为 UniFormerV2.

• 它包含了全新的 local 和 global aggregators, 并且通过集成 ViTs 和 Uniformer 的优点, 在计算代价和准确率上取得很好的平衡.

• UniFormerV2 在 8 个主流的视频 benchmarks 上获得了 SOTA 性能, 包括 K400/600/700 和 SSV1/V2. 它是第一个在 Kinetics-400 数据集上获得 90% top-1 准确率的模型.

7. Model (UniFormerV2)

7.1 Overview

• 动机: UniFormer 和 ViT 的对比

  

  • $L$$L$ 是 Local MHRA, $G$$G$ 是 Global MHRA, $T$$T$ 是 Temporal MHSA.

  • SL 是有监督学习, CL 是对比学习, MIM 是 Mask Image Modeling.

  • ViTs 在图像预处理方面有优势, 而 UniFormer 在视频微调方面有优势.

  • 本模型旨在融合 ViTs 和 UniFormer 的优势.

• UniFormerV2 模型的设计和与 SOTA 的对比

  

  比起 ViTs 的主要改进:

  • 将 Temporal MHSA 改为 Loacl MHRA.

  • 将 ViT 层的输出传给 Global MHRA 并将它们在最后融合.

下图: 模型框架.

7.2 Local UniBlock

Local UniBlock 的行为:

其中, $\text{LT\_MHRA}(\cdot )$$\text{LT\_MHRA}(\cdot)$ 表示使用局部亲和力的 MHRA. 即 $A_n({\mathbf{\text{X}}}_i,{\mathbf{\text{X}}}_j)$$A_n(\textbf X_i,\textbf X_j)$ 的计算方式为式 (12) (13).

其中, $\text{GS\_MHRA}(\cdot )$$\text{GS\_MHRA}(\cdot)$ 表示使用全局亲和力的 MHRA. 即 $A_n({\mathbf{\text{X}}}_i,{\mathbf{\text{X}}}_j)$$A_n(\textbf X_i,\textbf X_j)$ 的计算方式为式 (14).

[⚠️注: 留意到, Local UniBlock 的输入和输出形状是完全一样的.]

7.3 Global UniBlock

Global UniBlock 的行为:

其中, $\text{DPE}(\cdot )$$\text{DPE}(\cdot)$ 是 Dynamic Position Embedding, 同式 (4).

其中, $\text{C\_MHRA}(\mathbf{\text{q}},\cdot )$$\text{C\_MHRA}(\textbf q,\cdot)$ 是互注意力风格的 MHRA. 它的第 $n$$n$ 个关系聚合器头 (Relation Aggregator) $R_n^C(\mathbf{\text{q}},\cdot )$$R_n^C(\textbf q,\cdot)$ 表示如下, 其中 $\mathbf{\text{q}}\in {\mathbb{R}}^{1\times C}$$\textbf q\in\mathbb R^{1\times C}$ 是一个可学习的 query. [⚠️注: $\mathbf{\text{q}}$$\textbf q$ 的迭代和 Multi-stage Fusion 模块有关]

其中, $A_n^C(\mathbf{\text{q}},\mathbf{\text{X}})$$A_n^C(\textbf q,\textbf X)$ 是互注意力亲和力矩阵, 用于学习 $\mathbf{\text{q}}$$\textbf q$ 和 $\mathbf{\text{X}}$$\textbf X$ 的关系.其计算如下:

其中, $V_n(\cdot )$$V_n(\cdot)$, $Q_n(\cdot )$$Q_n(\cdot)$, $K_n(\cdot )$$K_n(\cdot)$ 都是全连接层.

同式 (15), $\text{C\_MHRA}(\mathbf{\text{q}},\cdot )$$\text{C\_MHRA}(\textbf q,\cdot)$ 的计算为:

最后, 计算 ${\mathbf{\text{X}}}^G$$\textbf X^G$:

7.4 Multi-stage Fusion

这一模块的任务: 将 Local UniBlock 输出的 ${\mathbf{\text{X}}}^L$$\textbf X^L$ 和 Global UniBlock 输出的 ${\mathbf{\text{X}}}^G$$\textbf X^G$ 融合, 并更新 query $\mathbf{\text{q}}$$\textbf q$. 下图: 四种设计思路

记第 $i$$i$ 个 Global UniBlock 的输出为 ${\mathbf{\text{X}}}_i^G=G_i({\mathbf{\text{q}}}_i,{\mathbf{\text{X}}}_i^L)$$\textbf X_i^G=G_i(\textbf q_i, \textbf X_i^L)$.

• Sequential 方式: 更新 ${\mathbf{\text{q}}}_i$$\textbf q_i$ 为上一个 Global UniBlock 的输出, 即:

  $$\begin{array}{}\text{(25)}& {\mathbf{\text{q}}}_1\leftarrow \text{Random}()\end{array}$$
  $$\begin{array}{}\text{(26)}& {\mathbf{\text{q}}}_i={\mathbf{\text{X}}}_{i-1}^G\end{array}$$
  $$\begin{array}{}\text{(27)}& {\mathbf{\text{X}}}_i^G=G_i({\mathbf{\text{q}}}_i,{\mathbf{\text{X}}}_i^L)=G_i({\mathbf{\text{X}}}_{i-1}^G,{\mathbf{\text{X}}}_i^G)\end{array}$$
  $$\begin{array}{}\text{(28)}& \mathbf{\text{F}}={\mathbf{\text{X}}}_N^G\end{array}$$

• Parallel 方式: 随机初始化所有的 ${\mathbf{\text{q}}}_i$$\textbf q_i$, concat 所有 ${\mathbf{\text{X}}}_i^G$$\textbf X^G_i$ 并映射到对应的形状, 即:

  $$\begin{array}{}\text{(29)}& {\mathbf{\text{q}}}_1,{\mathbf{\text{q}}}_2,\cdots ,{\mathbf{\text{q}}}_N\leftarrow \text{Random}()\end{array}$$
  $$\begin{array}{}\text{(30)}& {\mathbf{\text{X}}}_i^G=G_i({\mathbf{\text{q}}}_i,{\mathbf{\text{X}}}_i^L)\end{array}$$
  $$\begin{array}{}\text{(31)}& \mathbf{\text{F}}=\text{Concat}({\mathbf{\text{X}}}_1^G;{\mathbf{\text{X}}}_2^G;\cdots ;{\mathbf{\text{X}}}_N^G)U^F\end{array}$$

  其中, $U^F\in {\mathbb{R}}^{N\times C}$$U^F\in\mathbb R^{N\times C}$ 是可学习的.

• Hierarchical KV 方式: 随机初始化所有的 ${\mathbf{\text{q}}}_i$$\textbf q_i$, 但将上一个 Global UniBlock 的输出作为当前 Global UniBlock 的输入之一, 即:

  $$\begin{array}{}\text{(32)}& {\mathbf{\text{q}}}_1,{\mathbf{\text{q}}}_2,\cdots ,{\mathbf{\text{q}}}_N\leftarrow \text{Random}()\end{array}$$
  $$\begin{array}{}\text{(33)}& {\mathbf{\text{X}}}_i^G=G_i({\mathbf{\text{q}}}_i,[{\mathbf{\text{X}}}_{i-1}^G,{\mathbf{\text{X}}}_i^L])\end{array}$$
  $$\begin{array}{}\text{(34)}& \mathbf{\text{F}}={\mathbf{\text{X}}}_N^G\end{array}$$

• Hierarchical Q 方式: 随机初始化所有的 ${\mathbf{\text{q}}}_i$$\textbf q_i$, 但将上一个 Global UniBlock 的输出作为当前 Global UniBlock 的 $\mathbf{\text{q}}$$\textbf q$ 之一, 即:

  $$\begin{array}{}\text{(35)}& {\mathbf{\text{q}}}_1,{\mathbf{\text{q}}}_2,\cdots ,{\mathbf{\text{q}}}_N\leftarrow \text{Random}()\end{array}$$
  $$\begin{array}{}\text{(36)}& {\mathbf{\text{X}}}_i^G=G_i(\mathbf{\text{[}}{\mathbf{\text{X}}}_{i-1}^G,q_i],{\mathbf{\text{X}}}_i^L])\end{array}$$
  $$\begin{array}{}\text{(37)}& \mathbf{\text{F}}={\mathbf{\text{X}}}_N^G\end{array}$$

实验表明, Sequential 方式和 Hierarchical Q 方式最有效 (见消融实验部分), 模型采用 Sequential 方式作为基准.

8. Experiment & Result (UniFormerV2)

8.1 和 SOTA 的对比

下图: 在 K400 上的性能比较.

下图: 在 K600/700 上的性能比较.

下图: 在 MiT V1 上的性能比较.

下图: 在 SSV2 上的性能比较.

原文在 SSV1, ActivityNet, HACS 上也做了实验, 均达到了 SOTA.

8.2 消融实验

9. Conclusion (UniFormerV2)

主要创新点:

• 将 UniFormer 的设计放入 ViT 框架中.

• 为 ${\mathbf{\text{X}}}^G$$\textbf X^G$ 的融合设计了四种方式.

10. 总结

UniFormer 设计了多头关系聚合器 (MHRA), 很好地处理了视频数据的局部冗余和全局冗余; UniFormerV2 在此基础上, 将 UniFormer 中大获成功的 MHRA 设计成块, 设计了 Local UniBlock 和 Global UniBlock, 并且设计了 Global Cross MHRA, 在 Multi-stage Fusion 中提出四种方式, 结合了 ViTs 的优点, 大大提高了视频理解的性能.