Efficient Video Transformers with Spatial-Temporal Token Selection阅读笔记

摘要

Video Transformers在主要视频识别基准测试中取得了令人印象深刻的结果，但其计算成本很高。 在本文中，我们提出了 STTS，这是一种令牌选择框架，它根据输入视频样本在时间和空间维度上动态选择一些信息丰富的令牌。 具体来说，我们将令牌选择制定为排名问题，通过轻量级评分器网络估计每个令牌的重要性，并且只有得分最高的令牌才会用于下游评估。 在时间维度上，我们保留与动作类别最相关的帧，而在空间维度上，我们识别特征图中最具辨别力的区域，而不影响大多数视频中以分层方式使用的空间上下文 变压器。 由于标记选择的决策是不可微的，因此我们采用基于扰动最大的可微 Top-K 运算符进行端到端训练。 我们主要使用最近推出的视频变压器主干 MViT 在 Kinetics-400 上进行广泛的实验。 我们的框架实现了类似的结果，同时需要减少 20% 的计算量。 我们还证明我们的方法对于不同的变压器架构和视频数据集是通用的。 代码可在 https://github.com/wangjk666/STTS 获取。

Introduction

直接将这样的想法从图像transformer推广到视频transformer虽然很吸引人，但并非易事。 视频中的标记通常采用 3D 立方体的形式，视频transformer中的标记化层通常将所有立方体展平为 3D 向量序列。 因此，简单地学习使用基于采样的方法在序列中保留什么内容[49,46]不可避免地会在空间和时间上产生一组不连续的标记，从而破坏视频中的结构信息。 这也与视频transformer的最新设计相冲突，视频transformer以分层方式处理 3D 标记，同时保留空间和时间上下文 [14,39]。 相反，我们认为视频transformer中时空标记的选择应该以顺序方式处理——首先关注整个时间范围内的显着帧，然后深入到这些帧中寻找最重要的空间区域。

在本文中，我们介绍了空间-时间令牌选择（STTS），这是一个轻量级的即插即用模块，可以学习在视频transformer中在空间和时间上分配计算资源。 特别是，STTS 由时间标记选择网络和空间标记选择网络组成，相互协作以使用尽可能少的标记。 每个选择网络都是一个多层感知器 (MLP)，用于预测每个标记的重要性得分，并且可以附加到 Transformer 模型的任何位置。 根据这些分数，我们选择一些分数较高的标记进行下游处理。 更具体地说，给定一系列输入标记，STTS 首先在整个时间范围内选择一些重要的帧。 然后，对于每一帧，我们将标记序列分割成具有规则形状的锚点，并仅选择对视频识别贡献最大的一个锚点。 然而，值得指出的是，选择得分最高的标记是不可微分的，因此给训练带来了挑战。 为了缓解这个问题，我们采用最近提出的可微分 Top-K 选择算法 [4]，使用扰动最大值方法进行端到端的选择训练。 这也允许我们明确控制使用多少token。

我们使用 MViT [14] 和 VideoSwin [39] 作为主干，对两个大型视频数据集 Kinetics- 400 [28] 和 Something-Something-v2 [21] 进行了广泛的实验，以说明我们方法的有效性。 实验结果表明，STTS 可以有效提高效率，但精度略有损失。 特别是，通过仅保留 50% 的输入标记，STTS 通过使用 MViT 将千兆浮点运算 (GFLOP) 中测量的计算成本降低了 33% 以上，并且在 Kinetics-400 上精度下降了 0.7% 以内 [14 ] 作为骨干。 当分配更多计算资源时，STTS 实现了与原始模型相似的性能，但节省了 20% 以上的计算量。

相关工作

高效的视频识别 在过去的几年里，人们对视频识别进行了广泛的研究，调查了卷积神经网络[18,22,17,67]和变压器模型[3,14,39]的使用。 然而，这些模型通常计算成本很高，这刺激了高效视频识别方法的发展[69,76,65,5,40,74,56,64]以加快视频识别的推理时间。 多项研究 [76,62,17,16,31] 探索通过压缩 3D CNN 来设计用于视频识别的轻量级模型。 尽管显着节省了内存占用，但他们仍然需要关注输入视频的每个时间剪辑，因此在计算复杂性降低方面没有带来任何好处。 另一方面，最近的一些方法[32,5,65,63]建议选择最显着的时间片段输入到骨干模型以实现资源高效的视频识别。 与这些专注于加速基于 CNN 的视频识别模型的方法不同，据我们所知，我们是第一个探索视频变压器高效识别的方法。 还值得指出的是，STTS 是正交的，并且是对设计高效视觉转换器的最新方法的补充 [58,30]。

