# PolicyAttention

Softmax attention as a policy-improvement step for in-context control. A QRS@NTU research project (Quantitative Research Society, Singapore; Nanyang Technological University).

- Overview: https://qrsntu.org/projects/policyattention/
- How it works: https://qrsntu.org/projects/policyattention/details/
- arXiv: https://arxiv.org/abs/2609.30500
- PTA workshop paper, NeurIPS 2026 (OpenReview): https://openreview.net/forum?id=zbGWnk0kY3
- EvoRobust workshop paper, NeurIPS 2026 (OpenReview): https://openreview.net/forum?id=VgswuGrqp6
- ICAIF 2026 market-making pages: https://qrsntu.org/projects/policyattention/icaif/

Plain-text summary for AI agents and readers, generated from the evidence file behind the pages; the pages are the full account.

## Design

Negative-entropy policy mirror descent (PMD) updates a policy by pi_new(a|s) proportional to pi(a|s) exp(eta Q(s,a)), which is exactly softmax(log pi + eta Q). PolicyAttention uses one softmax attention row, with scores log pi + eta Q, as the actor of an in-context controller. A causal loop then runs actor, environment (fresh samples under the new policy) and a softmax critic (TD evaluation). A fixed 2-block, 2-head normalization-free decoder runs both calls, and its routing, critic and sampling errors propagate to a bound on the loss of the policy the loop returns, assuming a fresh per-state-action generative oracle. The robustness paper shows where finite attention breaks this: other states' memories keep a little mass, and a larger step can reverse improvement.

## Key results

- Exact two-state family (margin kappa = 8, gamma = 0.8): raising eta from 0.8 to 2.4 raises ideal PMD gain from 0.950 to 2.084 but turns the implemented gain from +0.916 to -2.009; the sign flips at eta*c = (1-gamma)*kappa. A state-local mask makes the decoder equal PMD in all 1,803 tested cells (840 negative without it, 0 with it).
- Coverage relaxed, same decoder, 64 small MDPs: a single on-policy stream of the oracle's size leaves 3.86 of 12 pairs unvisited per round and raises median normalized returned-policy loss from 0.009 to 0.029.
- Learned pre-LN Transformer trained without direct PMD labels: held-out row-L1 error to PMD 0.0357 (closest other fixed rule 0.112), TD MAE 0.121.
- Repeated control, 5 separately trained actors with an exact critic, 20 rounds: loss ratio to an exact PMD reference 0.918 to 1.396, median 1.052, all within the registered 1.5x band.
- Learned stress screen: on low-entropy inputs mean row error 0.059 but sample maximum 1.977; a matched audit shows large row error need not mean value loss.

## Limitations

- The guarantee needs a fresh per-state-action generative oracle after every actor call; ordinary trajectories do not provide it, and the coverage experiment covers one small, forgiving MDP family with the exact decoder.
- The exact construction is normalization-free; the learned pre-LN model is compared with it only through measured errors.
- No learned critic on trajectory data and no learned closed-loop rollout with matched errors.
- The reversal threshold is exact only on the matched two-state family; a tilted prior moves the first negative step from about 1.6 to between 1.22 and 2.36.
- The archived stress checkpoint was not kept; the matched audit uses a retrained checkpoint, so the two tables come from different checkpoints.
- A paired reward-only actor had lower task loss than the PMD-trained actor (-0.005963 [-0.009794, -0.003191]), so task loss alone does not identify the PMD mechanism.

## Not claimed

- No guarantee for trajectory data: the theorem assumes a fresh per-pair generative oracle after every actor call.
- No guarantee for layer-normalized Transformers: the exact construction is normalization-free, and the learned pre-LN model is compared with it only through measured errors.
- No closed-loop rollout of a learned model with matched errors, and no learned critic run on trajectory data.
- No state-of-the-art or benchmark claim: later comparisons with external methods are mixed or exploratory.
- The reversal threshold eta*c = (1-gamma)*kappa is exact only on the matched two-state family; outside it the threshold moves.

## Papers

- PolicyAttention: Softmax Attention Implements Policy Mirror Descent for Closed-Loop Control. Yuhe Sui, Yingzhi Tang, Shufang Chen. arXiv 2609.30500, 2026.
- PolicyAttention: When Softmax Attention Becomes Policy Improvement. PTA Workshop, NeurIPS 2026. Accepted.
- PolicyAttention: Robustness and Failure Modes of Softmax Policy Improvement. EvoRobust Workshop, NeurIPS 2026. Accepted.
- PolicyAttention: A Certifiable and Testable Update Rule for Attention-Based Allocation Controllers. IAFM'26 workshop, ICAIF 2026. Accepted.

## A question to start with

Explain the idea to a newcomer, then give the three strongest criticisms of the evidence and one experiment that would settle each.