# PolicyAttention: Softmax Attention Implements Policy Mirror Descent for Closed-Loop Control

arXiv: https://arxiv.org/abs/2609.30500 · HTML: https://mlsanity.com/paper/2609.30500/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 13 of its 13 references that are in the map.
- First posted on arXiv 24 Sep 2026.

## Builds on (most cited first, 13 of 13)

- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [A Theory of Regularized Markov Decision Processes](https://mlsanity.com/paper/1901.11275/) (2019)
- [In-context Reinforcement Learning with Algorithm Distillation](https://mlsanity.com/paper/2210.14215/) (2022)
- [Policy Mirror Descent for Reinforcement Learning: Linear Convergence, New Sampling Complexity, and Generalized Problem Classes](https://mlsanity.com/paper/2102.00135/) (2021)
- [On the Convergence Rates of Policy Gradient Methods](https://mlsanity.com/paper/2201.07443/) (2022)
- [Structured State Space Models for In-Context Reinforcement Learning](https://mlsanity.com/paper/2303.03982/) (2023)
- [Transformers Can Learn Temporal Difference Methods for In-Context Reinforcement Learning](https://mlsanity.com/paper/2405.13861/) (2024)
- [Optimal Convergence Rate for Exact Policy Mirror Descent in Discounted Markov Decision Processes](https://mlsanity.com/paper/2302.11381/) (2023)
- [Beyond Linear Attention: Softmax Transformers Implement In-Context Reinforcement Learning](https://mlsanity.com/paper/2605.07333/) (2026)
- [On the Convergence of Policy Mirror Descent with Temporal Difference Evaluation](https://mlsanity.com/paper/2509.18822/) (2025)
- [Transformers Provably Implement In-Context Reinforcement Learning with Policy Improvement](https://mlsanity.com/paper/2605.05755/) (2026)

## Abstract (arXiv)

Can causal softmax attention implement policy mirror descent as a repeated controller rather than a one-step algebraic identity? Negative-entropy policy mirror descent (PMD) has the statewise update operatornamePMDη(π,Q)=operatornamesoftmax(logπ+η Q). Building on the known Q-TD-PMD recursion, we construct one fixed causal-softmax actor--environment--one-step-critic protocol with explicit actor, routing, sampling, and normalization residuals, and propagate them to the policy actually returned. The construction states the finite-logit/full-support domain, the external tokenization and sampling boundary, and the mean-zero LayerNorm carrier conditions required by the normalized compilation. Separately trained pre-LN Transformers recover the target computation empirically. A frozen one-step audit model is closest to PMD among the tested fixed rules; in a preregistered five-run S=4 repeated-control test, the learned actor with an exact one-step critic reaches median returned-policy loss 1.052× the Exact PMD oracle and retains the criterion across four no-retraining shifts. The same checkpoints with their learned critic give descriptive median 1.050× the oracle (no registered margin). At S=8, replacing the exact critic by the learned critic raises median T=20 loss to 0.0225 yet leaves the Liang--Lai and Algorithm Distillation adaptations 20.2--24.2× higher-loss; this is a one-sided sampled-critic bound because PolicyAttention consumes 144 generative transitions per round versus 20 on-policy transitions for the adaptations. The strict 20-transition comparison remains open. At S=8,16, the exact-critic common-harness comparison remains 17.7--28.2× lower-loss than those adaptations, with the information asymmetry stated locally.