VideoMamba：用于高效视频理解的状态空间模型

论文题目： 

VideoMamba: State Space Model for Efficient Video Understanding 

论文链接：

​​https://arxiv.org/abs/2403.06977​​

代码链接：

​​https://github.com/OpenGVLab/VideoMamba​​

Huggingface：

​​https://huggingface.co/OpenGVLab/VideoMamba​​

Online Demo：

​​https://huggingface.co/spaces/OpenGVLab/VideoMamba​​

一、Motivation

视频表征学习长期以来存在两大痛点，一是短clip里存在大量的时空冗余，二是长上下本需要复杂的时空关联。曾经风靡一时的3D CNN和video transformer，分别使用卷积和自注意力机制解决了两大难题。在我们之前的工作UniFormer[1]里，我们尝试将卷积和自注意力无缝地结合，尽管它能同时解决两大难题，但对于长视频仍力不从心。而Gemini[2]和Sora[3]的爆火，使得长视频理解与生成成为了研究的重心，这亟需更高效的视频表征模型。

幸运的是，NLP领域这两年涌现了不少高效算子，如S4[4], RWKV[5]和RetNet[6]。而Mamba[7]提出动态状态空间模型(S6)，能以线性复杂度进行长时的动态建模。这引领了一系列视觉任务的适配，如Vision Mamba[8]和VMamba[9]，提出了多向SSM机制用于处理2D图片，这些模型不仅能与基于注意力的架构媲美，而且大大减小显存开销。

考虑到视频产生的超长token序列，一个自然而然的问题便是，Mamba对视频理解是否同样有效？答案是肯定的。

二、 Methods

2.1 Architecture

在进入VideoMamba结构的介绍之前，我们先看看用于1D序列的Mamba block，和用于视觉任务的双向Mamba block。这里我们不再赘述SSM和Mamba的底层原理，感兴趣的同学可以通过油管视频（https://www.youtube.com/watch?v=8Q_tqwpTpVU）学习。

双向Mamba在单向Mamba的基础上，引入了对反向序列的SSM，这使得双向Mamba能更好地对2D序列建模，从而提升对视觉输入的感知能力。基于双向Mamba，我们按照ViT[10]的设计，引入[CLS] token和空间位置编码，并针对视频建模，引入3D patch embedding和空间位置编码，提出了如下所示的VideoMamba：

为了应用双向Mamba处理时空信息，我们拓展原本的2D扫描到不同的双向3D扫描：

其中空间优先扫描最简单，实验证明效果也最好。基于该架构，我们提出了三种不同size的模型，VideoMamba-Ti，VideoMamba-S和VideoMamba-M。

但在实验里，当我们增大VideoMamba规模时，非常容易过拟合，导致大模型的结果甚至差于小模型。为此，我们提出了Self-Distillation策略，使用训练好的小模型当老师，引导大模型训练，有效地避免模型过拟合，而只需少量额外的开销。

2.2 Masked Modeling

近来，VideoMAE[11]引入掩码建模，显著增强了模型对细粒度时序的理解能力，而UMT[12]进一步提出高效的掩码对齐策略，不仅大大减小了训练开销，还使得模型能鲁棒地处理各种单模态和多模态任务。为了增强VideoMamba对时序的敏感性，同时验证它和文本模态的兼容性，我们借鉴UMT的方式，引入CLIP-ViT当teacher，进行两阶段蒸馏训练。

不同于UMT使用多层对齐，由于VideoMamba和ViT存在架构差异，我们只对齐模型最后一层，考虑到Mamba block对连续token更友好，我们设计了逐行掩码策略：

同时我们也考虑了注意力掩码策略，这能保持语义性更强的邻近token。

三、Experiments

3.1 Scale Up

我们首先在ImageNet上进行了图像分类实验如下所示

可见在没有Self-Distillation (SD)时，VideoMamba-M和VideoMamba-B都会在训练的最后过拟合，其中VideoMamba-B尤为严重。而在引入SD后，VideoMamba-M收敛符合期望，且明显强于老师模型VideoMamba-S。为了避免老师模型带偏训练，我们引入了Early Stop策略，即提前移除蒸馏引导，实验发现并无提升。完整ImageNet对比如下：

和无层次化结构的其他模型相比，VideoMamba优于其他CNN和ViT的模型，如ConvNeXt和ViT。随着模型规模和分辨率放大，性能稳定提升。

3.2 Short-term Video Understanding

在上述K400和SthSthV2的短视频分类任务中，我们同样观察到VideoMamba良好的放缩性，且显著优于基于注意力的视频模型如TimeSformer和ViViT，与结合卷积和自注意力的UniFormer性能相当。再者，在引入掩码训练后，VideoMamba性能显著提升，在细粒度动作分类SthSthV2数据集上，显著好于基于ViT的UMT。

进一步的消融实验表明，spatial-first扫描方案效果最好。不同于ImageNet上性能随分辨率逐渐提升，视频数据集上分辨率对性能影响有限，而帧数对性能影响明显。对于掩码建模，逐行掩码优于随机掩码策略，且注意力掩码策略最有效；对齐最后一层效果最好；合适的掩码比例和Droppath能较好提升训练效果。

3.3 Long-term Video Understanding

我们在Breakfast，COIN和LVU上评估了VideoMamba对长时视频的理解能力，相较于以往feature-based的方法，VideoMamba仅需要输入稀疏采样的32-64帧，效果便大幅领先，且模型规模更小。

3.4 Multi-modality Video Understanding

我们将VideoMamba和BERT连接，构造多模态模型，并使用大规模多模态数据进行预训练，在多个视频文本检索任务上进行了性能评估。实验揭示VideoMamba同样能很好地作为多模态的视觉编码器，随着预训练数据的增加，能持续提升多模态理解的能力，且由于以ViT为视觉编码器的UMT，尤其是在包含长视频（ANet和DiDeMo）和更复杂场景（LSMDC）的数据集上。

四、Conclusion

我们提出了仅基于状态空间模型的视频理解架构VideoMamba，全面的实验表明VideoMamba对视频理解具有一系列良好特性，我们希望它可以为未来长视频的表征学习指明道路。

References

[1] UniFormer: 

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

[2] Gemini: 

​​https://blog.google/technology/ai/google-gemini-next-generation-model-february-2024/​​

[3] Sora: ​​https://openai.com/sora​​

[4] S4: ​​https://github.com/state-spaces/s4​​

[5] RWKV: ​​https://www.rwkv.com/​​

[6] RetNet: 

​​https://github.com/microsoft/unilm/tree/master/retnet​​

[7] Mamba: ​​https://github.com/state-spaces/mamba​​

[8] Vision Mamba: ​​https://github.com/hustvl/Vim​​

[9] VMamba: ​​https://github.com/MzeroMiko/VMamba​​

[10] ViT: 

​​https://github.com/google-research/vision_transformer​​

[11] VideoMAE: 

​​https://github.com/MCG-NJU/VideoMAE​​

[12] UMT: 

​​https://github.com/OpenGVLab/unmasked_teacher​​

Illustration From IconScout By 22

本文转载自​​将门创投​​，作者：黎昆昌 ​​​​