DriveMamba: Task-Centric Scalable State Space Model for Efficient End-to-End Autonomous Driving

本文提出了 DriveMamba，一种基于任务中心可扩展状态空间模型（Mamba）的单阶段端到端自动驾驶框架，它通过线性复杂度的序列建模和双向轨迹引导扫描机制，有效解决了现有 Transformer 方法在信息损失、累积误差及计算效率方面的局限，并在 nuScenes 和 Bench2Drive 数据集上展现了优越的性能与泛化能力。

要点

这篇论文介绍了一个名为 DriveMamba 的新系统，它是为了让自动驾驶汽车变得更聪明、更快速、更省电而设计的。

想象一下，以前的自动驾驶系统像是一个分工明确的流水线工厂，而 DriveMamba 则像是一个全能的超级大脑。

下面我用几个生活中的比喻来给你讲讲它是怎么工作的，以及它为什么厉害。

1. 以前的痛点：流水线工厂的“传话游戏”

以前的自动驾驶系统（比如 UniAD）通常是这样工作的：

• 第一步（感知）： 先让一个专门的“眼睛”看路，把看到的画成一张详细的地图（这叫 BEV 特征）。
• 第二步（预测）： 把地图交给“预测员”，让他猜别的车会怎么动。
• 第三步（规划）： 最后交给“司机”，让他决定怎么开。

问题出在哪？
这就像玩“传话游戏”。信息从“眼睛”传到“预测员”再传到“司机”，每传一次，信息就会丢失一点，或者被误解一点（累积误差）。而且，如果“眼睛”看错了，后面全错。更糟糕的是，这种流水线太慢了，因为每个环节都要把数据整理得整整齐齐（生成密集的地图），非常消耗电脑算力，就像为了送一杯水，非要先把水倒进一个巨大的游泳池里再舀出来一样。

2. DriveMamba 的解决方案：全能超级大脑

DriveMamba 打破了这种流水线，它把“看”、“猜”、“开”全部整合进了一个统一的超级大脑里。

核心魔法一：把“大海”变成“珍珠” (稀疏表示)

以前的系统喜欢把整个视野里的所有像素都处理一遍，就像在沙滩上把每一粒沙子都捡起来数一遍，太累了。
DriveMamba 很聪明，它只捡重要的珍珠（比如：前面的车、路边的线、自己车的轨迹）。它把图像直接变成一个个“任务令牌”（Token），只关注关键信息。

• 比喻： 就像你找东西，以前是把整个房间翻个底朝天；现在你只盯着手里那张“寻宝图”，直接去拿宝藏，速度快多了。

核心魔法二：Mamba 引擎 (线性复杂度)

这是论文最牛的地方。以前的系统用的是 Transformer（像 Transformer 模型），处理长序列信息时，计算量是平方级增长的（数据多一倍，计算量翻四倍）。这就像你每多认识一个人，就要和所有认识的人重新握手一次，人多了就累瘫了。
DriveMamba 用的是 Mamba（一种新的状态空间模型）。

• 比喻： Mamba 就像是一个高效的记忆管家。它不需要每次都重新和所有人握手，而是把重要的信息记在脑子里，随着新信息的到来，只更新最关键的部分。它的计算量是线性的（数据多一倍，计算量只多一倍）。
• 结果： 处理同样的路况，DriveMamba 比以前的系统快 3.2 倍，而且省了 68.8% 的内存。

核心魔法三：双向扫描与“轨迹引导” (怎么读信息)

Mamba 原本是用来处理文字（一维）的，怎么用来处理图像（二维）和自动驾驶（三维 + 时间）呢？
作者设计了一种混合扫描法：

1. 局部到全局 (Local-to-Global)： 就像开车时，你首先关注车头正前方的路（局部），然后慢慢把视线扩展到整个路口（全局）。
2. 轨迹引导 (Trajectory-Guided)： 系统会根据你打算怎么开（比如要变道还是直行），动态调整它“看”的顺序。

• 比喻： 以前的系统是死板地按顺序读报纸（从头读到尾）；DriveMamba 是根据你要找什么新闻，直接跳到那一版去读。如果你要变道，它就优先扫描旁边的车道，而不是死板地扫描整个天空。

