Understanding Parallel Samplers in Masked Diffusion via Random Walks on Graphs
Abstract
In this paper, we propose using random walks on graphs as a verifiable sandbox to study different parallel sampling strategies in masked diffusion models (MDMs). We train an MDM on random walk samples from a fixed graph. The graph or the transition kernel is never shown to the model explicitly and plays the role of latent structure in the sequences, albeit one that is controllable and can be used for quantitative evaluation. Thus, this framework enjoys a validity check to verify that an output is a valid walk and the estimated transition kernel from the walks yields a measure of distribution fidelity. Using simple graphs, we theoretically prove that parallel unmasking via widely used scores like lowest entropy is not uniformly better than a random parallel sampler; the performance critically depends on the structure of the underlying graph, a separation difficult to isolate in benchmarks like Sudoku. We also develop new bisection samplers for MDMs, which take logarithmic steps in the sequence length and are provably exact for random walks under perfect training. Experiments on various graph-walk tasks confirm that different parallel samplers are better for different graphs even in practice. Our initial experiments on a pretrained OpenWebText MDLM show that the bisection-style samplers improve speed-quality tradeoffs even for language generation. Together, these results position graph random walks as a mechanistic benchmark for diagnosing and designing parallel samplers for masked diffusion models.