可区分的令牌选择令牌选择的决策是离散的，因此不适合端到端训练。 为了克服这个限制，一种解决方案是采用 Gumbel-Softmax 技巧 [26] 来决定是否应该修剪每个标记 [46,49]。 然而，这无法明确控制训练期间保留标记的数量，因此与视频转换器的常见分层设计相冲突[14,39,34]。 另一种方法是将令牌选择制定为排名问题，其中可以采用基于最优传输的方法[66,9]来匹配支持递增值的辅助概率度量。 在本文中，我们按照[8]应用扰动极大值方法，该方法经[8]验证优于Sinkhorn算子[66]。

Spatial-Temporal Token Selection

图 2：提议的 STTS 概述。 时间标记选择网络和空间标记选择网络可以插入到Transformer块之间，以分别在时间和空间维度上执行标记选择。 N1、N2、N3 可以为 0。

3.1

给定视频 V ∈ RT ×H ×W ×3 由 T 个大小为 H × W 的 RGB 帧组成，视频转换器通常采用两种可能方法之一将视频帧映射到补丁嵌入序列。 第一种是通过 2D 卷积独立地对每个帧内的 2D patch 进行标记，并将所有标记沿时间维度连接在一起 [3]，而另一种是直接从输入视频中提取 3D 管，并使用 3D 卷积将它们线性投影到 3D 嵌入[14,39]。 在这两种情况下，令牌的数量都与输入视频的时间长度和空间分辨率成正比。 我们将所得的时空补丁嵌入表示为 x ∈ RM×N×C，其中 M 和 N 分别表示时间和空间维度上令牌序列的长度，C 是嵌入维度。 位置编码被添加到 x 以注入位置信息[3,14,39]。 为了对视频中的外观和运动线索进行建模，补丁嵌入x被馈送到一堆变压器块中，这些变压器块联合计算空间和时间自注意力[14,39]或单独[2,3]。 自注意力通常表述为：

3.2

从公式1可以看到视频transformer的计算复杂度随着自注意力块中使用的 token 数量呈二次方增长。 考虑到视频中固有的空间和时间冗余，减少计算的一种自然方法是减少标记的数量。 然而，确定保留或丢弃哪些标记并非易事，这与输入样本和手头的目标任务密切相关。 受到最近关于高分辨率图像识别的补丁选择工作的启发[8]，我们将标记选择制定为排名问题——首先使用轻量级评分器网络估计输入标记的重要性分数，然后将前 K 个评分标记表示为 然后选择进行下游加工。 下面我们详细介绍这两个步骤，并分别介绍如何将它们应用于空间和时间标记选择。这与输入样本和手头的目标任务密切相关。 受到最近关于高分辨率图像识别的补丁选择工作的启发[8]，我们将标记选择制定为排名问题——首先使用轻量级评分器网络估计输入标记的重要性分数，然后将前 K 个评分标记表示为 然后选择进行下游加工。 下面我们详细介绍这两个步骤，并分别介绍如何将它们应用于空间和时间标记选择。

3.3 打分网络

给定输入标记序列 q ∈ RL×C ，评分器网络的目标是为每个标记生成输入条件的重要性得分。 这里，L 表示展平序列长度，C 是嵌入维度。 我们采用标准的两层全连接（FC）网络来生成这样的分数。 特别是，输入标记首先通过线性投影映射到局部表示 fl：

其中 θ2 是网络权重，s ∈ RL 是使用最小-最大标准化 [8] 标准化的所有标记的得分向量。 值得一提的是，与通过修剪无信息标记节省的计算成本相比，评分器网络带来的额外计算开销可以忽略不计，如4.2 节所示。

可微分的 Top-K 选择。
给定评分器网络生成的重要性分数，我们选择 K 个最高分数并提取相应的标记。 我们将此过程表示为 Top-K 运算符，它返回 K 个最大条目的索引： y = Top-K(s) ∈ NK 。 (4) 我们假设对索引进行排序以保留输入序列的顺序。 为了使用矩阵乘法实现标记选择，我们将 y 转换为一堆 K 个单热 L 维向量 Y = [Iy1,Iy2,…,IyK] ∈ {0,1}L×K。 因此，具有前 K 个分数的标记可以被提取为 q′ = YTq。 请注意，此操作是不可微的，因为 Top-K 和 one-hot 操作都是不可微的。
为了使用端到端训练来学习记分器网络的参数而不引入任何辅助损失，我们采用扰动最大值方法[4,8]来构建可微分的Top-K算子。 特别是，选择 Top-K 令牌相当于求解以下形式的线性程序：

