©作者 | 黎昆昌

單位 | 上海AILab

研究方向 | 影片行為識別

介紹

我們提出了一個僅基於狀態空間模型 (SSM) 的高效影片理解架構 VideoMamba，並透過大量的實驗證明了它具備一系列良好的特性，包括 1）Visual Domain Scalability；2）Short-term Action Sensitivity；3）Long-term Video Superiority；4）Modality Compatibility。這使得 VideoMamba 在一系列影片 benchmark 上取得不俗的結果，尤其是長影片 benchmark，為未來更全面的影片理解提供了更高效的方案。

論文標題：

VideoMamba: State Space Model for Efficient Video Understanding

論文連結：

https://arxiv.org/abs/2403.06977

程式碼連結：

https://github.com/OpenGVLab/VideoMamba

模型連結：

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 對影片理解是否同樣有效？答案是肯定的。

Method

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 策略，使用訓練好的小模型當老師，引導大模型訓練，有效地避免模型過擬合，而只需少量額外的開銷。

掩碼建模

近來，VideoMAE [11] 引入掩碼建模，顯著增強了模型對細粒度時序的理解能力，而 UMT [12] 進一步提出高效的掩碼對齊策略，不僅大大減小了訓練開銷，還使得模型能魯棒地處理各種單模態和多模態任務。為了增強 VideoMamba 對時序的敏感性，同時驗證它和文字模態的相容性，我們借鑑 UMT 的方式，引入 CLIP-ViT 當 teacher，進行兩階段蒸餾訓練。

不同於 UMT 使用多層對齊，由於 VideoMamba 和 ViT 存在架構差異，我們只對齊模型最後一層，考慮到 Mamba block 對連續 token 更友好，我們設計了逐行掩碼策略：

同時我們也考慮了注意力掩碼策略，這能保持語義性更強的鄰近 token。

Experiments

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。隨著模型規模和解析度放大，效能穩定提升。

Short-term Video Understanding

在上述 K400 和 SthSthV2 的短影片分類任務中，我們同樣觀察到 VideoMamba 良好的放縮性，且顯著優於基於注意力的影片模型如 TimeSformer 和 ViViT，與結合卷積和自注意力的 UniFormer 效能相當。再者，在引入掩碼訓練後，VideoMamba 效能顯著提升，在細粒度動作分類 SthSthV2 資料集上，顯著好於基於 ViT 的 UMT。

進一步的消融實驗表明，spatial-first 掃描方案效果最好。不同於 ImageNet 上效能隨解析度逐漸提升，影片資料集上解析度對效能影響有限，而幀數對效能影響明顯。對於掩碼建模，逐行掩碼優於隨機掩碼策略，且注意力掩碼策略最有效；對齊最後一層效果最好；合適的掩碼比例和 Droppath 能較好提升訓練效果。

Long-term Video Understanding

我們在 Breakfast，COIN 和 LVU 上評估了 VideoMamba 對長時影片的理解能力，相較於以往 feature-based 的方法，VideoMamba 僅需要輸入稀疏取樣的 32-64 幀，效果便大幅領先，且模型規模更小。

Multi-modality Video Understanding

我們將 VideoMamba 和 BERT 連線，構造多模態模型，並使用大規模多模態資料進行預訓練，在多個影片文字檢索任務上進行了效能評估。實驗揭示 VideoMamba 同樣能很好地作為多模態的視覺編碼器，隨著預訓練資料的增加，能持續提升多模態理解的能力，且由於以 ViT 為視覺編碼器的 UMT，尤其是在包含長影片（ANet 和 DiDeMo）和更復雜場景（LSMDC）的資料集上。