3. 它有多强？(实验结果)

作者在两个著名的自动驾驶测试场（nuScenes 和 Bench2Drive）上做了测试：

• 更稳： 撞车率极低（在 Bench2Drive 上只有 0.13% 的碰撞率）。
• 更准： 规划路线非常平滑，像老司机一样。
• 更快： 即使是小版本的 DriveMamba，也能达到 17.9 FPS（每秒处理 17.9 帧），这意味着它反应极快，能在普通的车载芯片上流畅运行。
• 更省： 它的“小身材”（Tiny 版本）性能甚至超过了别人巨大的“大模型”（Large 版本）。

总结

DriveMamba 就像是给自动驾驶汽车换装了一个既聪明又省油的“新引擎”。
它不再笨拙地把所有数据都过一遍，而是学会了抓重点（稀疏表示），用高效的方法（Mamba）处理信息，并且懂得根据驾驶意图（轨迹引导）来调整注意力。

这让自动驾驶不仅开得更安全、更流畅，而且让原本需要昂贵超级计算机才能跑动的系统，现在也能在普通的汽车芯片上跑得飞快。这对于未来让自动驾驶汽车真正走进千家万户，是一个巨大的进步。

技术摘要

DriveMamba 技术总结

1. 研究背景与问题 (Problem)

端到端自动驾驶（E2E-AD）旨在将感知、预测和规划任务统一到一个可微分框架中。然而，现有的主流方法（如 UniAD, VAD 等）主要存在以下局限性：

• 串行范式与信息损失：大多数方法遵循“感知 - 预测 - 规划”的串行顺序，依赖人工定义的模块连接。这种设计不可避免地导致信息在模块传递中丢失，并产生累积误差，缺乏模块间灵活多样的关系建模。
• 计算效率与扩展性瓶颈：
  • 基于密集 BEV（鸟瞰图）特征的方法计算成本高昂，难以处理长距离感知。
  • 基于稀疏 Query 的方法虽然减少了计算量，但通常仍依赖 Transformer 的注意力机制，其**二次方复杂度（Quadratic Complexity）**限制了模型在处理长时序输入和高分辨率图像时的扩展性。
  • 现有的并行化方法（如 ParaDrive）虽然提高了灵活性，但往往忽略了模块间的动态关系建模，且缺乏闭环评估。
• 长时序建模困难：现有的时序融合方法通常存储历史 BEV 特征，这在长期建模中既不切实际也消耗大量显存。

2. 核心方法论 (Methodology)

作者提出了 DriveMamba，一种以任务为中心（Task-Centric）、可扩展的端到端自动驾驶范式。其核心是利用 Mamba（状态空间模型，SSM） 的线性复杂度特性，构建统一的解码器。

2.1 整体架构

DriveMamba 摒弃了生成密集 BEV 特征的传统做法，采用稀疏表示：

1. Token 化（Tokenization）：
   • 图像 Token：多视角传感器图像通过视觉编码器（如 ResNet, ViT, VMamba）提取特征，并转换为 Token 序列。
   • 任务 Token：预定义三种 Query 用于不同任务：Agent Query（动态物体检测与预测）、Map Query（静态地图构建）、Ego Query（自车规划）。
   • 位置编码：引入空间（Spatial）、时间（Temporal）和任务（Task）三种位置编码，区分不同来源的 Token。
2. 统一 Mamba 解码器（Unified Mamba Decoder）：
   • 所有 Token 被送入一个统一的解码器，并行处理，无需串行依赖。
   • 利用 Mamba 的线性复杂度特性，高效处理长序列 Token。

2.2 关键技术创新

• 混合时空扫描策略（Hybrid Spatiotemporal Scan）：
  • 为了在 Mamba 的一维序列建模中保留 2D/3D 空间局部性，设计了双向扫描机制。
  • 轨迹引导的“局部到全局”（Trajectory-Centric Local-to-Global）：根据自车未来轨迹的插值点，动态调整周围任务 Token 的扫描顺序。距离自车路径越近的物体优先级越高，符合驾驶注意力分布，有利于规划。
  • 水平/垂直交替扫描：在不同解码层交替使用水平优先和垂直优先扫描，以优化视图对应关系学习。
