中文

超线性注意力:Softmax Transformer实现原位强化学习

机器学习 2026-05-19 v2

摘要

原位强化学习(ICRL)研究后训练后能通过条件于额外上下文而无需参数更新来适应新任务的智能体。现有的ICRL理论分析大多依赖于线性注意力,该方法用标准注意力中的恒等映射取代softmax函数。本文提供了首个在不作出不现实的线性注意力简化的情况下的ICRL理论理解。具体而言,我们考虑实际中使用的标准softmax注意力。我们表明,在特定参数下,带有此类softmax注意力的Transformer的层级前向传播等价于加权softmax时序差 TD(temporal difference)学习算法的迭代更新。这里,加权softmax TD是一种新的 RL算法,采用核空间中的策略评估,并将线性 TD 和表格 TD 作为特殊情况。我们还证明了在某种收缩条件下,策略评估误差随层数的增长而衰减,揭示了上述参数。最后,我们证明这些参数是预训练损失的全局最小化器,解释了它们在数值实验中的出现情况。

关键词

引用

@article{arxiv.2605.07333,
  title  = {Beyond Linear Attention: Softmax Transformers Implement In-Context Reinforcement Learning},
  author = {Zixuan Xie and Xinyu Liu and Claire Chen and Shuze Daniel Liu and Rohan Chandra and Shangtong Zhang},
  journal= {arXiv preprint arXiv:2605.07333},
  year   = {2026}
}