[2502.03540] Path Planning for Masked Diffusion Model Sampling

Computer Science > Machine Learning arXiv:2502.03540 (cs) [Submitted on 5 Feb 2025 (v1), last revised 5 Mar 2026 (this version, v5)] Title:Path Planning for Masked Diffusion Model Sampling Authors:Fred Zhangzhi Peng, Zachary Bezemek, Sawan Patel, Jarrid Rector-Brooks, Sherwood Yao, Avishek Joey Bose, Alexander Tong, Pranam Chatterjee View a PDF of the paper titled Path Planning for Masked Diffusion Model Sampling, by Fred Zhangzhi Peng and 7 other authors View PDF HTML (experimental) Abstract:Any order generation of discrete data using masked diffusion models (MDMs) offers a compelling alternative to traditional autoregressive models, especially in domains that lack a natural causal ordering of data. However, current popular MDMs depart from their successful continuous diffusion model counterparts with simplified masked inference wherein unmasked tokens cannot be iteratively refined -- even if there is a mistake. In this paper, we extract the full power of MDMs by introducing a novel inference sampling strategy termed Path Planning (P2) that decomposes each generation step into two sub-stages: planning and denoising. Under P2, the planner at every step selects appropriate tokens that are marked to be updated, which can then be sampled using the denoiser. We demonstrate that P2 generalizes all existing sampling strategies for MDMs and critically enhances generative quality through the new capability of refining and updating existing unmasked tokens. We theoretically prove t...