MambaFlow：用于高效光流估计的跨帧状态空间建模方法

《Neurocomputing》：MambaFlow: Cross-frame state space modeling for efficient optical flow estimation

　　摘要
近期，Mamba架构在多种计算机视觉任务中取得了显著成果。然而，其在光流估计领域的应用尚未被探索。本文提出了MambaFlow，这是一种专为端到端光流估计设计的基于Mamba的框架。该框架包含两个核心组件：（1）PolyMamba，通过双Mamba架构增强特征表示，结合用于帧内建模的Self-Mamba模块与用于帧间交互的Cross-Mamba模块；（2）PulseMamba，它利用注意力引导聚合器（AGA）以自适应方式整合特征，并采用动态学习的权重而非简单的串联方式，再借助Mamba的循环机制迭代进行自回归流解码。大量实验表明，MambaFlow在基准数据集上的表现可与主流方法相媲美。与SEA-RAFT相比，MambaFlow在Sintel基准测试中具有更高的精度，同时保持相当的推理效率，展现出在实际光流估计中良好的精度-效率平衡。

引言
光流估计是计算机视觉中的基础问题，旨在计算视频序列中连续帧之间像素的运动向量。它在动作识别[59]、[68]、视频插值[16]、[28]、[54]、自动驾驶[62]等领域有着重要且广泛的应用。基于卷积神经网络（CNN）的架构在光流领域取得了开创性进展。作为首个端到端的基于CNN的光流网络，FlowNet[5]展现了这种学习范式的巨大潜力。然而，它在处理大位移和复杂运动时存在局限性。随后，PWC-Net[45]通过引入金字塔特征提取和成本体积构建方法，提高了复杂运动场景下的估计精度，有效捕捉了多尺度局部特征。但它在大位移场景下仍面临挑战。为解决这些问题，RAFT[47]引入了循环迭代优化机制，显著提升了光流估计的精度。不过，这一创新也带来了新的问题，即其优化过程需要大量迭代，导致参数数量增加和推理时间延长。此后出现了许多RAFT算法的改进版本[3]、[18]、[19]、[46]、[51]、[69]，如MS-RAFT和SEA-RAFT，但它们处理大位移和复杂运动的能力仍有待提升。为应对上述问题，具有强大全局建模能力的Transformer架构也被引入到光流估计任务中[15]、[30]、[43]、[44]。这些方法大多利用注意力机制增强特征表示，从而捕捉长距离依赖关系，显著提高精度。值得注意的是，与基于RNN的流解码算法[20]、[47]不同，GMFlow[57]利用注意力机制迭代解码最终光流，从而实现更快速的流优化。然而，它们在训练和推理阶段都面临着Transformer的二次计算复杂度问题[15]、[30]、[43]、[44]、[57]，这严重限制了其在实际应用中的效率。

近期，Mamba[8]、[11]已被广泛应用于各种计算机视觉任务，如图像分类[7]、[60]、[61]、图像分割[6]、[27]、[71]、物体检测[32]、[33]、[48]、图像生成[13]、[23]、视频理解[21]、[24]、[55]等。这些显著成果表明，Mamba能够有效解决CNN和Transformer存在的难题，包括CNN的局部感知局限以及Transformer的二次计算复杂度问题[58]、[66]。它得益于选择性状态空间机制，能够高效捕捉全局信息，并以线性复杂度处理复杂的时空依赖关系[66]。具体而言，由于SSM中内置的RNN机制，Mamba被认为非常适合处理长序列和自回归任务[63]。受Mamba在长序列建模和自回归更新方面的能力启发，我们提出了一种新的光流估计框架——MambaFlow。在我们的设计中，Mamba主要用于增强特征和迭代优化流场，而全局匹配阶段则仍依赖成对相关性来提供初始对应关系估计。如图1所示，我们的方法在保持竞争力的精度的同时降低了推理延迟。该框架包含两个核心组件：（1）PolyMamba模块，它利用Self-Mamba、Cross-Mamba和多层感知器（MLP）以轻量级计算实现全局特征优化；（2）PulseMamba模块，它首先使用注意力引导聚合器（AGA）通过动态学习的权重自适应整合运动特征，而非简单串联，然后利用Mamba的自回归能力迭代解码和优化光流，有效解决了大位移和遮挡区域等难题。

具体实现上，我们首先采用CNN骨干网络从输入图像帧中提取稳定特征，这与传统方法类似[20]、[51]、[57]。随后，我们利用PolyMamba增强提取的特征，从而捕捉复杂的全球信息和长距离依赖关系。在增强特征的基础上，我们进一步使用全局匹配机制[57]、[67]计算相似度以获得初始流场。最后，我们借助PulseMamba通过循环优化得到最终精确的流场[57]。我们的算法在编码和解码阶段均以Mamba为核心组件，是一种以Mamba为中心的端到端光流估计算法。通过利用Mamba在全局建模和自回归方面的能力以及其高效的计算特性，该方法实现了精度与性能之间的良好平衡。

