1. 项目概述：当视频理解遇上状态空间模型

最近在视频理解领域，一个名为VideoMamba的新模型架构引起了不小的关注。作为一名长期关注计算机视觉前沿动态的从业者，我习惯性地去研读了它的论文，并动手进行了复现和测试。简单来说，VideoMamba的核心思路，是将近年来在语言和长序列建模中表现抢眼的 状态空间模型 ，首次系统性地引入到视频理解任务中，试图解决传统视觉Transformer在处理长视频时面临的计算效率和长程依赖建模的瓶颈。

视频数据天然具有时空冗余性，一秒钟30帧的画面里，相邻帧之间的变化可能微乎其微。传统的基于Transformer的方法，如TimeSformer或Video Swin Transformer，需要为每一帧的每一个图像块计算自注意力，这导致了计算复杂度随视频长度和空间分辨率呈二次方甚至三次方增长。处理几分钟的长视频，对显存和算力都是巨大的挑战。VideoMamba的提出，正是瞄准了这个痛点。它借鉴了Mamba在语言模型中的成功经验，利用SSM的选择性扫描机制，实现了对视频序列的线性复杂度建模，同时保持了全局感受野。这听起来很美好，但实际效果如何？在复现和实验的过程中，我发现了许多论文图表之外的有趣细节和实操陷阱。这篇文章，我就结合自己的笔记和代码实践，来深入聊聊VideoMamba的设计精髓、复现要点以及它给视频理解领域带来的新思考。

2. 核心思路拆解：为什么是SSM？为什么是现在？

2.1 视频建模的“老大难”问题

在深入VideoMamba之前，我们必须先理解它要解决的核心问题。视频理解，尤其是动作识别、长视频理解等任务，本质上是处理一个高维的时空信号。早期的3D卷积网络通过堆叠3D卷积核来同时捕捉空间和时序信息，但计算量巨大，且感受野有限。ViT的横空出世，将图像切分为块，通过自注意力机制实现了强大的全局建模能力。将其扩展到视频领域，最直观的想法就是引入时序维度，形成“时空块”，然后应用Transformer编码器。

然而，问题随之而来。假设我们将视频处理为 T×H×W 个时空块，那么自注意力机制的计算复杂度是 O((T×H×W)^2) 。即使采用分治策略，如将空间和时间注意力分离（像TimeSformer那样先做空间注意力再做时间注意力），其复杂度也是 O(T*(H*W)^2 + H*W*T^2) ，仍然难以处理很长的 T 。当视频长度达到数百甚至上千帧时，显存消耗会变得不可承受。此外，标准的自注意力机制是“内容感知”但“位置盲”的，它通过计算所有元素对的相似度来聚合信息，但对视频中至关重要的时序顺序和因果关系的建模并不直接。

2.2 SSM的“入场券”：线性复杂度与选择性

状态空间模型为上述问题提供了一个颇具潜力的解决方案。SSM本质上描述了一个将输入序列 x(t) 通过一个隐含状态 h(t) 映射到输出序列 y(t) 的连续系统。当这个连续系统被离散化后，可以转化为一个类似于RNN的递归计算形式，或者利用其线性时不变特性，通过全局卷积的方式进行高效并行训练。

Mamba模型的突破在于引入了“选择性”机制。传统的SSM参数（如状态转移矩阵）是输入无关的，而Mamba让这些参数成为输入的函数。这意味着，模型可以动态地决定记住哪些信息、忽略哪些信息，从而实现了类似注意力机制的“内容感知”能力。其核心的扫描操作（Scan）复杂度是线性的 O(T) ，这相比Transformer的二次方复杂度是巨大的优势。

对于视频数据，这种选择性扫描机制具有天然的吸引力：

1. 处理长序列 ：线性复杂度使得处理长视频片段成为可能，无需进行激情的下采样或分段处理。
2. 建模时空动态 ：选择性机制可以让模型专注于视频中发生显著变化的帧和区域，忽略静态背景，这与人类观看视频的直觉相符。
3. 因果建模 ：SSM的递归形式天然适合建模时序上的因果关系，这对于动作预测、视频生成等任务很重要。

