Softmax attention as policy improvement, and where it breaks

1Why softmax, not linear attention?

Transformers can run learning algorithms in context. For reinforcement learning, the constructions that improve a policy, and not only evaluate it, have used linear attention: its unnormalized bilinear products can carry a semi-gradient or actor–critic update. Standard softmax attention behaves differently. Its outputs are positive, normalized mixtures, and the constructive results for it have stopped at policy evaluation (temporal-difference learning).

For one well-known control update that normalization is not an obstacle but the update itself. Negative-entropy policy mirror descent (PMD) reweights each action by the exponential of its value and renormalizes:

State by state, that is exactly a softmax of the old log-probabilities plus the scaled values:

The identity is classical. What is new is treating it as the native actor of a causal Transformer and asking which control guarantees follow, and where they stop.

2Three calls, three errors

Each round runs in a fixed order: the actor must form a policy before the environment can sample under it, and the critic may only evaluate that policy after the samples exist. The order separates three kinds of error.

A greedy-defect recursion then carries ${\zeta }_k$ and ${\delta }_k$ through the rounds into a bound on $\Vert Q^{\star }-Q^{{\widehat{\pi }}_T}{\Vert }_{\infty }$, the loss of the policy the loop finally returns (§13). The last critic call does not enter, because the returned policy is formed before it.

3Why a stronger step can backfire

A larger step $\eta$ makes ideal PMD greedier, which helps. It also spreads the scores further apart, and a finite margin has to keep the other states’ memories down against that spread. In a two-state family where every quantity is exact, the competition reduces to one ratio, $\rho =\eta c/((1-\gamma )\kappa )$ for reward scale $c$: the implemented update improves the policy when $\rho \<1$ and makes it worse when $\rho >1$.

Within one state, the common level of the $Q$-values cancels out of PMD. Across states it does not: it shifts the two states’ softmax normalizers in opposite directions, by $\eta c/(1-\gamma )$ in total. So a larger step uses up the margin $\kappa$ faster than it sharpens the intended row. The critic is exact throughout: the wrong action is chosen from a perfect value table.

4The decoder in one picture

• Small and fixed. 2 blocks, 2 heads, head width S+A+3 and model width 6S+6A+14, with a gated ReLU feed-forward map of 8 active units. The same weights serve both calls; the step size ${\eta }_k$ arrives as token content.
• Normalization-free. No layer normalization: the value channels must keep controlled amplitudes while the key channels carry score-dependent routing, and normalizing tokens would change both at once.
• Declared before freezing. Step-size caps, packet sizes and coverage floors are fixed before the margins are chosen and the weights frozen; the bound holds for that declared horizon.

5Where it runs

The two NeurIPS 2026 workshop papers split the work: one builds the loop and its guarantee and checks a learned Transformer against it; the other maps where the loop fails. Two ICAIF 2026 workshop papers take the update rule into synthetic finance settings.

6How we test it

Every experiment uses small finite MDPs where the true values can be computed exactly, so each error is measured against the truth rather than estimated. Two kinds of evidence are kept apart.

• The exact construction. The fully specified decoder, run as written. On the two-state family every number is a deterministic computation; in the coverage study the decoder runs on 64 fresh random MDPs with $S=4$, $A=3$, $\gamma =0.8$, $\eta =0.8$ and $T=20$ rounds, its protocol fixed in a file before the run.
• Learned Transformers. Ordinary pre-LN Transformers trained on generated MDPs with disjoint train, validation and test splits, at the same $S$, $A$ and $\gamma$. The actor gets no PMD probabilities as labels: it is trained with a KL-proximal objective whose row-wise optimum is PMD, and the critic with Bellman (TD) targets. Learned models are compared with the construction through measured errors; they are not the construction.

Losses are reported as the normalized returned-policy loss $\Vert Q^{\star }-Q^{{\widehat{\pi }}_T}{\Vert }_{\infty }/\Vert Q^{\star }-Q^{{\pi }_0}{\Vert }_{\infty }$: 0 is optimal, 1 is no better than the starting policy.

7Results: the exact decoder

Take the two-state family at margin $\kappa =8$ and raise the step from 0.8 to 2.4. Ideal PMD gets better at every step. The implemented update, through the decoder’s own forward map, gets worse and then reverses.

Across the 1,803 cells of the grid ($\eta$ from 0 to 3 in steps of 0.005, three margins), the unmasked decoder has 840 cells with negative gain and the masked decoder 0. In this family the actor’s row error also predicts the value lost, exactly: $\Delta V_{\text{implemented}}=\Delta V_{\text{PMD}}-\frac{c}{2(1-\gamma )}\zeta$.

