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

TL;DR

This paper introduces BoE Steering, a gradient-based inference method for Masked Diffusion Models that approximates long-term decision effects efficiently, improving generation coherence and scalability.

Contribution

It proposes a novel gradient-guided inference framework with the Token Influence Score and ActiveQueryAttention for scalable, trajectory-aware generation in MDMs.

Findings

BoE outperforms existing methods on inference scalability.

Gradient guidance reduces incoherence in generated outputs.

The approach offers a mathematically principled solution for non-autoregressive generation.

Abstract

Masked Diffusion Models (MDMs) have emerged as a promising non-autoregressive paradigm for generative tasks, offering parallel decoding and bidirectional context utilization. However, current sampling methods rely on simple confidence-based heuristics that ignore the long-term impact of local decisions, leading to trajectory lock-in where early hallucinations cascade into global incoherence. While search-based methods mitigate this, they incur prohibitive computational costs ($O(K)$ forward passes per step). In this work, we propose Backward-on-Entropy (BoE) Steering, a gradient-guided inference framework that approximates infinite-horizon lookahead via a single backward pass. We formally derive the Token Influence Score (TIS) from a first-order expansion of the trajectory cost functional, proving that the gradient of future entropy with respect to input embeddings serves as an optimal control signal for minimizing uncertainty. To ensure scalability, we introduce \texttt{ActiveQueryAttention}, a sparse adjoint primitive that exploits the structure of the masking objective to reduce backward pass complexity. BoE achieves a superior Pareto frontier for inference-time scaling compared to existing unmasking methods, demonstrating that gradient-guided steering offers a mathematically principled and efficient path to robust non-autoregressive generation. We will release the code.