VideoMamba正是抓住了这一点，将Mamba块作为核心构建模块，替代了Transformer中的自注意力层，构建了一个纯SSM架构的视频基础模型。

注意 ：这里有一个关键理解点。SSM的线性复杂度优势在 训练 的并行计算阶段（使用卷积模式）和 推理 的递归阶段都非常明显。但在训练时，为了利用硬件并行性，需要将递归展开为卷积，这会消耗与序列长度成正比的显存（尽管计算是线性的）。不过，其显存增长是线性的 O(T) ，而Transformer是二次方的 O(T^2) ，这在长序列下依然有巨大优势。

2.3 VideoMamba的整体架构蓝图

VideoMamba的架构设计非常简洁清晰，遵循了“分而治之”的视觉编码思想，可以概括为以下几个步骤：

1. 输入切分与嵌入 ：和ViT一样，将输入视频（形状为 T×C×H×W ）沿空间维度切分成不重叠的块（Patch）。每个块经过线性投影后，得到一个时空令牌（Token）。同时，会添加一个可学习的分类令牌（[CLS]）用于最终分类。此外，为了保留位置信息，会添加标准的可学习一维位置编码（针对时空序列）。
2. VideoMamba编码器 ：这是模型的核心。编码器由L个相同的VideoMamba块堆叠而成。每个块的结构非常类似于Transformer的预归一化设计：
   • 输入 ：令牌序列 X 。
   • 步骤一 ：层归一化（LayerNorm）。
   • 步骤二 ：通过一个残差分支，先经过一个线性投影，然后进入 双向SSM（Bi-SSM） 层。这是与语言Mamba最大的不同之一。为了充分捕捉视频上下文信息，VideoMamba采用了双向扫描，即同时进行前向和后向扫描，然后将两个方向的输出相加。这允许每个令牌都能感知整个序列的上下文。
   • 步骤三 ：另一个残差分支，经过另一个层归一化和一个多层感知机。
   • 最终输出是这两个残差分支的和。
3. 头部 ：取[CLS]令牌对应的输出，通过一个线性分类头，得到最终的预测结果。

这种设计使得VideoMamba在形式上与ViT高度相似，只是把核心的Multi-Head Self-Attention层替换成了双向的SSM层。这种简洁性也大大降低了其复现和理解的难度。

3. 核心模块深度解析与复现要点

理解了宏观架构，我们深入到最关键的VideoMamba块和双向SSM的实现细节。这部分是复现过程中的重中之重，也是容易产生偏差的地方。

3.1 双向状态空间模型（Bi-SSM）的实现

单向的SSM（如Mamba）在语言建模中很常见，因为语言有强烈的从左到右的因果顺序。但视频理解更接近于“完形填空”，需要全局上下文。因此，实现双向扫描是关键。

在代码层面，一个双向SSM层可以这样实现（以PyTorch风格伪代码展示核心逻辑）：

import torch
import torch.nn as nn
from mamba_ssm import Mamba  # 假设使用官方或兼容的Mamba实现

class BiDirectionalSSM(nn.Module):
    def __init__(self, d_model, d_state=16, d_conv=4, expand=2):
        super().__init__()
        # 前向SSM
        self.ssm_forward = Mamba(
            d_model=d_model,
            d_state=d_state,
            d_conv=d_conv,
            expand=expand,
        )
        # 后向SSM。注意：在初始化时，它们应该是独立的参数。
        self.ssm_backward = Mamba(
            d_model=d_model,
            d_state=d_state,
            d_conv=d_conv,
            expand=expand,
        )

    def forward(self, x):
        """
        x: (B, L, D)  Batch, Sequence Length, Model Dimension
        """
        # 前向传播
        y_forward = self.ssm_forward(x)
        # 反向序列进行后向传播
        # 方法1：直接反转序列维度
        x_reversed = torch.flip(x, dims=[1])
        y_backward_rev = self.ssm_backward(x_reversed)
        y_backward = torch.flip(y_backward_rev, dims=[1])
        # 方法2：更高效的做法可能是在SSM内部实现双向逻辑，但这里为清晰起见使用反转。
        # 将两个方向的输出相加
        y = y_forward + y_backward
        return y

