Understanding Parallel Samplers in Masked Diffusion via Random Walks on Graphs

Masked diffusion models (MDMs) allow for any-order generation, meaning tokens can be revealed in parallel. However, parallel decoding introduces statistical risks because replacing true block conditionals with a product of independent one-coordinate marginals is only valid if the chosen tokens are conditionally independent given the current context. This paper investigates when parallel unmasking is accurate and how to design samplers that respect the conditional dependence structure of the target distribution.

The authors propose using random walks on graphs as a controlled, verifiable sandbox. By training an MDM on random walk samples, the graph structure acts as a latent, controllable constraint that the model never sees explicitly. This framework allows for a Sudoku-like validity check: because the graph transition kernel is known, researchers can verify if a generated sequence is a valid walk and measure the fidelity of the learned distribution. The authors use this to compare different parallel unmasking policies, such as lowest-entropy selection versus random selection, and to develop a new bisection-style sampler.

The study shows that no single parallel unmasking heuristic is uniformly optimal. For example, entropy-guided unmasking may perform well on tree-like structures by identifying key anchors, but it can fail on bottleneck graphs where it prematurely commits to dependent choices. To address this, the authors introduce a bisection sampler that reveals middle blocks of tokens to separate the past and future, effectively breaking the sequence into conditionally independent subproblems. This approach is provably exact for order-k Markov walks and requires only logarithmic parallel depth. Experiments on both graph walks and pretrained language models demonstrate that this coarse-to-fine schedule improves the speed-quality tradeoff compared to standard greedy or random parallel samplers.