Decision MetaMamba: Enhancing Selective SSM in Offline RL with Heterogeneous Sequence Mixing

📄 arXiv: 2602.19805v1 📥 PDF

作者: Wall Kim, Chaeyoung Song, Hanul Kim

分类: cs.LG, cs.AI

发布日期: 2026-02-23


💡 一句话要点

Decision MetaMamba:异构序列混合增强离线强化学习中的选择性SSM

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

关键词: 离线强化学习 Mamba模型 选择性扫描 序列混合 决策模型 强化学习

📋 核心要点

  1. Mamba模型在离线强化学习中表现出色,但其选择机制在关键步骤缺失时会损害性能。
  2. Decision MetaMamba (DMM)通过密集层序列混合器替代Mamba的token混合器,并调整位置结构以保留局部信息,从而避免信息损失。
  3. 实验表明,DMM在多种强化学习任务中实现了最先进的性能,且参数量较小,具有实际应用潜力。

📝 摘要(中文)

基于Mamba的模型在离线强化学习中备受关注。然而,当强化学习序列中的关键步骤被省略时,其选择机制通常会产生不利影响。为了解决这些问题,我们提出了一种简单而有效的结构,称为Decision MetaMamba (DMM),它用基于密集层的序列混合器取代了Mamba的token混合器,并修改了位置结构以保留局部信息。通过在Mamba之前执行同时考虑所有通道的序列混合,DMM可以防止由于选择性扫描和残差门控导致的信息丢失。广泛的实验表明,我们的DMM在各种强化学习任务中提供了最先进的性能。此外,DMM以紧凑的参数规模实现了这些结果,展示了在实际应用中的强大潜力。

🔬 方法详解

问题定义:离线强化学习中,基于Mamba的模型虽然表现出色,但其选择性扫描机制在处理不完整的序列数据时会丢失关键信息,导致性能下降。现有方法难以在选择性和信息完整性之间取得平衡。

核心思路:DMM的核心思路是在Mamba模型之前,通过一个密集层实现的序列混合器,对所有通道的信息进行融合。这样可以确保在进行选择性扫描之前,所有通道的信息都得到了充分的考虑,从而避免关键信息的丢失。同时,修改位置结构以保留局部信息,进一步提升模型性能。

技术框架:DMM模型主要由两部分组成:首先是一个基于密集层的序列混合器,用于融合所有通道的信息;然后是改进的Mamba模块,用于进行选择性扫描。输入序列首先经过序列混合器,然后输入到Mamba模块中进行处理。模型的整体训练目标是最大化离线数据集上的策略性能。

关键创新:DMM的关键创新在于使用密集层序列混合器替代了Mamba原有的token混合器。这种替代使得模型能够在进行选择性扫描之前,充分考虑所有通道的信息,从而避免了信息丢失。此外,对位置结构的修改也有助于保留局部信息,提升模型性能。

关键设计:序列混合器使用一个或多个密集层实现,其输入是整个序列的所有通道的信息,输出是融合后的序列表示。Mamba模块的位置结构被修改为更关注局部信息。损失函数通常采用行为克隆或策略约束等离线强化学习中常用的损失函数。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,Decision MetaMamba (DMM) 在多个离线强化学习任务中取得了最先进的性能。与现有方法相比,DMM 在性能上取得了显著提升,同时保持了较小的参数规模。例如,在某些任务上,DMM 的性能提升超过了10%,并且参数量远小于其他基于Transformer的模型。

🎯 应用场景

Decision MetaMamba (DMM) 在离线强化学习领域具有广泛的应用前景,例如在机器人控制、游戏AI、推荐系统等领域,可以利用离线数据训练出高性能的策略。该方法尤其适用于数据不完整或存在缺失的场景,能够有效提升模型的鲁棒性和泛化能力。未来,DMM有望在更多实际应用中发挥重要作用。

📄 摘要(原文)

Mamba-based models have drawn much attention in offline RL. However, their selective mechanism often detrimental when key steps in RL sequences are omitted. To address these issues, we propose a simple yet effective structure, called Decision MetaMamba (DMM), which replaces Mamba's token mixer with a dense layer-based sequence mixer and modifies positional structure to preserve local information. By performing sequence mixing that considers all channels simultaneously before Mamba, DMM prevents information loss due to selective scanning and residual gating. Extensive experiments demonstrate that our DMM delivers the state-of-the-art performance across diverse RL tasks. Furthermore, DMM achieves these results with a compact parameter footprint, demonstrating strong potential for real-world applications.