8Relaxing the coverage oracle

The guarantee assumes the critic receives fresh samples for every state–action pair after each actor call. An acting agent does not get that: it sees one trajectory. The theorem says nothing once the oracle is removed, so we ran its own decoder there. Every arm gets the same budget of 12$m$ fresh transitions per round; only how they are spread differs.

• One on-policy stream of the oracle’s size leaves 3.86 of 12 pairs unvisited per round on average, so the critic error rises from 0.214 to 1.219 at $m=6$, and the median loss from 0.009 to 0.029 (paired difference 0.018 [0.014, 0.022]). Four times the data brings it down only to 0.016.
• Carrying the old $Q$ forward for unvisited pairs brings the loss back to 0.010 although the critic error stays large, so the loss appears to come mainly from overwriting unvisited pairs with leaked outputs. That patch is not part of the decoder or the theorem.

9Results: learned Transformers

Does an ordinary Transformer, trained without PMD labels, learn the same map? On held-out MDPs a pre-LN actor’s output is closest to PMD by a wide margin among five fixed rules, and it responds to the step size it is given.

One-step correspondence does not show that an actor stays useful when its own outputs become its next inputs. A later study with separately trained actors (5 fresh runs of 51,200 training steps, seeds 12100–12104) put them in a closed loop on 64 held-out MDPs, with an exact one-step critic, so that it tests the actor alone. The reference is exact PMD at each actor’s fitted step, with the same critic.

The same checkpoints, without retraining, kept all 5 run ratios below 1.5 on four declared environment families of the same shape (cross-run medians 0.927, 1.004, 1.014, 0.997). Replacing the exact critic by the learned one, given 6 fresh samples per pair under the returned policy, leaves a small resolved gap to the reference: median 0.000518 [0.000008, 0.002155].

10Stress test and matched audit

Are the learned errors uniformly small? A screen of a separately trained pre-LN checkpoint, 512 examples per condition from 8 generated MDPs, says no: on low-entropy inputs the mean row error to PMD is 0.059 but the largest is 1.977, close to the maximum possible of 2 (all conditions in §15).

Two follow-ups separate what that screen can and cannot say.

• It is not the routing failure. This model’s actor packet for a state holds only that state’s memories, so the competing group of §3 does not exist and a state-local mask would change nothing.
• Large row error is not value loss. On a retrained checkpoint we ran the actor on every state and solved exact values for the policy it returns and for ideal PMD on the same inputs. At low entropy 79 of 512 inputs have negative implemented gain (by 0.126 on average, 1.426 at most), while the input with the largest row error gains 1.775 more than PMD.

To see whether cross-state visibility opens the channel of §3 in learned models, we trained paired actors on a packet holding all 12 memories, with and without a state-local mask (same seed, data and initial weights; 3 seeds, 3,200 steps). Replacing only the other states’ memories, which leaves the PMD target unchanged, moved the visible actor’s row by 0.255 on average on IID inputs (0.320 at low entropy) and the masked actor’s by exactly 0. The masked actor was also closer to PMD: row error 0.022 against 0.084 on IID inputs.

11Limitations and what didn’t work

Scope

• A generative oracle, not trajectories. The guarantee assumes fresh samples for every state–action pair after each actor call, with at least $m_{min,k}$ per pair. Trajectory data do not provide this: pairs go unvisited, counts depend on the data, samples are Markov-dependent and successor actions go stale. Step sizes and packet caps are supplied from outside; the decoder does not choose its own schedule.
• Normalization-free. The exact construction has no layer normalization, and nothing is proved for pre-LN, RMSNorm, post-LN or final-LN models. The learned pre-LN model is compared with the construction only through measured errors. The arXiv version also sketches a compilation into a layer-normalized model; its proof leaves constants unnamed and is not relied on here.
• Small generated MDPs. All learned results use $S=4$, $A=3$, $\gamma =0.8$ and generated task families with held-out splits: generalization within the family, not across environments.
• The reversal boundary is exact on one family only. It needs the uniform prior, self-loops, an exact critic and visible cross-state memories. At $\kappa =8$ a tilted prior moves the first negative step from about 1.6 to between 1.22 and 2.36, and a 5% cross-state transition moves it to 2.34. These are numerical sensitivities, not a law. The mask removes this mechanism only; it does not bound critic error, within-state error or learned-actor error, and it needs a state index that a general transcript may not expose.

What the evidence does not show