我们的主要贡献可总结如下：
• 我们提出了MambaFlow，这是一种基于Mamba框架的新型光流估计架构，既用于特征增强也用于流场传播。据我们所知，这是首个专为端到端光流估计设计的以Mamba为核心的架构。我们的方法在降低推理延迟的同时达到了与最先进算法相当的精度，因此有望用于对效率要求较高的光流应用。
• 我们引入了PolyMamba模块，该模块由Self-Mamba、Cross-Mamba和MLP按顺序组成，用于全局增强参考图像及其与匹配图像的交叉特征的特征表示。正如我们将要展示的，它有效地利用了Mamba的长序列建模能力，同时保持了线性的计算负载。
• 我们设计了PulseMamba，它结合了用于自适应特征融合的注意力引导聚合器以及基于Mamba的自回归解码器，解决了类似RNN的循环优化方法中存在的计算负担过重、大位移和遮挡问题。通过利用Mamba的自回归建模能力和高效的状态空间更新，它提升了迭代流优化中的精度-效率平衡。

章节节选
光流估计
传统上，光流估计被定义为能量最小化问题[1]、[12]、[39]，早期的变分方法依赖于亮度恒定性和空间平滑性等假设。随着深度学习的兴起，早期的基于CNN的方法，如FlowNet[5]、PWC-Net[45]和SpyNet[41]，采用了从粗到细的流程来进行密集运动估计。尽管这些方法效果不错，但往往会在金字塔层级间出现误差累积，尤其是在……

预备知识
SSM源于线性系统，通过隐藏状态将输入序列映射到输出。该模型可以表示为线性常微分方程（ODE）：
其中， 和 表示从给定参数 、 、 及时间尺度参数 导出的离散化参数；此外， 还是用于将隐藏状态 映射到输出的投影矩阵。最后，模型通过全局卷积操作生成输出，其定义为：
其中， 表示……

数据集
遵循之前的方法，我们首先在FlyingChairs（Chairs）[5]和FlyingThings3D（Things）[35]数据集上进行预训练。在微调阶段，我们使用了由KITTI[36]、HD1K[22]、FlyingThings3D和Sintel[2]组成的混合训练集。最后，我们在KITTI训练集上进一步微调模型，并在相应的在线基准测试上报告性能。

评估指标
在实验中，我们使用了光流估计领域常用的评估指标。首先，我们……

结论
近期，Mamba在计算机视觉中的长距离高效建模方面展现出巨大潜力。基于这一认识，本文提出了MambaFlow，这是一种以Mamba为核心的端到端光流估计框架。通过结合基于Mamba的特征增强与自回归流优化，该方法实现了精度与效率之间的良好平衡。实验结果表明，MambaFlow在Sintel基准测试上的表现尤为出色，同时……

CRediT作者贡献声明
周志虎：撰写——初稿、项目管理、形式分析、数据整理、概念设计。王竹成浩：软件实现、研究、形式分析。杜俊天：撰写——审阅与编辑、研究、形式分析、数据整理、概念设计。陈品一：资金获取、形式分析、数据整理。叶伟平：可视化、资源协调、研究。徐瑞基：监督、资源协调、数据整理。赵永标：撰写——审阅与编辑、软件实现。

利益冲突声明
作者声明不存在任何可能影响本文工作的已知财务利益或个人关系。

致谢
本研究得到了浙江省“先锋”和“领雁”研发计划的支持，项目编号分别为2026C02A1019、2026C02A1222、2026C02A1018；浙江省自然科学基金的支持，项目编号为LTGG23F020002；浙江工业大学校企合作项目的支持，项目编号为KYY-HX-20250798、KYY-HX-20230493、KYY-HX-20250381、KYY-HX-20250423；以及浙江省医疗卫生科学计划的支持，项目编号为WKJ-ZJ-26092。周志虎目前是浙江工业大学人工智能专业的博士候选人，同时也是该校的特聘教授。他还是一名高级工程师，被誉为“杭州数字工匠”，并担任杭州云健智能科技有限公司（YunjianAI）的总经理。他拥有西南大学计算机与信息科学学院的学士学位，以及浙江大学计算机视觉方向的硕士学位。
周志虎|王竹成浩|杜俊天|陈品一|叶伟平|徐瑞基|赵永标|毛克基
中国杭州，浙江工业大学信息工程学院