[论文解读] Think Before You Act: Decision Transformers with Working Memory

Jikun Kang, Romain Laroche+4 more|arXiv (Cornell University)|May 24, 2023

Explainable Artificial Intelligence (XAI)被引用 5

一句话总结

本文提出决策变换器记忆模块（DT-Mem），一种增强记忆的强化学习架构，通过内容寻址的工作记忆模块显式存储和检索特定任务的经验。通过集成低秩微调（LoRA）实现高效微调，DT-Mem仅使用模型参数的10%即在Atari游戏和Meta-World任务上实现了最先进水平的泛化性能，显著提升了样本效率和适应能力。

ABSTRACT

Decision Transformer-based decision-making agents have shown the ability to generalize across multiple tasks. However, their performance relies on massive data and computation. We argue that this inefficiency stems from the forgetting phenomenon, in which a model memorizes its behaviors in parameters throughout training. As a result, training on a new task may deteriorate the model's performance on previous tasks. In contrast to LLMs' implicit memory mechanism, the human brain utilizes distributed memory storage, which helps manage and organize multiple skills efficiently, mitigating the forgetting phenomenon. Inspired by this, we propose a working memory module to store, blend, and retrieve information for different downstream tasks. Evaluation results show that the proposed method improves training efficiency and generalization in Atari games and Meta-World object manipulation tasks. Moreover, we demonstrate that memory fine-tuning further enhances the adaptability of the proposed architecture.

研究动机与目标

• 解决基于大语言模型的智能体依赖隐式参数化记忆而产生的效率低下问题，以及由此引发的灾难性遗忘现象。
• 通过显式建模工作记忆（受人类认知过程启发），提升跨任务训练效率与泛化能力。
• 通过在记忆模块上应用低秩微调（LoRA），实现仅用极少微调数据即可高效适应新任务。
• 在保持或提升性能的同时，相比现有方法降低计算与数据需求。
• 为决策变换器中的参数化记忆提供更具可解释性与模块化的替代方案。

提出的方法

• 该方法引入一个内部工作记忆模块，以内容寻址的矩阵形式实现，用于存储特定任务的经验。
• 通过基于注意力的机制执行记忆更新，计算输入序列与记忆键之间的相关性，并利用注意力权重更新值。
• 通过基于内容的寻址实现记忆检索，根据查询与存储键的相似度定位并检索相关信息。
• 使用低秩微调（LoRA）对记忆模块进行微调，通过一组少量可学习参数调节记忆输出，以适应新任务。
• 架构利用预训练的决策变换器获取通用知识，而微调过程中仅调整记忆模块，从而保持参数效率。
• 系统采用双机制：记忆更新用于存储新经验，记忆检索用于基于过往相关任务信息条件化决策。

实验结果

研究问题

• RQ1 显式内部工作记忆是否能提升多任务强化学习中的泛化能力与样本效率？
• RQ2 与隐式参数化记忆相比，内容寻址的记忆检索在灾难性遗忘与性能表现方面有何差异？
• RQ3 低秩微调（LoRA）在仅使用极少数据的情况下，能在多大程度上实现对记忆模块的高效微调？
• RQ4 所提出的记忆增强架构是否在Atari与Meta-World基准测试中均优于现有方法（如MDT、HDT与PDT）？
• RQ5 记忆大小与微调数据量对模型性能与适应能力有何影响？

主要发现

• DT-Mem在Meta-World ML45基准测试中达到最先进性能，微调后测试得分为0.95±0.10，优于HDT、PDT与MDT。
• 在Meta-World中，DT-Mem在无微调情况下的测试性能较HDT提升8%，微调后性能提升3%。
• 在Atari游戏中，DT-Mem仅使用10%的模型参数，即实现了与MDT-Top3相当或更优的泛化性能。
• 当在Pong数据集20%的样本上微调200k步后，DT-Mem Top3的性能较MDT-Top3提升1.154%，表明其具备有效适应能力。
• 性能下降幅度从10%数据下的0.978%降至20%数据下的1.154%提升，表明增加微调数据可有效改善泛化能力。
• 采用LoRA微调的记忆模块仅使用147K参数，显著少于HDT的147K微调参数与230万参数的超网络，显著降低了计算开销。