arXiv:2409.06985cs.LG2024-09ICML被引 2

发现预训练语言模型中的马尔可夫注意力头,提升离线强化学习长期任务表现。

Unveiling Markov Heads in Pretrained Language Models for Offline Reinforcement Learning

  • 分析注意力头发现马尔可夫头导致仅关注最后一个输入
  • 提出GPT2-DTMA方法,在微调中引入多注意力机制
  • 在长时序任务中显著缩小性能差距,适合长期决策场景

近期将预训练语言模型(PLMs)知识融入决策变压器(DTs)在离线强化学习中备受关注。这些模型在任务中表现优异,引发一个关键问题:传递了何种知识?本文首次定量分析各注意力头,发现PLMs中存在马尔可夫头——该头对最后一个输入过度关注,仅在短期环境中有效。进一步证明,重训练嵌入层或微调无法改变此现象。为此,提出通用方法GPT2-DTMA,为预训练DT引入注意力混合(MoA),以适应微调中多样化的注意力需求。大量实验验证理论并表明:在短时环境表现相当,同时显著缩小长时环境下的性能差距。

原文摘要 · Abstract (English)

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.

强化学习注意力机制语言模型长期决策

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。