mlsanityData to 14 Oct 2026

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

arXiv 2609.30500 ↗First posted 24 Sep 2026cs.LG cs.AIcited by 0

Not on the map. Why?

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.

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

13 of 13 references.

References · 13

  1. 2022Transformers learn in-context by gradient descentcited by 732
  2. 2022What Can Transformers Learn In-Context? A Case Study of Simple Function Classescited by 715
  3. 2022What learning algorithm is in-context learning? Investigations with linear modelscited by 684
  4. 2019A Theory of Regularized Markov Decision Processescited by 286
  5. 2022In-context Reinforcement Learning with Algorithm Distillationcited by 190
  6. 2021Policy Mirror Descent for Reinforcement Learning: Linear Convergence, New Sampling Complexity, and Generalized Problem Classescited by 156
  7. 2022On the Convergence Rates of Policy Gradient Methodscited by 119
  8. 2023Structured State Space Models for In-Context Reinforcement Learningcited by 111
  9. 2024Transformers Can Learn Temporal Difference Methods for In-Context Reinforcement Learningcited by 23
  10. 2023Optimal Convergence Rate for Exact Policy Mirror Descent in Discounted Markov Decision Processescited by 18
  11. 2026Beyond Linear Attention: Softmax Transformers Implement In-Context Reinforcement Learningcited by 5
  12. 2025On the Convergence of Policy Mirror Descent with Temporal Difference Evaluationcited by 3
1 more
  1. 2026Transformers Provably Implement In-Context Reinforcement Learning with Policy Improvementcited by 3

Built on by

Not cited yet.