Saved in:
| Main Authors: | , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2026
|
| Subjects: | |
| Online Access: | https://arxiv.org/abs/2605.07333 |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866914574128644096 |
|---|---|
| author | Xie, Zixuan Liu, Xinyu Chen, Claire Liu, Shuze Daniel Chandra, Rohan Zhang, Shangtong |
| author_facet | Xie, Zixuan Liu, Xinyu Chen, Claire Liu, Shuze Daniel Chandra, Rohan Zhang, Shangtong |
| contents | In-context reinforcement learning (ICRL) studies agents that, after pretraining, adapt to new tasks by conditioning on additional context without parameter updates. Existing theoretical analyses of ICRL largely rely on linear attention, which replaces the softmax function in the standard attention with an identity mapping. This paper provides the first theoretical understanding of ICRL without making the unrealistic linear attention simplification. In particular, we consider the standard softmax attention used in practice. We show that, with certain parameters, the layerwise forward pass of a Transformer with such softmax attention is equivalent to iterative updates of a weighted softmax temporal difference (TD) learning algorithm. Here, weighted softmax TD is a new RL algorithm that performs policy evaluation in kernel space and adopts both linear TD and tabular TD as special cases. We also prove that under a certain contraction condition, the policy evaluation error decays as the number of layers grows, with the identified parameters above. Finally, we prove that those parameters are a global minimizer of a pretraining loss, explaining their emergence in our numerical experiments. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2605_07333 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | Beyond Linear Attention: Softmax Transformers Implement In-Context Reinforcement Learning Xie, Zixuan Liu, Xinyu Chen, Claire Liu, Shuze Daniel Chandra, Rohan Zhang, Shangtong Machine Learning In-context reinforcement learning (ICRL) studies agents that, after pretraining, adapt to new tasks by conditioning on additional context without parameter updates. Existing theoretical analyses of ICRL largely rely on linear attention, which replaces the softmax function in the standard attention with an identity mapping. This paper provides the first theoretical understanding of ICRL without making the unrealistic linear attention simplification. In particular, we consider the standard softmax attention used in practice. We show that, with certain parameters, the layerwise forward pass of a Transformer with such softmax attention is equivalent to iterative updates of a weighted softmax temporal difference (TD) learning algorithm. Here, weighted softmax TD is a new RL algorithm that performs policy evaluation in kernel space and adopts both linear TD and tabular TD as special cases. We also prove that under a certain contraction condition, the policy evaluation error decays as the number of layers grows, with the identified parameters above. Finally, we prove that those parameters are a global minimizer of a pretraining loss, explaining their emergence in our numerical experiments. |
| title | Beyond Linear Attention: Softmax Transformers Implement In-Context Reinforcement Learning |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2605.07333 |