改进Mamba模型的序列选择机制,提升离线强化学习性能
Decision MetaMamba: Enhancing Selective SSM in Offline RL with Heterogeneous Sequence Mixing
- 用全通道并行混合取代逐令牌选择,避免关键信息丢失
- 在多种任务上达到当前最优性能,参数量小适合实际应用
- 适合需要高效低延迟决策的离线强化学习场景
基于Mamba的模型在离线强化学习中受到广泛关注。然而,其选择性机制在关键步骤被遗漏时往往带来负面影响。为解决此问题,我们提出一种简单而有效的结构——决策元Mamba(DMM),将Mamba的令牌混合器替换为基于全连接层的序列混合器,并调整位置编码结构以保留局部信息。DMM在Mamba处理前对所有通道同时进行序列混合,防止因选择性扫描和残差门控导致的信息损失。大量实验表明,DMM在多种强化学习任务中均实现最先进性能,且参数量紧凑,展现出强大的实际应用潜力。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。