LLM2D
基于Mamba的决策模型中集成多模态输入令牌混合器:元Mamba决策
Integrating Multi-Modal Input Token Mixer Into Mamba-Based Decision Models: Decision MetaMamba
作者: Wall Kim
发布日期: 11/28/2024
arXiv ID: oai:arXiv.org:2408.10517v3

摘要

基于状态空间模型(SSM)的序列建模在各种任务中已展现出超越Transformer的性能,这使得人们对它在离线强化学习(RL)中超越决策Transformer及其改进版本的潜力抱有期待。然而,基于最先进SSM模型Mamba的决策模型未能取得优于这些改进型决策Transformer的性能。我们假设这种局限性源于选择性扫描阶段的信息丢失。为了解决这个问题,我们提出了决策MetaMamba(DMM),它在其输入层中为Mamba添加了一个令牌混合器。该混合器明确考虑了离线RL输入的多模态特性,包括状态、动作和剩余回报。DMM在显著减少参数数量的同时,提高了性能。值得注意的是,使用简单的线性令牌混合器也能获得类似的性能提升,这突显了保留来自相邻时间步的信息比令牌混合器的具体设计更为重要。这种对Mamba输入层的改进代表了与Transformer中使用的传统基于时间戳的编码方法的偏离。通过增强Mamba在离线RL中的性能(其特点是内存效率高且推理速度快),这项工作为其在未来RL研究中的更广泛应用开辟了新的途径。