Trajectory Mamba: Efficient Attention-Mamba Forecasting Model Based on Selective SSM

📄 arXiv: 2503.10898v1 📥 PDF

作者: Yizhou Huang, Yihua Cheng, Kezhi Wang

分类: cs.CV

发布日期: 2025-03-13

备注: Accepted by CVPR 2025

💡 一句话要点

提出Trajectory Mamba，基于选择性SSM的高效轨迹预测模型

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 轨迹预测 状态空间模型 选择性SSM 自动驾驶 运动预测

📋 核心要点

  1. 传统基于注意力的轨迹预测模型计算复杂度高，难以应用于动态环境。
  2. Trajectory Mamba利用选择性SSM重构注意力机制，实现线性复杂度并提升预测精度。
  3. 实验表明，该模型在推理速度和参数效率上均优于现有方法，并在Argoverse数据集上取得SOTA结果。

📝 摘要（中文）

本文提出了一种名为Trajectory Mamba的新型高效轨迹预测框架，该框架基于选择性状态空间模型（SSM）。传统的基于注意力机制的模型面临着计算成本随目标数量呈二次方增长的挑战，这阻碍了它们在高度动态环境中的应用。为了解决这个问题，我们利用SSM重新设计了编码器-解码器架构中的自注意力机制，从而实现了线性时间复杂度。为了解决修改注意力机制可能导致的预测精度下降问题，我们提出了一种联合折线编码策略，以更好地捕捉静态和动态上下文之间的关联，最终提高预测精度。此外，为了平衡预测精度和推理速度，我们采用了与编码器完全不同的解码器。通过跨状态空间注意力，所有目标代理共享场景上下文，允许SSM在解码期间与共享场景表示进行交互，从而推断出未来预测步骤中的不同轨迹。我们的模型在Argoverse 1和Argoverse 2数据集上实现了最先进的推理速度和参数效率。与现有方法相比，它的FLOPs减少了四倍，参数数量减少了40%以上，同时超过了绝大多数先前方法的性能。这些发现验证了Trajectory Mamba在轨迹预测任务中的有效性。

🔬 方法详解

问题定义：现有的基于注意力机制的轨迹预测模型，其计算复杂度随着预测目标的数量呈平方级增长，这使得它们难以应用于需要快速响应的高度动态环境，例如自动驾驶。因此，如何降低计算复杂度，同时保持甚至提升预测精度，是本文要解决的核心问题。

核心思路：本文的核心思路是利用选择性状态空间模型（Selective State Space Model, SSM）来替代传统注意力机制。SSM具有线性时间复杂度，可以显著降低计算负担。同时，为了弥补替换注意力机制可能带来的精度损失，引入了联合折线编码策略，以增强模型对静态和动态上下文信息的理解。

技术框架：Trajectory Mamba采用编码器-解码器架构。编码器部分使用SSM替代自注意力机制，对输入轨迹进行编码。解码器部分则采用与编码器不同的结构，通过跨状态空间注意力，使所有目标代理共享场景上下文，并利用SSM与共享场景表示进行交互，从而预测未来轨迹。

关键创新：最重要的技术创新点在于使用SSM来替代传统的自注意力机制，从而将计算复杂度从平方级降低到线性级。此外，联合折线编码策略也是一个重要的创新，它能够更好地捕捉静态和动态上下文之间的关联，从而提高预测精度。

关键设计：解码器部分的关键设计在于使用跨状态空间注意力，使得所有目标代理能够共享场景上下文信息。这种设计允许模型在解码过程中考虑到不同目标之间的相互作用，从而更准确地预测它们的未来轨迹。此外，具体的参数设置、损失函数和网络结构等细节在论文中应该有更详细的描述（未知）。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

Trajectory Mamba在Argoverse 1和Argoverse 2数据集上取得了SOTA结果，并在推理速度和参数效率方面显著优于现有方法。具体而言，该模型相比现有方法，FLOPs减少了四倍，参数数量减少了40%以上，同时超过了绝大多数先前方法的性能。

🎯 应用场景

Trajectory Mamba在自动驾驶领域具有广阔的应用前景，可以用于车辆、行人等交通参与者的轨迹预测，从而提高自动驾驶系统的安全性。此外，该模型还可以应用于机器人导航、智能交通管理等领域，实现更高效、更智能的运动规划和决策。

📄 摘要（原文）

Motion prediction is crucial for autonomous driving, as it enables accurate forecasting of future vehicle trajectories based on historical inputs. This paper introduces Trajectory Mamba, a novel efficient trajectory prediction framework based on the selective state-space model (SSM). Conventional attention-based models face the challenge of computational costs that grow quadratically with the number of targets, hindering their application in highly dynamic environments. In response, we leverage the SSM to redesign the self-attention mechanism in the encoder-decoder architecture, thereby achieving linear time complexity. To address the potential reduction in prediction accuracy resulting from modifications to the attention mechanism, we propose a joint polyline encoding strategy to better capture the associations between static and dynamic contexts, ultimately enhancing prediction accuracy. Additionally, to balance prediction accuracy and inference speed, we adopted the decoder that differs entirely from the encoder. Through cross-state space attention, all target agents share the scene context, allowing the SSM to interact with the shared scene representation during decoding, thus inferring different trajectories over the next prediction steps. Our model achieves state-of-the-art results in terms of inference speed and parameter efficiency on both the Argoverse 1 and Argoverse 2 datasets. It demonstrates a four-fold reduction in FLOPs compared to existing methods and reduces parameter count by over 40% while surpassing the performance of the vast majority of previous methods. These findings validate the effectiveness of Trajectory Mamba in trajectory prediction tasks.