《X-ViT》-基于时空混合attention的视频Transformer,大幅度降低计算复杂度

作者 | 小马

编辑 | CV君

报道 | 我爱计算机视觉(微信id:aicvml)

【写在前面】

本文介绍了利用Transformer进行的视频识别问题。最近Transformer在视频识别领域的尝试在识别精度方面展现出了非常不错的结果,但在许多情况下,由于时间维度的额外建模,会导致显著的计算开销提升。在这项工作中,作者提出了一个视频Transformer模型,该模型的复杂度与视频序列中的帧数呈线性的关系,因此与基于图像的Transformer模型相比,不会产生额外的计算开销。

为了实现这一点,本文的视频Transformer对 full space-time attention进行了两个方面的近似:

(a)它将时间注意力限制在一个局部时间窗口 ,并利用Transformer的深度来获得视频序列的全时间覆盖(这一点类似CNN中用小卷积核和深层结构获得全局空间建模的思想很像)。

(b)它使用有效的时空混合来联合建模空间和时间信息 ,与仅空间注意模型相比没有产生任何额外的计算成本。作者通过实验证明了,本文的模型在视频识别数据集上产生了非常高的精度。

1. 论文和代码地址

Space-time Mixing Attention for Video Transformer

论文地址:arxiv.org/abs/2106.0596

代码地址:未开源

2. Motivation

视频识别任务指的是识别视频序列中感兴趣的事件,如人类活动的问题。随着Transformer在处理序列数据方面取得的巨大成功,特别是在自然语言处理(NLP)任务中,视觉Transformer最近被证明在图像识别方面也优于CNN。在此基础上,作者提出了一个视频Transformer模型来提高视频识别任务的准确性。

视觉Transformer在时空领域的一个直接的、自然的扩展是在所有的S个空间位置和T个时间位置上共同进行Self-Attention。但是全时空注意具有[公式]的复杂性,这使得模型计算起来很沉重,甚至比3D CNN的复杂度还要高。因此,在本文中,作者的目标是在保留对视频流中存在的时间信息建模的同时,最小化Transformer框架内的计算负担,以实现有效的视频识别。

这个问题的一个基本解决方案是只考虑空间注意,然后在时间维度上平均,它具有[公式]的复杂度(上图(b))。还有一些方法(上图(c))在视频识别精度方面已经显示出了很不错的结果,但在大多数情况下,由于时间信息的额外建模,与baseline(仅空间建模)方法相比,它们还是会导致显著的计算开销。

本文的主要贡献是提出了一个具有[公式]复杂度的视频Transformer模型,因此,与baseline模型一样高效,它在效率(即精度/FLOP)方面显著优于最近提出的工作。为了实现这一点,本文的模型对视频Transformer中使用的全时空注意力进行了两个方面的近似(如上图(d)):

(a)它将时间注意力限制在一个局部时间窗口,并利用Transformer的深度来获得视频序列的全时间覆盖。

(b)它使用有效的时空混合来联合建模空间和时间位置,并且不在仅空间注意的模型上产生任何额外的计算成本。

作者在本文中展示了如何集成两个非常轻量级的机制来实现全局时间attention,以最小的计算成本提供额外的精度提升。实验表明本文的模型在各种视频识别数据集(Something-Something-v2, Kinetics ,Epic Kitchens)上产生非常高的识别精度,同时比其他视频Transformer更加高效。

3. 方法

3.1. Video Transformer

给定一个视频clip[公式]。基于ViT的处理方法,每一帧被分成K×K个不重叠的patch,然后使用线性embedding层[公式]将patch映射到视觉token。

由于Self-Attention是排列不变的,为了保留每个patch在空间和时间内的位置信息,作者还用了两个可学习的位置embedding,一个用于空间:[公式],一个用于时间:[公式]。然后将这些token添加到初始的视觉token中,用L个Transformer layers进行信息的建模。

第l层、空间位置s和时间位置t处的视觉token可以表示为:

除了从视频中提取的[公式]个视觉token外,在token序列中还包括一个特殊的分类token[公式]。第[公式]个Transformer layer使用一系列多头自注意(MSA)、层归一化(LN)和MLP([公式])对上一层的视觉token[公式]进行处理,表示如下:

一个全时空自注意(SA) head的计算可以表示为:

其中[公式]是由[公式]使用embedding矩阵[公式]计算的query、key和value向量。最后,利用embedding矩阵[公式]对h个head的输出concat后进行投影。

完整模型的复杂性为:[公式](qkv投影)+[公式](h个head的MSA)+[公式](多个head的投影)+[公式](MLP)。在本文中,作者的目标是降低h个head的MSA的复杂度([公式])。为了清晰起见,在后面的介绍中,我们先忽略[公式]这个维度,专注于时间和空间维度。因此,整个时空注意的复杂度为[公式]

3.2. Baseline

本文的Baseline模型在每个Transformer层应用仅空间注意,对时空注意进行了一个简单的近似,计算方式如下:

其复杂度是[公式]。为了捕获时间上的信息,作者对时间维度进行了一个简单的平均操作,以获得输入线性分类器的单一特征:

3.3. Recent Work

最近的工作主要集中在减少全时空注意的计算成本[公式]。文献[1]提出了由因子分解后的attention,也就是将时间和空间进行分开attend:

其中[公式]是基于[公式]计算的query、key和value向量。上述模型将复杂性降低到[公式]。然而,对于固定的空间位置,不同帧之间的同一空间信息发生巨大偏差时(比如镜头切换等场景),时间注意是无效的。