关键细节与避坑指南 ：

• 参数独立性 ：前向和后向SSM必须是两个完全独立的实例，拥有各自的参数。不能共享参数，因为它们在处理不同方向的信息。
• 序列反转的正确性 ：确保 torch.flip 操作在正确的维度（序列维度，通常是dim=1）上进行。输入输出都需要反转，以保证与原始序列对齐。
• 与Mamba原论文的差异 ：标准的Mamba是因果的（单向）。在VideoMamba中，双向扫描破坏了因果性，但这对于视频分类等非自回归任务是可接受的，并且带来了性能提升。
• 效率考量 ：上述实现进行了两次完整的前向传播。在实际高效的实现中，可能会将双向扫描融合到一次计算中，以优化速度。但作为研究和理解，分开实现更清晰。

3.2 VideoMamba块的完整结构

将Bi-SSM嵌入到残差块中，就构成了完整的VideoMamba块。其结构比Transformer的编码器块更简单，因为它没有显式的“注意力”和“FFN”之分，SSM层本身融合了序列混合的功能。

class VideoMambaBlock(nn.Module):
    def __init__(self, dim, d_state=16, d_conv=4, expand=2, drop_path=0.):
        super().__init__()
        # 第一个归一化层和SSM分支
        self.norm1 = nn.LayerNorm(dim)
        self.ssm = BiDirectionalSSM(dim, d_state, d_conv, expand)
        # 第二个归一化层和MLP分支
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, dim * expand),
            nn.GELU(), # 原文可能使用Swish/GELU
            nn.Dropout(0.1),
            nn.Linear(dim * expand, dim),
            nn.Dropout(0.1)
        )
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()

    def forward(self, x):
        # 残差连接1: SSM分支
        x = x + self.drop_path(self.ssm(self.norm1(x)))
        # 残差连接2: MLP分支
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x

实操心得 ：

1. 归一化位置 ：这里采用了“Pre-Norm”结构，即在SSM或MLP之前进行LayerNorm。这种结构通常比原始的“Post-Norm”更稳定，更容易训练深层模型。在复现时务必检查归一化层的位置，这与训练动态密切相关。
2. DropPath（随机深度） ：在深层网络中（如VideoMamba-Base有24层），使用DropPath（也称为Stochastic Depth）是防止过拟合、提升模型泛化能力的关键技术。它会在训练时随机“跳过”某些块，相当于一种集成学习。在推理时，所有块都参与。
3. MLP的扩展比 ： expand 参数控制了MLP中间层的维度，通常为4。这是一个重要的超参数，影响模型的容量和计算量。

3.3 输入处理：从视频到令牌序列

这部分与ViT高度一致，但对于视频，需要特别注意时空位置的编码。

class PatchEmbed(nn.Module):
    """ 将视频切分为块并嵌入 """
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, num_frames=8):
        super().__init__()
        self.img_size = (img_size, img_size)
        self.patch_size = (patch_size, patch_size)
        self.num_patches = (img_size // patch_size) ** 2
        self.num_frames = num_frames
        self.proj = nn.Conv3d(in_chans, embed_dim,
                              kernel_size=(1, patch_size, patch_size), # 时间维kernel为1
                              stride=(1, patch_size, patch_size))

    def forward(self, x):
        # x: (B, C, T, H, W)
        B, C, T, H, W = x.shape
        assert H == self.img_size[0] and W == self.img_size[1], f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
        # 投影: (B, C, T, H, W) -> (B, D, T, H//P, W//P)
        x = self.proj(x)
        # 重排维度: (B, D, T, H//P, W//P) -> (B, T * (H//P)*(W//P), D)
        D = x.shape[1]
        x = x.flatten(2).transpose(1, 2)  # (B, L, D)
        return x

# 位置编码：使用标准的可学习1D位置编码
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches * num_frames + 1, embed_dim)) # +1 for [CLS] token

注意事项 ：

