Fast Parallel Sampling for Masked Diffusion Models via Dependence Partitions Prediction

Masked diffusion models enable parallel generation of discrete data, but simultaneously revealing multiple tokens can introduce error by neglecting their conditional dependencies. Recent methods address this challenge through adaptive unmasking schedules and dependency-aware token selection. However, explicit dependence checks can incur additional model evaluations, while inexpensive dependency estimates need not provide rigorous guarantees on sampling accuracy. In this work, we develop a theoretical framework for sampling with partition prediction. A partition oracle groups masked coordinates into bounded-size blocks, which our sampler processes in parallel while performing successive conditional updates within each block. Under quantitative assumptions on score estimation and partition quality, we establish nonasymptotic KL error bounds that separate initialization error, score estimation error, and residual cross-block dependence and quantify how the block size controls the trade-off between sampling accuracy and parallel query depth. We further construct partition oracles for ferromagnetic Ising models with bounded susceptibility, and masked language-model predictors with stable neural network.