1. 从Transformer到Mamba:视频理解为何需要新架构?
最近在复现和测试一些视频理解相关的模型时,我重新仔细读了一遍《VideoMamba: State Space Model for Efficient Video Understanding》这篇论文。说实话,第一次看到“Mamba”这个词出现在视频领域时,我的第一反应是好奇和怀疑。毕竟,Transformer架构在视觉、尤其是视频理解领域,经过VIT、TimeSformer等工作的锤炼,几乎已经成为事实上的标准。大家习惯了用自注意力机制(Self-Attention)去建模空间和时序上的长距离依赖。那么,为什么还需要一个基于状态空间模型(State Space Model, SSM)的VideoMamba?它到底解决了什么Transformer解决不了,或者解决起来很“疼”的问题?
要理解VideoMamba的价值,我们得先回到视频理解任务本身的核心挑战上。视频是一系列在时间上连续的图像帧,其数据量巨大。一个几秒钟的短视频,按30帧/秒算,就是上百张高分辨率图片。Transformer的自注意力机制计算复杂度是序列长度的平方级(O(N²))。对于图像,我们可以把图像块(patch)展平成一个序列,这个序列长度已经不小了(例如224x224的图像切成16x16的块,序列长度是196)。对于视频,如果我们简单地把所有帧的所有图像块拼接成一个超长序列,那么序列长度N = 帧数T * 每帧图像块数L。这个N会变得极其庞大,导致计算和内存开销完全无法承受。
因此,现有的视频Transformer(如TimeSformer)通常采用“因式分解”的注意力机制,比如将空间注意力和时间注意力分开。这虽然降低了计算量,但也引入了一个很强的先验假设:空间和时间依赖是可以完全解耦的。然而,现实世界中的动作和事件,空间和时间信息往往是紧密交织、相互影响的。一个挥手的动作,其空间轨迹(手的位置变化)本身就定义了时间信息。强行分开建模,可能会损失这种联合表征的能力。
另一方面,状态空间模型(SSM),特别是像Mamba这样的选择性状态空间模型,提供了一种完全不同的思路。它本质上是一个线性时不变(LTI)系统,通过一个隐藏状态h_t来递归地处理序列信息,其核心计算是线性递归,复杂度是线性的O(N)。这意味着,无论视频有多长,理论上Mamba都能以线性增长的计算成本来处理整个序列,这是对Transformer平方复杂度瓶颈的一个根本性突破。更重要的是,Mamba引入了“选择性”机制,让模型能够根据输入内容动态地决定哪些信息需要被记住(持久化到状态中),哪些可以忽略。这对于视频理解至关重要,因为视频中充斥着大量冗余信息(如静态背景),关键的动作变化往往只发生在局部区域和短暂时刻。
所以,VideoMamba的出现,并不是为了取代Transformer,而是为了在视频这个对长序列建模效率和全局依赖捕捉都有极高要求的特定领域,提供一个更优雅、更高效的解决方案。它试图用线性复杂度的SSM,去捕获视频中全局的时空依赖,同时避免因式分解注意力可能带来的信息损失。接下来,我们就深入论文,拆解VideoMamba是如何具体实现这一目标的。
2. VideoMamba的核心设计:将图像Mamba扩展到时空领域
VideoMamba的整体架构设计非常清晰,其核心思想是:将针对图像设计的Mamba模型,以一种尽可能自然的方式,扩展到视频数据上。它没有重新发明轮子,而是巧妙地利用了Mamba在图像分类(如Vision Mamba)上已经证明有效的设计,并针对视频的时序特性进行适配。
2.1 输入编码:从2D图像块到3D时空立方体
和Vision Transformer(ViT)一样,VideoMamba的第一步也是将输入视频进行“分块”。假设输入视频大小为 T x H x W x 3(T帧,高H,宽W,3通道)。VideoMamba将其划分为不重叠的时空立方体(Spatio-Temporal Cubes)。每个立方体的大小是 t x p x p,其中t是时间维度的跨度,p是空间维度的跨度。例如,一个常见的设置是 t=2, p=16。这意味着,模型将每连续2帧图像,在空间上切分成16x16的小块,这样一个立方体就包含了2帧内同一个空间位置上的图像信息。
这个过程可以理解为在时空维度上做了一次均匀采样。然后将每个立方体展平,并通过一个线性投影层映射到一个D维的嵌入向量。同时,会加上标准的可学习位置编码(或固定位置编码),以保留时空位置信息。最终,我们得到了一个长度为 N = (T/t) * (H/p) * (W/p) 的令牌(token)序列。这个序列同时包含了空间和时间的信息,并且其长度相比“帧级拼接”的方案已经大大减少,为后续的线性复杂度建模奠定了基础。
这里有一个关键的设计考量:时间下采样(t>1)。在Vision Mamba中,处理单张图像,序列完全由空间块构成。而在VideoMamba中,通过引入t>1,我们在建模的一开始就对时间维度进行了轻度聚合。这样做的好处是直接缩短了需要建模的序列长度,降低了计算负担。风险在于可能会损失一些非常细微的时间动态(毫秒级变化)。但论文中的实验表明,对于大多数以识别和理解为主的任务(如动作识别),t=2是一个在效率和效果上很好的平衡点。
2.2 骨干网络:Video Mamba Block的堆叠
得到令牌序列后,它们被送入一系列堆叠的Video Mamba Block中。每个Block的结构继承了Mamba的核心组件,并做了视频化的调整:
- 归一化层(Norm):首先对输入序列进行层归一化(LayerNorm)。
- 1x1卷积前向投影:这是一个可选的、但实践中非常有效的设计。在进入SSM核心计算前,使用一个1x1的卷积(等价于全连接层)对令牌序列进行通道混合和升维。这有助于模型融合不同通道的信息。
- 深度卷积(Depthwise Convolution):接着是一个深度可分离卷积,其作用是进行局部的时空特征交互。这是从ConvNeXt等现代卷积网络中借鉴来的思想,能为模型提供强大的局部归纳偏置(inductive bias)。在视频中,这意味着相邻的时空立方体(在时间和空间上都相邻)的信息会先进行一步混合,帮助SSM更好地理解局部上下文。
- 激活与投影:经过卷积后,通过SiLU激活函数,并再次使用1x1卷积进行投影。
- 选择性状态空间模型(Selective SSM):这是整个Block的心脏。经过前述处理的序列,被送入Mamba的SSM层。如前所述,SSM通过一个隐藏状态h_t进行线性递归计算。其“选择性”体现在系统参数(如离散化步长Δ)是由输入数据通过一个线性层动态生成的。这使得模型能够根据当前输入的时空内容,决定以多大的“步幅”或“分辨率”来更新其内部状态。对于视频,这可以理解为:当画面静止或变化缓慢时(如背景),模型选择“长时记忆”,状态更新慢;当画面出现快速、重要的动作时(如人物突然举手),模型选择“短时记忆”,状态快速更新以捕捉细节。
- 残差连接:SSM的输出与Block的原始输入进行残差相加,构成这个Block的最终输出。
这一系列Block堆叠起来,就构成了VideoMamba的骨干网络。信息在这个网络中流动时,一方面通过深度卷积进行快速的局部时空交互,另一方面通过选择性SSM进行高效的全局长程依赖建模。两者互补,形成了强大的特征提取能力。
注意:这里容易产生一个误解,认为SSM完全替代了注意力。实际上,在VideoMamba Block中,并没有显式的注意力计算。SSM通过其线性递归和选择性机制,隐式地实现了对序列中重要信息的聚焦和长距离信息的传递,这是一种与注意力机制机理不同但目标相似的建模方式。
2.3 输出头与任务适配
骨干网络提取的时空特征序列,最后会根据下游任务的不同,被送入不同的输出头。
- 视频分类:这是最直接的任务。通常采用全局平均池化(Global Average Pooling)将整个序列聚合为一个全局特征向量,然后通过一个全连接层分类器输出类别概率。
- 时序动作定位:这类任务需要模型不仅知道视频中发生了什么,还要知道动作发生的起止时间。VideoMamba可以输出每个时序位置(对应一组时空立方体)的特征,然后接上专门的时间边界检测头(例如,基于锚点或Transformer的解码头)来预测动作片段。
- 时空动作检测:这需要同时定位动作发生的时间和空间位置(在视频的哪一帧、哪个区域)。由于VideoMamba的令牌序列本身就保留了时空网格结构,理论上可以通过对输出序列进行重塑和上采样,得到逐像素或逐区域的特征图,再接入检测头(如Faster R-CNN的RPN)。
论文中主要验证了其在视频分类任务上的有效性,但这套架构的设计是通用的,为扩展到更复杂的视频理解任务留下了接口。
3. 选择性SSM在视频中的工作原理与优势
要真正理解VideoMamba为何有效,我们必须深入其灵魂——选择性状态空间模型(Selective SSM),并看看它在处理视频数据时,是如何发挥作用的。
3.1 状态空间模型(SSM)的直观理解
我们可以用一个非常简单的类比来理解SSM:把它想象成一个有“记忆”的系统。这个系统在每个时间步(对应视频中的一个时空令牌)接收一个输入x_t,然后根据当前的内部状态h_t和输入,更新自己的状态到h_{t+1},并产生一个输出y_t。其数学形式通常表示为:
h_{t+1} = A * h_t + B * x_t y_t = C * h_t + D * x_t其中A, B, C, D是系统的参数矩阵。A决定了状态如何自我演化(衰减或保持),B决定了当前输入如何影响状态,C决定了如何从状态生成输出,D是输入到输出的直接通路。
这个系统的关键特性是线性递归。计算h_{t+1}只需要知道h_t和x_t,与之前所有的历史输入{x_1, x_2, ..., x_{t-1}}无关,因为这些历史信息已经被“压缩”到了状态h_t中。这使得它的计算可以按时间步顺序进行,复杂度是O(N)。同时,由于状态h_t理论上可以携带从序列开始到当前的所有信息(尽管受A矩阵影响会有衰减),因此它具备建模长距离依赖的潜力。
3.2 “选择性”为何是视频建模的关键
传统SSM(如S4模型)的参数A, B, C是固定的,与输入无关。这意味着它对所有输入序列都采用相同的“记忆策略”。这在处理像视频这样内容高度动态、信息重要性分布极不均匀的数据时,显得不够灵活。
Mamba的核心创新就是引入了“选择性”(Selectivity)。具体来说,模型参数B, C和最重要的离散化参数Δ,都是由当前时间步的输入x_t通过一个线性变换动态生成的:
Δ_t, B_t, C_t = f_θ(x_t)这意味着:
- Δ_t(步长):控制着系统状态更新的“速度”或“分辨率”。Δ_t大,意味着系统将当前输入的影响快速“写入”长期状态(状态更新快);Δ_t小,则意味着系统更“惰性”,当前输入对长期状态影响小。在视频中,面对快速变化的动作帧,模型可以生成大的Δ_t,迅速更新状态以记住这个关键变化;面对静止的背景帧,则生成小的Δ_t,避免状态被无关信息污染。
- B_t, C_t:动态的B和C使得系统能够根据输入内容,决定“记住什么”(B控制输入如何影响状态)以及“输出什么”(C控制状态如何贡献给输出)。这实现了内容感知的信息路由。
这种选择性机制,让VideoMamba像一个拥有动态焦点的观察者。它不再平等地对待视频中的每一帧每一个区域,而是能够主动地、根据内容重要性,分配其有限的“注意力”和“记忆资源”。这与Transformer的自注意力在功能上异曲同工——都是让模型关注重要的部分。但实现路径不同:自注意力是通过所有令牌两两计算相似度(平方复杂度)来实现显式的聚焦;而选择性SSM是通过动态调整线性递归系统的参数,来实现隐式的、高效的内容选择。
3.3 与Transformer在视频建模上的对比
我们可以从几个维度对比VideoMamba和基于Transformer的视频模型:
- 计算复杂度:这是最显著的差异。Transformer自注意力是O(N²),而SSM是O(N)。对于长视频序列,这个优势是指数级的。这使得VideoMamba能够处理更长的视频片段,或者使用更高的时空分辨率,而无需复杂的因式分解或窗口化技巧。
- 建模能力:
- 全局依赖:标准的Transformer(全局注意力)能建模序列中任意两个位置的关系。VideoMamba通过状态的递归传递,理论上也能建立任意长距离的依赖,但其信息传递是沿着时序单向进行的(因果性),不像注意力那样是“全连接”的。对于视频理解,这种单向性通常不是问题,因为时间本身就有方向。
- 时空联合建模:因式分解的Transformer(如TimeSformer)将空间和时间注意力分离。VideoMamba的输入是时空令牌的混合序列,其SSM和深度卷积天然地在同一套参数下对时空信息进行联合处理,可能更有利于捕捉时空不可分的特征。
- 训练与推理:
- 训练:SSM的线性递归特性使其难以进行高效的并行训练(因为第t步依赖第t-步的状态)。Mamba通过其“扫描”(Scan)算法和硬件优化的实现,巧妙地解决了这个问题,实现了训练时的并行化。
- 推理:SSM在推理时具有内存效率高的优势。由于其状态是固定维度的向量,与序列长度无关,因此处理极长视频时,内存占用几乎恒定。而Transformer在推理时(尤其是自回归生成或需要缓存键值对时),内存消耗会随序列长度增长。
- 归纳偏置:Transformer的注意力机制是一种非常通用的关系建模器,但需要大量数据来学习有用的模式。SSM(尤其是结合了深度卷积的VideoMamba)内置了局部性(通过卷积)和序列性(通过递归)的归纳偏置,这可能使其在数据量不是极端庞大的情况下,更容易训练和泛化。
下表总结了核心对比:
| 特性 | VideoMamba (基于选择性SSM) | 标准视频Transformer (如ViViT) | 因式分解视频Transformer (如TimeSformer) |
|---|---|---|---|
| 核心操作 | 选择性线性递归 + 深度卷积 | 全局自注意力 | 空间注意力 + 时间注意力(分离) |
| 时空建模 | 联合建模(时空令牌混合) | 联合建模(时空令牌混合) | 分离建模 |
| 计算复杂度 | O(N) | O(N²) | O(N_s² + N_t²), 其中N_s, N_t分别是空间、时间序列长度 |
| 长序列处理 | 非常高效,内存占用恒定 | 代价高昂 | 相对高效,但可能损失全局时空交互 |
| 内容感知 | 通过动态参数实现隐式选择 | 通过注意力权重实现显式选择 | 通过注意力权重实现显式选择 |
| 主要优势 | 线性复杂度,适合极长视频,内存高效 | 强大的全局关系建模能力 | 平衡了计算开销和建模能力 |
| 主要挑战 | 需要理解递归和选择性机制,并行化实现复杂 | 计算和内存成本高 | 时空解耦的假设可能不总是成立 |
4. 实验设置、结果分析与实战启示
论文在多个标准的视频理解基准数据集上对VideoMamba进行了评估,主要包括Kinetics-400、Something-Something V2、Charades等。这些数据集涵盖了从大规模互联网视频分类(Kinetics)到精细时序动作理解(Something-Something)等不同场景。
4.1 主要实验结果与发现
- 在主流基准上达到有竞争力的性能:在Kinetics-400上,VideoMamba与基于Transformer的同类模型(如VideoSwin, MViT)相比,在参数量和计算量(FLOPs)相近的情况下,取得了相当甚至略优的精度。这证明了选择性SSM架构在视频分类任务上的有效性,它并非一个“理论玩具”,而是具备实战能力的模型。
- 在长视频理解上展现潜力:论文特别强调了VideoMamba在处理长序列视频上的优势。在一些需要更长时序上下文的任务或数据集上,由于其线性复杂度,VideoMamba能够以更低的计算成本处理更多的帧数,从而获得更好的性能。例如,在需要理解较长流程动作的数据集上,增加输入帧数带来的性能提升比Transformer模型更为明显。
- 效率优势:在相同的硬件条件下(如相同的GPU内存),VideoMamba能够处理更高分辨率或更长时长的视频输入。这对于实际部署,特别是资源受限的边缘场景,是一个重要的优势。
- 消融实验的启示:
- 时空立方体编码的重要性:实验表明,采用时空立方体(t>1)作为输入单元,比单纯在帧级应用图像Mamba(t=1)效果更好,且效率更高。这验证了在输入阶段进行轻度时序聚合的有效性。
- 深度卷积的作用:移除深度卷积层会导致性能下降,尤其是在Something-Something V2这种对细粒度时空关系要求高的数据集上。这说明局部归纳偏置对于视频SSM模型是必要的补充。
- 选择性的关键性:将Mamba替换为没有选择性的传统SSM(如S4),模型性能显著下降。这直接证明了“内容感知”的动态机制对于视频建模至关重要。
4.2 从论文到实践:复现与使用的思考
如果你和我一样,对VideoMamba感兴趣并想动手尝试,以下是一些基于论文和当前开源生态的实战思考:
1. 环境搭建与依赖:VideoMamba基于PyTorch实现,其核心依赖于Mamba官方库(mamba-ssm)。安装时需要注意CUDA版本兼容性。由于Mamba底层涉及CUDA内核的定制,对PyTorch和CUDA工具链的版本匹配要求较高。建议在干净的虚拟环境(如conda)中,严格按照官方仓库的说明进行安装。
# 示例,具体版本请参考最新官方文档 conda create -n videomamba python=3.10 conda activate videomamba pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 pip install causal-conv1d==1.1.1 # 依赖 pip install mamba-ssm # 核心SSM层 # 然后克隆VideoMamba代码仓库并安装其他依赖2. 数据预处理流程:视频理解模型的数据预处理管线通常比较重。你需要准备:
- 视频解码:使用Decord、PyAV或OpenCV从视频文件中抽帧。
- 时空采样:按照论文设定,均匀采样T帧(如16帧)。这里要注意采样策略(密集采样适合短时动作,稀疏采样适合长时动作)。
- 空间裁剪与缩放:通常采用随机裁剪(训练时)或中心裁剪(测试时)到固定分辨率(如224x224)。
- 构建时空立方体:这是VideoMamba特有的步骤。你需要将采样得到的T帧,在时间维度以步长t进行分组,并在空间上切割成p x p的块,最终形成令牌序列。这个过程可以集成在数据加载器中。
3. 模型训练技巧:
- 优化器与学习率:论文采用AdamW优化器,并配合余弦退火学习率调度器。这是视觉Transformer训练的标配。
- 正则化:权重衰减、DropPath(Stochastic Depth)是防止过拟合的关键。VideoMamba作为序列模型,可能也需要在SSM层应用特定的归一化或正则化技术。
- 混合精度训练(AMP):强烈推荐使用自动混合精度训练,这能显著减少GPU内存占用并加快训练速度,尤其对于VideoMamba这种包含自定义CUDA算子的模型。
- 长视频训练:如果你想利用VideoMamba处理超长视频的优势,可能需要调整优化策略。例如,使用梯度检查点(Gradient Checkpointing)来节省内存,以便使用更大的批次大小或更长的序列。
4. 可能遇到的坑与解决方案:
- CUDA内存溢出(OOM):即使SSM理论内存效率高,但视频数据本身很大。如果遇到OOM,首先检查输入帧大小(分辨率)和批次大小(batch size)。可以尝试梯度累积(Gradient Accumulation)来模拟更大的批次。
- 训练不稳定:SSM模型在训练初期可能比Transformer更敏感。确保使用了正确的初始化方法(如Mamba层特定的初始化),并监控训练损失曲线。适当降低初始学习率可能会有帮助。
- 自定义任务适配:论文主要提供了分类任务的代码。如果你要用于检测或分割,需要自行设计输出头。这里的关键是如何将输出的1D令牌序列重新映射回2D(空间)或3D(时空)的特征图。可以参考Vision Transformer中用于分割的“序列到空间”重塑方法。
5. 与现有生态的集成:目前,VideoMamba可以作为一个强大的视频特征提取器(backbone)集成到更大的系统中。例如,你可以:
- 在视频-语言多模态任务中,用VideoMamba替换掉原来的视频编码器(如CLIP的视觉编码器),来获取更高效、更强大的视频表征。
- 将其用于视频生成任务(如扩散模型)的条件编码部分,利用其长序列建模能力来提供更连贯的时序条件。
- 在工业界的视频内容分析平台中,作为动作识别、异常检测等模块的核心网络,其高效率特性有利于降低服务器成本。
VideoMamba为我们打开了一扇新的大门,它展示了状态空间模型在视觉乃至多模态领域的巨大潜力。它不是一个终点,而是一个起点。围绕如何更好地设计视频专用的SSM块、如何与注意力机制进行更高效的融合(例如混合架构)、如何将其应用于更复杂的视频生成和理解任务,还有大量的探索空间。对于从业者来说,理解其原理,掌握其用法,并将其应用到合适的场景中,是当前阶段最务实的选择。在我自己的实验中,将其用于一些长视频的摘要生成任务,初步看到了其在保持上下文连贯性上的优势,这让我对这条技术路线的未来充满期待。