• 三大核心功能模块集成：
  1. 视图对应学习（View Correspondence Learning, VCL）：任务 Token 直接从原始传感器特征中提取语义，隐式学习 2D 图像与 3D 任务输出的对应关系，无需显式 BEV 转换。
  2. 任务关系建模（Task Relation Modeling, TRM）：通过双向 Mamba 层，动态捕捉不同任务（如感知与规划）之间以及同类任务内部的依赖关系（例如：规划导向的感知）。
  3. 长时序融合（Long-term Temporal Fusion, LTF）：使用 FIFO 队列存储历史任务 Token（而非历史 BEV 特征），结合运动感知层归一化（MLN）进行运动补偿，实现高效的长时序建模。

2.3 训练策略

• 单阶段端到端训练：所有任务（检测、建图、运动预测、规划）并行优化，无串行依赖。
• 迭代优化与残差学习：每个解码块输出相对偏移量，更新参考点和位置编码，逐步细化预测结果。

3. 主要贡献 (Key Contributions)

1. 提出 DriveMamba 范式：首个将动态任务关系建模、隐式视图对应学习和长时序融合集成到统一 Mamba 解码器中的 E2E-AD 框架，实现了高效且可扩展的并行处理。
2. 设计混合时空扫描方法：提出了“轨迹引导的局部到全局”和“空间到时间”的双向扫描策略，在保留时空局部性的同时，优化了自车视角的交互规划。
3. 验证了 Mamba 在 E2E-AD 中的优越性：证明了线性复杂度的 SSM 在替代二次方复杂度的 Transformer 时，能显著降低显存占用并提升推理速度，同时保持甚至提升性能。

4. 实验结果 (Results)

实验在 nuScenes（开环）和 Bench2Drive（闭环）数据集上进行。

• 开环规划性能 (nuScenes)：
  • DriveMamba-Tiny 在仅使用 40.2M 参数和 ResNet-50 骨干网络的情况下，平均 L2 误差为 0.44m，碰撞率为 0.15%，推理速度达到 17.9 FPS。
  • 相比 UniAD 和 VAD，DriveMamba-Tiny 在 L2 误差上分别降低了 42.1% 和 38.9%，且推理速度快 10 倍和 4 倍。
  • 随着模型规模扩大（Base, Large），性能进一步提升，DriveMamba-Large 达到 0.33m L2 误差和 0.06% 碰撞率。
• 闭环规划性能 (Bench2Drive)：
  • 在闭环测试中，DriveMamba-Base 的驾驶分数（Driving Score）达到 65.50，优于 DriveTransformer-Large（63.46），且参数量更少、延迟更低。
  • DriveMamba-Tiny 在仅 55.8ms 延迟下实现了 53.54 的驾驶分数，展现了极高的效率。
• 效率对比：
  • 相比 Transformer 基线，DriveMamba 在扩展输入分辨率时，显存占用减少 68.8%，处理速度提升 3.2 倍。
• 中间任务性能：
  • 在物体检测（mAP 50.8）、在线建图（mAP 64.8）和运动预测（minADE 0.55）任务上，DriveMamba 均超越了现有的 SOTA 方法（如 UniAD, SparseDrive, DriveTransformer）。

5. 意义与影响 (Significance)

• 突破扩展性瓶颈：DriveMamba 通过线性复杂度的 Mamba 架构，解决了 E2E-AD 模型在处理长时序、高分辨率数据时的显存和计算瓶颈，为构建更大规模的自动驾驶模型提供了可行路径。
• 重新定义任务交互：打破了传统的串行模块设计，通过统一解码器实现了任务间的动态交互和协同优化，特别是“规划导向的感知”能力，显著提升了系统的整体安全性。
• 实际部署潜力：DriveMamba-Tiny 在保持高性能的同时，推理速度极快（17.9 FPS），显存占用低，非常适合在计算资源受限的车载芯片上部署，推动了端到端自动驾驶从实验室走向量产应用。
• 范式转移：该工作标志着自动驾驶领域从基于 Transformer 的二次方复杂度注意力机制，向基于状态空间模型（SSM）的线性复杂度机制的重要转变。