• No learned critic on trajectory data, and no learned closed-loop rollout with matched errors. Repeated control uses an exact one-step critic; with the learned critic it still uses the balanced oracle.
• The coverage study is one forgiving family. Its dense random transitions mix quickly: exact backup alone already gives loss 0.010, and the worst-case sampling bound at $\alpha =0.05$ (8.74 at $m=6$, 4.37 at $m=24$) is far above the measured critic errors. Hard-exploration MDPs are untested, and the reading that loss comes from overwriting unvisited pairs is an interpretation.
• Task loss does not identify the mechanism. A paired actor trained on rewards only had lower loss than the PMD-trained actor in repeated control (median difference −0.005963 [−0.009794, −0.003191]). The loss ratio of 1.052 is a registered usefulness band, and the paired difference to the reference, 0.000177 [−0.000137, 0.000748], establishes neither a disadvantage nor equivalence.
• “No PMD labels” is not discovery. The actor sees no PMD probabilities, but its objective encodes the PMD variational principle. The evidence supports close correspondence, not exact recovery of PMD’s coefficients, and the effective step shrinks towards the middle of the training range (slope 0.743).
• The cause of the low-entropy failures is unknown. Clipping the log-prior input to its training minimum changes little (13.7% of inputs still negative; chosen after looking at the worst row). Masking does not remove them either: 12.8% of the masked actor’s low-entropy inputs still lose value, against 22.7% visible. Learned results are diagnostics, not evidence of a general learned-model failure.
• No benchmark or state-of-the-art claim. The arXiv version compares against adaptations of two published in-context methods under a common harness in which PolicyAttention receives exact model information the others do not; we do not present that as a benchmark result. Later experiments on sparse interaction budgets and on external benchmarks were mixed or exploratory and support no performance claim.

Implementation and reproducibility

• Different checkpoints. The stress screen’s checkpoint was not kept. The matched audit retrained the same configuration and seed; two reruns agree with each other but not with the archive (low-entropy mean 0.0635, maximum 1.9885), so the screen and the audit come from different checkpoints. Repeated control also uses checkpoints other than those of §9’s one-step results.
• Carried, not rerun. The held-out correspondence numbers (0.035741, TD MAE 0.121) and the effective-step sweep are carried from the accepted submission and cannot be replayed from the current repository. A much smaller error reported in an early phase is a deterministic identity check, not a held-out result.
• Smaller than planned. The visible-versus-masked control was cut to 3 seeds at 3,200 steps for runtime. The critic in the stress screen was scored against targets sampled under ideal PMD rather than the returned policy; the matched audit corrects this for one critic call (0.154 against 0.155).
• Protocol records. The coverage protocol was written to a file before the run but is not externally timestamped, and its bound check on realized residuals is an indexing check, since it reuses those residuals.

12Open questions

These come from the papers’ own scope sections and reviewers’ remaining concerns; none has been run.

1. Learning without the oracle

   Run a learned critic on ordinary trajectory data, where pairs go unvisited and samples depend on each other, and measure what happens to the returned policy. The coverage study suggests where to look: how the critic treats pairs it has not seen.

2. Closed-loop audits of learned models

   Measure a learned model’s actor and critic errors against the policy it actually executes, round after round, rather than on fixed inputs, so that the bound’s residuals are observed in the loop they describe.

3. Routing failures in general models

   Find out whether reversals of the kind in §7 occur in learned Transformers that see other states’ memories. Visibility opens a sensitivity channel in the paired control, but that does not establish the reversal law in learned models.

13Definitions and the bound

Take a finite discounted MDP with $S$ states, $A$ actions, $0<\gamma \&lt;1$ and $|r|\le R_{max}$; write $B_Q=R_{max}/(1-\gamma )$ and ${\mathcal{F}}^{\pi }$ for the one-step Bellman operator under $\pi$. With $p_{k+1}={PMD}_{{\eta }_k}({\pi }_k,Q_k)$ the ideal update and ${\widehat{\pi }}_{k+1}$ the actor’s output, the two residuals are

For the greedy set $G_k(s)=\mathrm{\backslash arg}{max}_a{Q}_k(s,a)$, let $D_k(s)=-\log {\pi }_k(G_k(s)\mid s)$, ${\overline{D}}_k={max}_{s,a}{\mathbb{E}}_{S^{\prime }\mid s,a}{D}_k(S^{\prime })$ and $h_k={\overline{D}}_k/{\eta }_k+B_Q{\zeta }_k$. Choose $b_0\ge \Vert Q^{\star }-Q_0{\Vert }_{\infty }$ (for example $2B_Q$) and set $b_{k+1}=\gamma b_k+{\delta }_k+\gamma h_k$. Then, for a finite horizon $T$,