• 卷积核 ：使用3D卷积 Conv3d 时，时间维度的kernel和stride通常设为1，因为我们不希望跨帧混合信息，这一步只是空间上的嵌入。时空交互交给后面的SSM层。
• 序列长度 ：最终的序列长度 L = T * (H/P) * (W/P) 。这是决定计算和显存消耗的关键。VideoMamba的线性复杂度优势在 L 很大时会极其明显。
• [CLS]令牌 ：与ViT一样，需要一个额外的可学习令牌，它通过自注意力（在这里是SSM的全局扫描）聚合整个序列的信息，用于最终分类。

4. 训练策略与调参经验实录

论文中给出了在大型数据集（如ImageNet-1K， Kinetics）上的训练配置，但当我们自己在不同规模的数据集上复现或微调时，这些策略需要调整。

4.1 优化器与学习率调度

VideoMamba通常使用AdamW优化器，这是训练现代Transformer类模型的标准选择。权重衰减（weight decay）对于防止过拟合至关重要。

• 学习率（LR） ：基础学习率是一个需要仔细调校的超参。对于从头训练，论文可能使用类似 1e-3 的量级。对于微调，通常更小，如 5e-4 到 1e-4 。
• 学习率调度 ：余弦退火（Cosine Annealing）调度器是主流选择。它能让学习率从初始值平滑地衰减到0，有助于模型收敛到一个更平坦的极小值，提升泛化能力。
• 预热（Warmup） ：在训练开始时，使用一个较短周期的线性预热（例如1个epoch的5%步数）非常重要。这可以防止模型在初期因梯度不稳定而“跑偏”。
• 层学习率衰减（Layer-wise LR Decay） ：对于深层模型，这是一个有效的技巧。给靠近输入的层设置更小的学习率，给靠近输出的层设置更大的学习率。这基于一个直观假设：底层学习的是通用特征，微调时不应变化太大；高层学习的是任务特定特征，需要更大更新。在复现时，可以尝试设置一个衰减系数（如0.75）。

我的调参记录（以在较小数据集上微调VideoMamba-Tiny为例） ：

• 优化器：AdamW (betas=(0.9, 0.999), weight_decay=0.05)
• 基础学习率： 2e-4
• 调度器：Cosine Annealing with Warmup
• 预热epoch：1个epoch（约总步数的5%）
• 总epoch数：30
• 批量大小：根据显存调整，尽可能大（如8或16）
• 梯度裁剪（Gradient Clipping）：Norm设置为1.0，这对稳定SSM的训练有时有帮助。

4.2 数据增强与正则化

视频领域的数据增强比图像更复杂，因为需要考虑时序一致性。

• 空间增强 ：随机水平翻转、多尺度随机裁剪（RandomResizedCrop）、颜色抖动（ColorJitter）、RandAugment或AutoAugment。这些与图像训练类似。
• 时序增强 ：
  • 随机帧采样 ：训练时从视频中随机抽取固定数量（如8或16）的帧。这是最常用且必要的，它增加了时间维度上的多样性，模拟了不同的观察速度。
  • 时序抖动 ：在随机起始点附近进行微小偏移采样。
  • 时序裁剪 ：从长视频中随机裁剪出一段连续片段。
• 模型正则化 ：
  • Dropout ：在MLP层和注意力（SSM）之后使用Dropout。VideoMamba的MLP中通常包含Dropout。
  • DropPath ：如前所述，对深层网络至关重要。
  • 标签平滑（Label Smoothing） ：在分类损失中使用标签平滑（如系数0.1），可以减轻模型对训练标签的过度自信，提升泛化能力。

注意 ：过于激进的数据增强（尤其是时空混合的增强）可能会破坏视频的语义连贯性，需要根据具体任务进行平衡。对于动作识别，空间增强为主；对于需要精细时序理解的任务，则要谨慎使用时序增强。

4.3 混合精度训练与显存优化

VideoMamba的线性复杂度优势在长序列下能节省大量显存，但模型参数量本身并不小。为了训练更大的批量或更长的序列，混合精度训练（AMP）是必备技能。

