arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2409.06985cs.LG
收藏

揭示预训练语言模型中用于离线强化学习的 Markov 头

Unveiling Markov Heads in Pretrained Language Models for Offline Reinforcement Learning

Pennsylvania State University(宾夕法尼亚州立大学);Queen Mary University of London(伦敦玛丽女王大学);Microsoft Research Asia(微软亚洲研究院)

Wenhao Zhao, Qiushui Xu, Linjie Xu, Lei Song, Jinyu Wang, Chunlai Zhou, Jiang Bian

AI总结:

本文揭示预训练语言模型中的 Markov 头导致 DT 只关注最后 token,并提出带 MoA 的 GPT2-DTMA,在保持短期任务性能的同时显著改善长期任务。

摘要:

近来,将来自预训练语言模型(PLM)的知识融入决策 Transformer(DT),已在离线强化学习(RL)中引起广泛关注。这些 PLM 在 RL 任务上表现良好,从而引出一个有趣的问题:PLM 中的哪类知识被迁移到 RL 中,才取得了如此好的结果?本工作首先通过定量分析每个注意力头来深入研究这一问题,并指出 Markov 头这一存在于 PLM 注意力头中的关键组件。它会导致对最后一个输入 token 的极端关注,并且仅在短期环境中表现良好。此外,我们证明,这种极端关注无法通过重新训练嵌入层或微调来改变。受分析启发,我们提出一种通用方法 GPT2-DTMA,它为预训练 DT 配备 Mixture of Attention(MoA,注意力混合),以适应微调过程中多样化的注意力需求。大量实验证实了我们的定理,并证明了 GPT2-DTMA 的有效性:它在短期环境中取得可比性能,同时显著缩小了长期环境中的性能差距。

英文摘要:

Recently, incorporating knowledge from pretrained language models (PLMs) into decision transformers (DTs) has generated significant attention in offline reinforcement learning (RL). These PLMs perform well in RL tasks, raising an intriguing question: what kind of knowledge from PLMs has been transferred to RL to achieve such good results? This work first dives into this problem by analyzing each head quantitatively and points out Markov head, a crucial component that exists in the attention heads of PLMs. It leads to extreme attention on the last-input token and performs well only in short-term environments. Furthermore, we prove that this extreme attention cannot be changed by re-training embedding layer or fine-tuning. Inspired by our analysis, we propose a general method GPT2-DTMA, which equips a pretrained DT with Mixture of Attention (MoA), to accommodate diverse attention requirements during fine-tuning. Extensive experiments corroborate our theorems and demonstrate the effectiveness of GPT2-DTMA: it achieves comparable performance in short-term environments while significantly narrowing the performance gap in long-term environments.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