There is no ${\delta }_{T-1}$ term: the returned policy is formed before the backup that would produce $Q_T$. With state and role margins $\kappa ,\nu$, the actor’s routing error is bounded by

where

and, with leakage masses ${\lambda }_{1,k},{\lambda }_{2,k}$ and at least $m_{min,k}$ fresh samples per pair, with probability at least $1-\alpha$ over the horizon,

with the sampling term

A larger $\eta$ lowers the $\overline{D}/\eta$ term but raises the routing bound through $\eta B_Q$: the trade-off that §7 makes exact.

14The two-state family

States $s_{-}$ and $s_{+}$ loop on themselves, with action rewards $(-c,0)$ and $(c,0)$. From a uniform prior and its exact $Q^{\pi }$, the query at each state sees all four state–action memories with logits $\log \frac{1}{2}+\eta Q^{\pi }(\overline{s},a)+\kappa \mathbf{1}\{\overline{s}=s\}$ and one-hot action values, and itself with logit $-\nu$ and a uniform value. This row is exactly the actor call of the frozen decoder at $S=A=2$, up to a common score shift. At $s_{-}$, ideal PMD picks the better action with probability $\sigma (\eta c)$, and the wrong-state group, which prefers the other action, has partition ratio

so the sign of the implemented improvement $\Delta V(s_{-})=V^{\widehat{\pi }}(s_{-})-V^{\pi }(s_{-})$ is set by the margin against the step:

The query’s own uniform value lowers the gain but does not move the threshold. A state-local mask removes the wrong-state group ($Z_{\text{wrong}}=0$): without the query token the row is exactly PMD; with it, $\widehat{p}-\frac{1}{2}=Z_{\text{intended}}(p-\frac{1}{2})/(Z_{\text{intended}}+Z_q)\&gt;0$ for every $\eta >0$.

15Full tables

64 MDPs, $S=4$, $A=3$, $\gamma =0.8$, $\eta =0.8$, $T=20$, all margins 16. Unvisited pairs: mean per round, of 12. Critic error: mean per-round ${\delta }_k$ against the exact ${\mathcal{F}}^{{\widehat{\pi }}_{k+1}}$. Loss: median normalized returned-policy loss with a bootstrap interval over MDPs; the difference is the paired median against the balanced arm at the same budget.

Ratio of median normalized losses over 64 held-out MDPs, learned actor over the exact PMD reference at the actor’s fitted step, both with an exact one-step critic. Cross-run median 1.052; paired median difference 0.000177 [−0.000137, 0.000748] (95% hierarchical bootstrap).

512 examples per condition from 8 generated MDPs, $\eta =0.8$. Maxima are over the evaluated sample, not population suprema; critic errors are against six-sample TD targets drawn under ideal PMD.

512 inputs per condition. Gains are uniform-state means of $V-V^{\pi }$ under exact evaluation; ideal PMD gain is never negative. Critic error is against ${\mathcal{F}}^{\widehat{\pi }}{Q}^{\pi }$, with packets drawn under the returned policy; using a learned instead of an oracle critic in the next actor call changes value by −0.007.

16Reproducibility notes

Decoder

Every token coordinate, packet order, $Q/K/V/O$ matrix and feed-forward weight is listed in the EvoRobust supplement; at $S=A=2$ the weights have 60 non-zero entries. All 1,803 grid cells replay through the full forward map, and the repository circuit agrees on all weights and actor outputs.

Exact checks

Direct softmax and Bellman solves match the partition and error–value identities on the 1,803 grid cells and on 2,928 further checks over reward scale, discount and margin, all to machine precision (errors below 10⁻¹⁴).

Coverage study

64 MDPs drawn once from a fixed seed, uniform starting policy, $Q_0=0$; the decoder is the theorem’s own, unchanged.

Repeated control

5 runs of 51,200 steps (seeds 12100–12104), 64 held-out MDPs, 20 rounds, an exact one-step critic for both the learned actors and the reference.

Tests

The decoder’s 11 tests pass (shapes, score identity, masked exactness, envelope, critic bound, cross-call silence, feed-forward gate, shared weights, causal mask, witness replay, repository cross-check), as do the repository’s architecture tests.

17In brief

18Sources

PolicyAttention: Softmax Attention Implements Policy Mirror Descent for Closed-Loop Control. Yuhe Sui, Yingzhi Tang, Shufang Chen. arXiv 2609.30500, 2026. [arXiv]

The EvoRobust camera-ready version is titled PolicyAttention: When Stronger Softmax Policy Improvement Backfires. Results come from the camera-ready versions of the two NeurIPS workshop papers and their data, and from the records of the later repeated-control study.

Back to top ↑