Generation Order and Parallel Decoding in Masked Diffusion Models: An Information-Theoretic Perspective

TL;DR

The paper tackles the fundamental problem of how generation order and parallel decoding affect distributional correctness in Masked Diffusion Models. By building an information-theoretic framework, it distinguishes order sensitivity due to model error from parallelization bias due to factorized sampling, introducing forward KL, reverse KL, and incoherence as complementary diagnostics. It shows that Easy-First strategies become more beneficial as model error increases, while exact verification to remove sampling error incurs an exponential cost tied to conditional total correlation; remasking offers heuristic improvements without guaranteeing correctness. Empirically, controlled Block-HMM studies and arithmetic reasoning experiments validate the theory, highlighting a critical speed–fidelity trade-off and guiding principled decoding strategies for efficient yet reliable generation.

Abstract

Masked Diffusion Models (MDMs) significantly accelerate inference by trading off sequential determinism. However, the theoretical mechanisms governing generation order and the risks inherent in parallelization remain under-explored. In this work, we provide a unified information-theoretic framework to decouple and analyze two fundamental sources of failure: order sensitivity and parallelization bias. Our analysis yields three key insights: (1) The benefits of Easy-First decoding (prioritizing low-entropy tokens) are magnified as model error increases; (2) factorized parallel decoding introduces intrinsic sampling errors that can lead to arbitrary large Reverse KL divergence, capturing "incoherence" failures that standard Forward KL metrics overlook; and (3) while verification can eliminate sampling error, it incurs an exponential cost governed by the total correlation within a block. Conversely, heuristics like remasking, though computationally efficient, cannot guarantee distributional correctness. Experiments on a controlled Block-HMM and large-scale MDMs (LLaDA) for arithmetic reasoning validate our theoretical framework.

Figures (5)

• Figure 1: Order sensitivity and parallelization bias under imperfect factorization. Minimal examples illustrating that different generation orders ($\pi_1,\pi_2$) or parallel factorizations induce identical joint distributions under ideal conditions, but diverge under model error or conditional dependence.
• Figure 2: Forward vs. reverse KL under parallel factorization. (A1–A2) A distribution with a near-zero-support configuration. Although the factorized proposal $p_{\prod}(x_1,x_2)=p(x_1)p(x_2)$ preserves marginals, it assigns substantial probability mass to an implausible region (highlighted in red), resulting in small forward KL but large reverse KL and a non-negligible probability of incoherent samples. (B1–B2) A correlated but fully supported distribution. Here, factorization increases forward KL but does not introduce support violations, leading to small reverse KL and coherent parallel samples. All KL values are reported in nats; zero-probability entries are replaced by a small constant for numerical stability.
• Figure 3: Forward KL, reverse KL, and incoherence under parallel mean-field decoding in a Block-HMM with parity emissions ($B=8$). Mean-field decoding incurs severe sampling error and incoherence at low parity noise despite modest forward KL, while verified decoding achieves zero sampling error.
• Figure 4: Comparison of L2R and R2L decoding strategies under varying model scales.
• Figure 5: Teacher entropy $H_t$ vs. KL error $\epsilon_t$ (200 tokens). Higher entropy correlates with larger student error ($\alpha = 0.138$, $p < 0.001$, $R^2 = 0.26$).

Theorems & Definitions (10)

• Lemma 4.1: Rollout-weighted KL decomposition
• Proposition 5.1: Verification Requires Exponential Cost
• proof : Proof of \ref{['eq:weighted_objective']}
• Lemma 3.1: Reverse KL diverges under conditional support mismatch
• proof
• Corollary 3.2: Reverse conditional total correlation can be unbounded
• Lemma 3.3: A one-sided bound from reverse KL
• proof
• Proposition 3.5: Remasking Cannot Guarantee Elimination of Sampling Error
• proof