Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Shreshth Saini, Avinab Saha, Balu Adsumilli, Neil Birkbeck, Yilin Wang, Alan C. Bovik

TL;DR

BoE tackles trajectory lock-in in masked diffusion models by introducing a training-free, gradient-guided one-step lookahead that predicts which tokens to unmask to minimize future entropy. It derives a Token Importance Score from a first-order embedding-space surrogate and employs ActiveQueryAttention to enable a sparse, backward pass, delivering improved accuracy-efficiency tradeoffs for non-autoregressive decoding. Across reasoning, coding, and distribution-level benchmarks, BoE yields a stronger accuracy-compute Pareto frontier and better global coherence compared to strong baselines. This control-theoretic framework provides a reusable approach to principled inference in discrete diffusion models and invites extensions to longer-horizon planning and multimodal diffusion systems.

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.

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

TL;DR

BoE tackles trajectory lock-in in masked diffusion models by introducing a training-free, gradient-guided one-step lookahead that predicts which tokens to unmask to minimize future entropy. It derives a Token Importance Score from a first-order embedding-space surrogate and employs ActiveQueryAttention to enable a sparse, backward pass, delivering improved accuracy-efficiency tradeoffs for non-autoregressive decoding. Across reasoning, coding, and distribution-level benchmarks, BoE yields a stronger accuracy-compute Pareto frontier and better global coherence compared to strong baselines. This control-theoretic framework provides a reusable approach to principled inference in discrete diffusion models and invites extensions to longer-horizon planning and multimodal diffusion systems.

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 ( 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.
Paper Structure (37 sections, 3 theorems, 33 equations, 4 figures, 5 tables, 1 algorithm)

This paper contains 37 sections, 3 theorems, 33 equations, 4 figures, 5 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Method: Backward-on-Entropy (BoE) Steering
  4. One-Step Lookahead as Entropy-Regularized Control
  5. Continuous Relaxation and a Differentiable One-Step Surrogate
  6. Token Importance Score (TIS) from a Single Backward Signal
  7. Theoretical Guarantee for the First-Order Surrogate
  8. Efficient Backward Scoring via ActiveQueryAttention
  9. Anti-Collapse Regularizer
  10. BoE Steering Algorithm
  11. Experiments
  12. Experimental Setup
  13. Main Results
  14. Ablations
  15. Conclusion
  16. ...and 22 more sections

Key Result

Theorem 3.1

Fix $t$ and $i\in\mathcal{M}_t$. Assume $\widetilde{\mathcal{H}}_{t-1}(\alpha;i)$ is twice continuously differentiable on $[0,1]$ and $\sup_{\alpha\in[0,1]}|\frac{d^2}{d\alpha^2}\widetilde{\mathcal{H}}_{t-1}(\alpha;i)|\le M$. Define $\Delta\widetilde{\mathcal{H}}_{t-1}(i):=\widetilde{\mathcal{H}}_{t

Figures (4)

  • Figure 1: Trajectory-aware sampling via Backward-on-Entropy steering. Unlike heuristic schedules that unmask tokens with low current entropy, BoE selects positions that maximize expected reduction in future masked entropy, approximated with a single backward signal. ActiveQueryAttention restricts the backward computation to active positions, preserving inference efficiency.
  • Figure 2: BoE (on right) prioritizes tokens that reduce future uncertainty unlike greedy schedule on left. On reasoning tasks, BoE tends to unmask globally load-bearing positions earlier, yielding faster reduction in total masked entropy and reduced incorrect trajectories. Best viewed zoomed-in.
  • Figure 3: Performance of greedy sampling with various unmasking method, and two trajectory-aware methods LookUM lookum and BoE(ours) from LLaDA-8B.
  • Figure 4: Pass@1 accuracy vs. full max_gen_len NFE on LLaDA-8B. We report compute-matched scaling curves on HumanEval, MBPP, GSM8K, and MATH500 for four training-free samplers: Top-$k$ (entropy), Top-$k$ (confidence), LookUM, and BoE (ours). Following ebsampler evaluation protocol, we fix max_gen_len and measure NFE as the number of denoiser evaluations required to unmask all tokens in the generation region.

Theorems & Definitions (7)

  • Theorem 3.1: First-order entropy-reduction surrogate
  • proof
  • Corollary 3.2: Ordering stability under margin
  • proof
  • proof
  • Proposition A.1: Boundedness of gated TIS
  • proof