• AMP使用 ：PyTorch中通过 torch.cuda.amp.autocast 和 GradScaler 可以轻松实现。它能将部分计算转换为FP16，显著减少显存占用并加速计算。
• 梯度检查点（Gradient Checkpointing） ：对于极深的模型或极长的序列，即使线性复杂度，前向激活值也会占用大量显存。梯度检查点技术通过牺牲部分计算时间（重新计算中间激活）来换取显存节省。对于VideoMamba的深层架构，在资源紧张时可以考虑使用。
• 序列长度权衡 ：虽然SSM能处理长序列，但序列长度 L 依然直接影响激活值显存。在固定硬件下，需要在帧数（T）和空间分辨率（H, W, P）之间做权衡。例如，降低空间分辨率（增大patch size或缩小输入尺寸）可以容纳更多帧。

5. 复现与调试中的常见问题与解决方案

在将论文思路转化为可运行代码的过程中，我遇到了不少坑。这里总结一份“避坑指南”。

5.1 收敛困难或性能不达预期

• 问题现象 ：损失不下降，准确率远低于论文报告值。
• 排查思路 ：
  1. 数据流检查 ：首先确保数据加载和预处理流程100%正确。检查视频解码、帧采样、归一化（均值/标准差）是否与论文一致。一个常见的错误是归一化参数用错（ImageNet的mean/std vs. Kinetics的mean/std）。
  2. 模型初始化 ：检查模型参数初始化。SSM层的参数初始化有其特殊性（如A矩阵的初始化）。确保你使用的Mamba实现或自己实现的SSM层采用了正确的初始化策略。可以对比官方Mamba代码库的初始化方式。
  3. 学习率过高/过低 ：尝试一个广泛的学习率范围进行扫描（如从 1e-5 到 1e-2 ）。使用学习率查找器（LR Finder）工具可以帮助快速定位合适范围。
  4. 权重衰减过强 ：过大的 weight_decay 可能会抑制模型学习。尝试将其调小（如从0.05调到0.01或0.005）。
  5. 梯度爆炸/消失 ：监控梯度的范数。如果出现NaN，可能是梯度爆炸。可以尝试减小学习率、启用梯度裁剪（clip_grad_norm_）、或检查SSM中离散化步骤（如Delta、A、B矩阵的计算）的数值稳定性。

5.2 训练速度慢

• 问题现象 ：每个epoch耗时远超预期。
• 排查思路 ：
  1. 数据加载瓶颈 ：视频解码是I/O密集型操作。使用 prefetch_factor 和 num_workers 增加数据加载并行度。考虑将视频预提取为帧的图片序列或更高效的格式（如WebDataset格式的tar包）。
  2. SSM实现效率 ：确认使用的SSM内核是经过高度优化的。原版Mamba代码提供了CUDA优化的选择性扫描内核。如果使用纯PyTorch实现的SSM用于教学，其速度会慢很多。对于严肃的研究或应用，必须使用优化内核。
  3. 混合精度未生效 ：确认AMP已正确启用，并且主要计算（卷积、线性层、SSM）运行在FP16上。检查是否有某些操作被强制保持在FP32（如某些归一化层），这可能是必要的。
  4. 序列长度过长 ：虽然SSM是线性复杂度，但 O(T) 中的常数因子可能很大。如果序列长度 L 极大（例如处理非常高分辨率的视频），计算时间仍然可观。需要根据任务需求合理选择输入尺寸和帧数。

5.3 显存不足（OOM）

• 问题现象 ：即使批量很小，也出现CUDA out of memory错误。
• 排查思路 ：
  1. 激活值显存 ：这是主要占用。使用线性复杂度的SSM，激活值显存应与序列长度 L 成线性关系。检查你的 L = T * (H/P)*(W/P) 是否过大。尝试减少 T （帧数）、增大 P （patch size，如从16改为32）或减小 H, W （输入图像尺寸）。
  2. 梯度检查点 ：如前所述，对VideoMamba块启用梯度检查点。在PyTorch中，可以使用 torch.utils.checkpoint.checkpoint 。
  3. 优化器状态 ：AdamW优化器会为每个参数保存两个动量状态，这会使显存占用翻三倍。考虑使用像Adafactor或Shampoo这样更省显存的优化器，或者使用 混合精度训练 ，因为优化器状态可以部分存储在FP16中。
  4. 模型并行 ：对于极大的模型，可能需要将其拆分到多个GPU上。SSM的递归性质使其模型并行比Transformer的自注意力更复杂，但并非不可能。

