Optimizing Decoding Paths in Masked Diffusion Models by Quantifying Uncertainty

TL;DR

This work establishes Denoising Entropy as a principled tool for understanding and controlling generation, effectively turning the uncertainty in MDMs from a liability into a key advantage for discovering high-quality solutions.

Abstract

Masked Diffusion Models (MDMs) offer flexible, non-autoregressive generation, but this freedom introduces a challenge: final output quality is highly sensitive to the decoding order. We are the first to formalize this issue, attributing the variability in output quality to the cumulative predictive uncertainty along a generative path. To quantify this uncertainty, we introduce Denoising Entropy, a computable metric that serves as an internal signal for evaluating generative process. Leveraging this metric, we propose two algorithms designed to optimize the decoding path: a post-hoc selection method and a real-time guidance strategy. Experiments demonstrate that our entropy-guided methods significantly improve generation quality, consistently boosting accuracy on challenging reasoning, planning, and code benchmarks. Our work establishes Denoising Entropy as a principled tool for understanding and controlling generation, effectively turning the uncertainty in MDMs from a liability into a key advantage for discovering high-quality solutions.

Figures (7)

• Figure 1: Quantifying Path Uncertainty in MDMs with Denoising Entropy.Left: State Entropy ($h_{\texttt{DE}}$) measures per-state uncertainty, computed as the mean Shannon Entropy over the predictive distributions for all masked positions. $h_{\texttt{DE}}$ is then aggregated over the entire path to form the Path Entropy ($H_{\texttt{DE}}$). Right: Decoding process shows how different paths yield outputs of varying quality. Our key insight is that the lower $H_{\texttt{DE}}$ indicates path yeilding better output, providing a potent internal signal for generation quality.
• Figure 2: An overview of our entropy-guided decoding algorithms. While standard inference in subfigure (1) generates a single decoding path, our methods explore multiple candidate paths guided by Denoising Entropy. E-BoN in subfigure (2) performs post-hoc selection, choosing the best path from multiple independent candidates based on path entropy $H_{\texttt{DE}}$. E-SMC in subfigure (3) provides real-time guidance, iteratively pruning high-entropy paths and replicating low-entropy ones based on state entropy $h_{\texttt{DE}}$ of partial solutions.
• Figure 3: Trend of $H_{\texttt{DE}}$ vs. $\ln(\text{PPL})$ with Denoising Steps.
• Figure 4: Per-sample correlation between $H_{\texttt{DE}}$ and $\ln(\text{PPL})$.
• Figure 6: Average accuracy on code benchmarks. E-BoN and E-SMC consistently enhance performance of various baseline decoding strategies, demonstrating their broad applicability.

Theorems & Definitions (23)

• Definition 1: Oracle State Uncertainty
• Definition 2: State Entropy
• Definition 3: Path Entropy
• Proposition 1: $h_{\texttt{DE}}$ as an upper bound
• Definition 4: $\epsilon$-Accurate Denoising Model
• Proposition 2: Approximation of Normalized Loss by $h_{\texttt{DE}}$
• Proposition 3: Path Entropy Gap as a Quality Bound
• Lemma 1: Subadditivity of Conditional Entropy
• proof
• Lemma 2: Decomposition of Cross-Entropy Loss