其中 s1T 是重复 K 次的得分向量 s，C 是凸多胞形约束集，定义为：

我们按照[4]执行前向和后向运算来求解方程。 5. – 前向：等式中 Top-K 运算的平滑版本。 5 可以通过取随机扰动的期望来实现：

其中Z是从均匀高斯分布中采样的随机噪声向量，σ是控制噪声方差的超参数。 在实践中，我们运行 Top-K 算法 n（在我们所有的实验中设置为 500）次，并计算 n 个独立样本的期望。 – 向后：根据[1]，上述前向传递的雅可比行列式可以计算为：

在 Z 服从正态分布的特殊情况下，上式已得到简化。 这样，我们就可以通过 Top-K 操作进行反向传播。 我们以端到端的方式使用交叉熵损失来训练骨干模型和令牌选择网络。 在推理过程中，我们利用硬 Top-K（在 Pytorch 中以 torch.topk 的形式实现[47]）来进一步提高效率。 我们遵循[8]，其中仅执行单个 Top-K 操作（而不是 n 次扰动重复），并且通过张量切片来实现标记选择。 然而，在推理过程中应用硬 Top-K 会导致训练-测试差距。 为了解决这个问题，我们在训练期间将 σ 线性衰减到零。 请注意，当 σ = 0 时，可微分 Top-K 操作相当于硬 Top-K，并且流入记分器网络的梯度消失。

3.3 空间和时间的实例化

考虑到视频中外观和运动线索的独特结构，我们在空间和时间上分别执行标记选择——首先关注显着帧，然后深入这些帧以寻找最重要的空间区域。 类似的想法在之前的解耦 3D 卷积核 [53]、时空非局部块 [23] 和自注意力 [3] 的工作中已经被探索过。
时间选择。 给定输入标记 x ∈ RM×N×C，时间选择的目标是选择 M 帧中的 K 帧并丢弃其余帧。 我们首先沿空间维度平均池 x 以获得一系列基于帧的标记 xt ∈ RM×C，然后将其馈送到评分器网络和 Top-K 运算符（第 3.2 节）以生成指示矩阵 前 K 个最高得分的帧：Yt ∈ RM×K。 最后，我们将输入 x 重塑为 x ∈ RM×(N×C) 并使用指示矩阵提取所选的 K 帧：

所选令牌被重塑为 z ∈ RK×N×C 以进行下游处理。

空间选择。 与时间选择相反，空间选择是在每一帧上分别执行的，我们的目标是为每一帧选择 N 个标记中的 K 个。 更具体地说，我们首先将第 m 帧 xm ∈ RN×C 的标记馈送到记分器网络以生成重要性分数 sm ∈ RN 。 为了简洁起见，我们省略了下标 m，因为相同的操作适用于所有帧。 为了获得前 K 个空间标记，一种简单的方法是将 Top-K 运算符（第 3.2 节）直接应用于基于标记的分数 s。 然而，这种设计不可避免地破坏了输入标记的空间结构，这对于视频转换器中的空间选择尤其不合适。 首先，最近表现最好的视频转换器[14,38]涉及分层架构，该架构逐渐降低多个阶段的空间分辨率。 令牌的不连续丢弃行为不利于局部操作，例如用于空间下采样的卷积和池化。 其次，空间标记在时间维度上的错位使得时间建模[3]即使不是不可能，也更具挑战性。 我们在表 6 中提供了经验证据来支持我们的主张。

我们没有执行基于标记的选择，而是引入了一种新颖的基于锚的空间选择设计。 在获得每帧的重要性分数 s 后，我们首先将其重塑为 2D 分数图 ss ∈ R N× N 并将该图分割为重叠锚点的网格 ϰss ∈ RG×K ，锚点大小为 K，其中 √是锚点的数量，α 是步幅。 锚点的视觉示例 α 可以在图 1 和图 4 中找到。之后，我们通过平均池化聚合每个锚点内的分数，以获得基于锚点的分数 sa ∈ RG。 原来的 Top-K 选择现在转换为 Top-1 选择问题，我们在第 2 节中再次利用 Top-K 运算符（K = 1）。 3.2 获取指标矩阵，最终提取得分最高的anchor。

实验

