基于时空令牌选择的高效视频 Transformer

视频 transformer 在主要视频识别基准上取得了令人印象深刻的成果,但其计算成本高昂。本文提出 STTS,一种令牌选择框架,可根据输入视频样本动态地在时间和空间维度上选择少量信息丰富的令牌。具体而言,我们将令牌选择表述为排序问题,通过轻量级打分网络估计每个令牌的重要性,仅将得分最高的令牌用于下游评估。在时间维度上,我们保留与动作类别最相关的帧;在空间维度上,我们在特征图中识别最具判别性的区域,而不影响大多数视频 transformer 中以分层方式使用的空间上下文。由于令牌选择决策不可微分,我们采用基于扰动最大化的可微分 Top-K 算子进行端到端训练。我们主要在 Kinetics-400 上使用最近引入的视频 transformer 骨干 MViT 进行了大量实验。我们的框架在所需计算量减少 20% 的情况下取得了相似的结果。我们还证明了我们的方法对不同 transformer 架构和视频数据集具有通用性。代码见 https://github.com/wangjk666/STTS。