Transformer可无需参数更新实现强化学习策略优化,理论与实证均验证其有效性。
Transformers Provably Implement In-Context Reinforcement Learning with Policy Improvement

- 通过显式参数构造,Transformer块可实现半梯度SARSA与演员-评论家算法。
- 在随机生成的表格MDP上训练后,模型在未见任务中表现优异,达到最优参数流形。
- 首次给出ICRL中梯度流收敛性证明,适合关注机制解释与训练动态的研究者。
我们研究了Transformer在上下文强化学习(ICRL)中的能力,即模型需从轨迹数据中推断并执行学习算法,而无需参数更新。研究表明,线性自注意力Transformer块可通过显式参数构造,严格实现包括半梯度SARSA和演员-评论家在内的策略改进方法。超越存在性证明,我们设计了教师模仿训练流程,分析其梯度流动力学,并建立了ICRL领域的首个收敛性保证:在训练MDP分布满足适当丰富性条件下,梯度流局部且指数收敛至对应期望强化学习更新的最优参数流形。实验上,在随机生成的表格型MDP上训练Transformer,验证了上述预测:学习到的模型恢复了显式构造的参数结构,并在未见MDP上展现出强劲的上下文控制性能。这些结果揭示了Transformer架构如何在上下文中内化并执行经典强化学习算法,弥合了机制理解与训练动态之间的鸿沟。
原文摘要 · Abstract (English)
We investigate the ability of transformers to perform in-context reinforcement learning (ICRL), where a model must infer and execute learning algorithms from trajectory data without parameter updates. We show that a linear self-attention transformer block can provably implement policy-improvement methods, including semi-gradient SARSA and actor-critic, via explicit parameter constructions. Beyond existence, we design a teacher-mimicking training procedure, analyze its gradient-flow dynamics, and establish the first convergence guarantee in the ICRL literature: under suitable richness conditions on the training MDP distribution, gradient flow converges locally and exponentially to an optimal parameter manifold corresponding to the desired RL update. Empirically, training transformers on randomly generated tabular MDPs confirms these predictions: the learned models recover the parameter structure of our explicit constructions and, when deployed on unseen MDPs, deliver strong in-context control performance. Together, these results illuminate how transformer architectures internalize and execute classical reinforcement learning algorithms in context, bridging mechanistic understanding and training dynamics in ICRL.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。