在本节中，我们通过使用两个最新的视频转换器主干对两个大规模视频识别数据集进行广泛的实验来评估 STTS 的有效性。 我们在第二节介绍实验设置。在第 4 节中介绍主要结果。 并进行消融研究以验证第 4.2 节中不同成分的影响。

4.1实验设置

数据集和骨干网。 我们主要使用最先进的视频转换器 MViT-B16 [14] 作为基础模型，并评估 STTS 在 Kinetics-400 [28] 上的有效性。 为了证明我们的方法对于不同的 Transformer 架构和数据集是通用的，我们还在 Kinetics-400 和 Something-Something-V2 (SSV2) [21] 上使用 Video Swin Transformer [39] 进行实验。 具体来说，Kinetics-400 由属于 400 个动作类别的 240k 训练视频和 20k 验证视频组成。 SSV2 是一个时间敏感数据集，包含涵盖 174 个动作类别的 220,847 个视频。 我们用 B-TLR-SLR 表示 STTS 的不同变体，其中 B 表示骨干网络，T 和 S 分别表示沿时间和空间维度执行的令牌选择。 L和R表示令牌选择模块插入的位置以及相应的被选择令牌的比例。 例如，MViT-B16-T0.4-S40.6表示在第0个自注意力块之前以0.4的选择比率进行时间令牌选择，在第4个自注意力块之前以0.6的比率进行空间令牌选择，使用MViT- B16为基础型号。

评估指标。 为了衡量分类性能，我们报告了验证集上的 top-1 准确度。 我们用 FLOP（浮点运算）来衡量计算成本，这是一个独立于硬件的指标。 实施细节。 在我们的实验中，我们使用 STTS 模块对预先训练的视频转换器进行微调。 STTS 模块中的参数是随机初始化的。 我们在方程中设置 σ。 7 至 0.1。 主干层的学习率设置为 STTS 模块中学习率的 0.01 倍。 为了在 Kinetics-400 上训练 MViT 模型，我们按照 [14] 对 16 帧的剪辑进行采样，时间步长为 4，空间大小为 224×224。 该模型使用 AdamW [42] 进行 20 个时期的训练，其中前 3 个时期用于线性预热 [20]。 动态令牌选择模块的初始学习率设置为1e−4，骨干网络的初始学习率设置为1e−6，小批量大小为16。采用余弦学习率调度[41]。 为了实现 Video Swin Transformer，我们使用时间步长 2 和空间大小 224 × 224 从每个完整视频中采样 32 帧的剪辑。我们将选择网络的学习率设置为 3e−4，并使用 骨干模型的0.01×学习率。 批量大小设置为 64。在 Kinetics-400 上训练时，我们使用 AdamW 优化器 [29] 进行 10 个周期，使用余弦衰减学习率调度器和 0.8 个周期的线性预热。 在 Something-Something-V2 上进行训练时，我们采用 AdamW 优化器进行 20 个 epoch 的较长训练，其中包含 0.8 个 epoch 的线性预热。 为了进行推理，我们应用与原始骨干模型相同的测试策略以进行公平比较。

4.2 main results

STTS 的有效性。 我们首先将 STTS 与一些常见的标记选择基线进行比较：（1）随机（Rand.），它随机采样 K 个标记，而不考虑其视觉内容； (2)注意力(Atten.)，选择类别token为2的前K个最高注意力分数的token； (3) Gumbel-Softmax (GS)，它使用 Gumbel-softmax 技巧 [26] 进行标记选择。 请注意，由于最近的视频转换器中存在空间下采样，因此 GS 无法应用于空间标记选择。 请参阅补充材料了解更多详细信息。 我们在表 1 中总结了仅时间标记、仅空间标记和联合标记选择的结果。不可行的基线设置用 N/A 填充。 可以看出，在相似的计算预算下，与所有基线方法相比，STTS 实现了最佳精度。 特别是，STTS 的性能远远优于 GS，尽管它们共享相同的评分器网络设计（第 3.2 节），这表明我们的可微分 Top-K 运算符对于动态令牌选择的有效性。 我们进一步将 STTS 与兰德进行比较。 不同计算预算下的基线。 如图 3 所示，STTS 始终取得优异的结果，特别是对于计算减少率较高的设置。 例如，MViT-B16-T0.4-S40.6 的性能优于 Rand。 当使用类似的 12 GFLOP 时，性能降低 5.7%。 这验证了通过我们的动态令牌选择模块可以很好地保存信息令牌。 还可以观察到，我们的令牌选择模块的计算开销可以忽略不计——实际上，评分器网络的参数和 FLOP 仅为原始 MViT-B16 主干中的 1.0% 和 0.7%。 与现有技术的比较。 在表 2 中，我们将 STTS 与 Kinetics-400 [7] 上最先进的视频识别模型进行了比较，包括基于 CNN 和基于 Transformer 的模型。