5.4 与Transformer基线对比时的公平性

当你想验证VideoMamba是否真的比ViT或TimeSformer更好时，确保对比实验是公平的。

• 参数量与计算量（FLOPs）对齐 ：比较相同参数量级或相同FLOPs下的性能。VideoMamba-Tiny应与ViT-Tiny对比，而不是和ViT-Base比。
• 训练设置一致 ：使用完全相同的数据增强、学习率调度、优化器超参、训练epoch数。任何差异都可能导致结果偏差。
• 输入规格一致 ：帧数、分辨率、采样策略必须完全相同。
• 评估协议一致 ：使用相同次数的时序裁剪和空间裁剪进行测试（如常用的3次随机裁剪）。

我自己在复现时，最初因为给ViT基线使用了更强的数据增强，而VideoMamba用了默认配置，导致结果有误导性。统一所有设置后，才得到了可靠的对比结论。

6. VideoMamba的潜力、局限与未来展望

经过一系列实验和代码剖析，我对VideoMamba有了更立体的认识。

其优势是显而易见的 ：

1. 效率优势 ：在处理长视频序列时，其线性复杂度带来了实实在在的显存和速度优势。这在处理电影、监控视频、长教程等场景下潜力巨大。
2. 建模能力 ：选择性扫描机制使其能动态聚焦于相关信息，在实验中表现出对时序动态良好的捕捉能力，尤其在需要长程依赖的任务上。
3. 架构简洁 ：整体设计非常优雅，与ViT一脉相承，社区易于理解和接受，便于在其基础上进行各种改进和适配。

但同样，它也存在局限和挑战 ：

1. 训练不稳定性 ：SSM的训练有时比Transformer更敏感，对初始化、学习率等超参数的要求可能更高。需要更细致的调参。
2. 硬件优化程度 ：虽然有了CUDA内核，但其在各类硬件（尤其是非NVIDIA GPU）上的优化普及度远不如高度优化的Transformer注意力内核（如FlashAttention）。
3. 理论理解 ：相对于自注意力机制，SSM的内部工作机制可解释性稍弱，其“选择性”如何具体作用于视觉令牌，还需要更多的可视化分析工作。
4. 多模态扩展 ：当前工作主要集中在视频分类。如何将其有效地用于视频-语言多模态任务（如视频问答、描述生成），是一个开放且有趣的方向。这需要设计新的机制来融合视觉SSM和语言模型。

未来可能的演进方向 ：

• 更高阶的SSM ：探索更复杂的SSM变体，以捕捉视频中更复杂的时空动态。
• 与注意力的混合 ：也许纯粹的SSM并非终极答案。混合架构（如局部使用注意力处理关键帧，全局使用SSM处理长序列）可能结合两者优点。
• 无监督/自监督预训练 ：为VideoMamba设计像Masked Autoencoder (MAE) 或对比学习这样的预训练任务，以充分利用海量无标签视频数据，学习更通用的时空表示。

复现VideoMamba的过程，是一次深入理解状态空间模型如何与视觉问题结合的绝佳实践。它不仅仅是一个新模型的实现，更代表着一种用高效序列模型处理视觉信号的新范式。尽管前路仍有挑战，但它无疑为视频理解领域打开了一扇新的大门。对于研究者而言，这是一个值得深入探索的矿藏；对于工程师而言，当面临长视频处理的实际需求时，VideoMamba及其后续变体可能会成为一个有力的工具选项。我的建议是，不要仅仅停留在跑通代码，多去思考其设计背后的动机，尝试在不同的任务和数据集上测试其边界，你会有更多属于自己的发现。