3.4. Our model

本文的模型旨在保持[公式]复杂度的同时,更好地近似完全的时空自注意,即不增加仅空间自注意的复杂度。为了实现这一点,作者首先进行了一个近似操作来执行全时空注意,但将时间注意的范围限制在一个局部时间窗口[公式]

其中

[公式][公式]

是代表相应注意力权值的向量。对于单个Transformer层,[公式]是局部窗口[公式]中视觉token的时空组合。在k个Transformer Layer之后,[公式]将是局部时间窗口[公式]中视觉token的时空组合,因此经过多个Transformer Layer之后,时间维度上的感受野就变大了。

上式中局部自注意的复杂度为[公式]。为了进一步减少复杂度,作者在第一个近似之上做了第二个近似。上式中位置[公式][公式]的attention表示如下:

它需要计算[公式]个位置的attention,因此在时间维度上依旧多了计算局部时间attention的复杂度。在计算[公式]的时间区间上,作者采用了一个单一的attention,来减少计算复杂度。具体实现上,作者采用了“shift trick”,这使得不需要引入任何额外的参数和计算量,就能融合时间和空间维度。

每个时间区间[公式]的通道[公式]从通道[公式]中分配(即[公式],这一步就是指新的特征的通道维度有来自时间窗口内其他特征的信息,因此新的通道维度就同时融合了时间维度的信息,具体可见下面的图示 )。设[公式]表示从[公式]中取出下标为[公式]的通道。新的key向量表示为:

上图展示了新的key [公式]形成的过程,用相似的方法,还可以构建新的value[公式]最后,新的全时空注意的近似方法计算如下:

通过这两种近似方法,就可以获得[公式]的复杂度,因此它比以前提出的视频Transformer更有效。除了更小的计算量,本文的方法还能获得更高的精度。

3.5. Temporal Attention aggregation

模型使用class token [公式]来生成预测。作者尝试了以下几种方式来聚合时间的信息:

(1)在时间维度上简单进行时间维度的平均:

(2)直接平均的一个缺点是,输出就是每帧特征的集合,因此,完全忽略了它们之间的时间顺序。为了解决这个问题,作者使用了一种轻量级的时间注意(TA)机制,它将attend T个class token。在实现上,[公式]使用时间Transformer来attend到序列[公式],将其作为分类器的输入。

3.6. Summary token

作为TA的替代方案,作者还提出了一种简单的轻量级机制,用于网络中间层不同帧之间的信息交换。给定第t帧的token集合[公式],作者计算了一组新的长度为[公式]的token [公式],它对token信息进行了总结(记为“Summary” tokens)。然后将这些“Summary” tokens接到query,key,value上,以便query能够直接attend到带有“Summary” tokens的key。其中[公式]为简单的空间上的平均,如下所示:

综上,作者将具有时空混合注意和轻量级全局时间注意(或Summary token)的视频Transformer称为X-ViT 。

4.实验

4.2消融研究

4.2.1. 局部窗口大小的影响

上表显示了通过改变在时空混合注意中使用的局部窗口大小[公式]而得到的模型的准确性。可以看出[公式]时效果是最好的。

4.2.2. SA位置的影响

作者探索了在Transformer内,时空混合注意操作应该应用于哪些层比较好。上表显示了实验的结果,可以看出,在所有层都应用效果最好。

4.2.3. 时间聚合的影响

作者比较了两种用于时间聚合的方法:简单的时间平均和所提出的时间注意(TA)机制。上表显示了不同层的影响,可以看到,加了一层就能将Top-1准确率从62.4%提高到64.4%。

4.2.4. 时空混合qkv向量的影响

上表显示了时空混合qkv向量相比于不进行混合的向量,在最中的性能上能有很大的改进。

4.2.5. 时空混合量的影响

上表展示了不同混合通道的数量的实验结果。可以看出,50%的通道混合量是一个比较好选择。

4.2.6. Summary token的影响

作者在两个数据集比较了Summary token的影响,可以看出Summary token在Kinetics-400数据集上表现的较好。

4.2.7. 输入帧数的影响

从上表可以看出,将帧数从8帧增加到16帧,SSv2上Top-1的准确率提高了1.8%。

上表为 Kinetics-400的结果,结论相似。

上表为Epic-100的结果,结论相似。

4.2.8. 标记数量的影响

即使token的数量显著减少(ViT-B/32),本文的方法仍然能够达到不错的准确性。

4.2.9. 测试时Crop数量的影响

从上表可以看出,crops数量的增加,在2之后对最终结果影响不大。

4.3与SOTA的比较

上表展示了本文方法在 Kinetics-600数据集上和其他SOTA方法的对比。

5. 总结

在本文中,作者提出了一种用于视频识别的全时空注意近似。本文的近似使得模型具有与仅空间注意相同的计算成本,比最近提出的视频Transformer更有效。具体来说,作者做了两个方面的近似:(a)将时间注意力限制在一个局部时间窗口 ,并利用Transformer的深度来获得视频序列的全时间覆盖;(b)使用有效的时空混合来联合建模空间和时间信息 ,而不在仅空间注意模型上产生任何额外的计算成本。

参考文献

[1]. Gedas Bertasius, Heng Wang, and Lorenzo Torresani. Is space-time attention all you need for video understanding? arXiv preprint arXiv:2102.05095, 2021.

转载请注明:《《X-ViT》-基于时空混合attention的视频Transformer,大幅度降低计算复杂度