为了证明我们的方法可以推广到不同的转换器架构，我们使用 MViT [14] 和 VideoSwin [39] 作为基础模型。 我们报告了推理时的总体计算成本——单个视图的成本 × 空间和时间中的视图数量，以 Tera-FLOPs (TFLOPs) 为单位。 为了进行清晰的比较，我们将模型分为两组，并将 STTS 与具有可比 TFLOP 的模型进行比较。 请注意，我们的 STTS 的默认设置是 MViT-T0-S4 和 VideoSwin-B-T0，因为根据消融研究，它们实现了最具竞争力的结果，我们将在第 2 节中探索其他选项。 4.3. 此外，我们只对VideoSwin-B执行时间选择，因为Swin转换器中的窗口洗牌操作很复杂[72]，这使得空间选择特别具有挑战性。 最近的研究表明，现代架构中可能不需要移位操作[68]。 我们观察到配备 STTS 的模型在两个复杂性级别上表现出有利的复杂性/准确性权衡。 值得注意的是，我们的 MViT-B16-T0.9-S40.9 取得了与最近的 TimeSformer [3] 相当的结果，而仅需要 47% 的计算成本。 同样，我们的 VideoSwin-B-T0.6 优于最新的视频转换器（例如 TimeSformer-HR [3]、Mformer- HR [48]），同时节省了 50% 以上的计算量。 还值得一提的是，对于这两个基础模型，我们的 STTS 能够节省至少 10% 的计算成本，而精度下降不到 1%。

Discussion

不同地点/多个步骤的token选择。 我们的 STTS 模块的灵活性意味着可以使用不同的令牌选择配置来实现类似的计算减少率。 例如，为了将 MViT-B16 的计算量减少大约 50%，可以在早期阶段应用（1）仅空间/仅时间标记选择，并具有更高的选择率（例如，-S0.6 /-T0.6); （2）后期阶段仅进行空间/仅时间标记选择，选择率较低（例如 -S40.3 / -T40.3 ）； 或（3）联合令牌选择（例如，-T0.6-S40.9），可选地以多步骤方式（表5）。 我们在本节中以 MViT-B16 为例并对 Kinetics-400 进行评估，对这些选择进行深入分析。 我们报告单个视图的推理成本。 表 4 显示了在不同位置执行令牌选择时 STTS 的性能。 对于所有设置，都会调整所选令牌的比率，以确保总体计算成本降低约 50%。 我们观察到时间选择大大超过了空间选择，这表明视频中的时间冗余比空间冗余更重要。 此外，联合标记选择（早期层的时间标记选择和深层的空间标记选择，即-T0.6-S40.9）取得了最好的结果。 我们还以多步骤的方式执行令牌选择，即在变压器网络的不同层执行多个令牌选择，每个令牌选择比更高。 表 5 中的结果表明，尽管计算成本相似，但多步选择产生的结果比单步选择差。 我们假设多步选择会导致视频时空结构的频繁变化，并且更难以训练。

定性结果。 我们在图 4 中可视化了仅时间、仅空间和联合标记选择结果，其中 STTS 丢弃了屏蔽帧和区域。 我们观察到，我们的 STTS 不仅可以有效地识别视频剪辑中信息最丰富的帧，还可以定位每个帧内的判别区域。 通过在时间和空间维度上进行这种标记修剪，只有对于正确识别动作很重要的标记才会被保留并输入到后续的视频转换器中，从而节省了计算成本，并且分类精度的下降最小。

Conclusion

在本文中，我们引入了 STTS，一种动态时空标记选择框架，以减少视频转换器的时空冗余，以实现高效的视频识别。 我们将令牌选择制定为排名问题，其中通过轻量级选择网络预测不同令牌的重要性，并且仅保留具有 Top-K 分数的令牌以进行后续处理。 在时间维度上，我们选择了与动作类别相关性最高的帧子集，而在空间维度上，我们保留了每个帧内最具辨别力的区域，以保留输入帧的结构信息。 为了实现主干模型与令牌选择模块的端到端训练，我们利用了基于扰动最大的可微分 Top-K 算子。 对多个视频识别基准进行的广泛实验验证了我们的 STTS 可以实现有竞争力的效率与准确性权衡。