TABES: Trajectory-Aware Backward-on-Entropy Steering
for Masked Diffusion Models

TL;DR

Why unmasking order matters

A masked diffusion model (MDM) generates text by iteratively revealing masked positions. Once a token is written, it is locked in — the standard reverse process never revises it. So the schedule is a control decision: whatever you reveal becomes the conditioning context for everything that follows. Confidence-greedy schedules reveal the tokens the model is already sure about — which are precisely the tokens whose revelation carries the least information. The globally load-bearing pivots (an intermediate quantity in a derivation, a crucial entity binding) stay masked until it is too late, and one locally-plausible-but-wrong early commit cascades into global incoherence. We call this trajectory lock-in.

The method in four steps

At each denoising step t:

1. Prefilter. Take the top ρ·|M_t| masked positions by confidence as candidates C_t (ρ = 0.25 works well).
2. Surrogate. Run one forward pass at t−1 and measure the total predictive entropy H̃_t−1 over the still-masked positions, with the candidate embedding slots kept differentiable.
3. Backward. A single backward pass gives, for every candidate at once, the gradient g_i = ∇_eiH̃_t−1. ActiveQueryAttention stops gradients at all non-candidate query rows, so the adjoint scales with |C_t| instead of the sequence length.
4. Score & unmask. The Token Importance Score TIS_i = −⟨g_i, ẽ_i − e_mask⟩ is the first-order predicted reduction in future entropy from revealing i (with soft write ẽ_i = E^⊤π_i). Gate it by confidence, subtract a small anti-collapse penalty that discourages ultra-low-entropy reveals early on, and unmask the top scorers.

Cost per step: 2 forwards + 1 sparse backward, vs 1 forward for heuristics — and a first-order guarantee (paper, Thm. 3.1) that TIS matches the true one-step entropy reduction up to a curvature term.

Try it in two minutes

# CPU demo — trains a tiny MDM on 4x4 Sudoku, compares samplers (~5 min)
git clone https://github.com/shreshthsaini/TABES && cd TABES
pip install torch
python demo.py

# GPU demo — LLaDA-8B, confidence vs BoE on a GSM8K-style prompt
pip install transformers huggingface_hub
python demo_llada.py

Sudoku is the demo task because unmasking order is causally load-bearing there: a single wrong early commit provably forces constraint violations later — trajectory lock-in in its purest form. The demo also prints the per-step masked-entropy trace, where you can watch BoE drive future uncertainty down faster than confidence unmasking.

Results reported in the paper

Selected columns at generation length 128 from Table 2 of the paper (arXiv:2602.00250v2); see the paper for full tables, LLaDA-1.5 results, compute-matched Pareto curves, and ablations. This repository is the minimal demo companion; numbers depend on the evaluation protocol.

Limitations, honestly

BoE optimizes uncertainty, which correlates with — but is not — correctness: a model can be confidently wrong, and on tasks with brittle structure (e.g. code) the entropy objective can prefer reveal orders that hurt syntax even while it succeeds at reducing entropy. The backward pass also requires gradient support at inference time, which today's optimized serving stacks don't provide. See the paper's Limitations section, and the anti-collapse regularizer, for how we mitigate the first issue.

Citation

@article{saini2026tabes,
  title   = {TABES: Trajectory-Aware Backward-on-Entropy Steering
             for Masked Diffusion Models},
  author  = {Saini, Shreshth and Saha, Avinab and Adsumilli, Balu and
             Birkbeck, Neil and Wang, Yilin and Bovik, Alan C.},
  journal = {arXiv preprint arXiv:2602.00250},
  year    = {2026}
}