Path Planning for Masked Diffusion Model Sampling

• DOI: 10.48550/arxiv.2502.03540
• Publication Date: 2025-02-05
• Authors: Peng Fei; Zachary Bezemek; Sawan Patel; S S Yao; Jarrid Rector-Brooks; Alexander Tong; Pranam Chatterjee
• Keywords: FOS: Computer and information sciences; Computer Science - Machine Learning; Artificial Intelligence (cs.AI); Computer Science - Artificial Intelligence; Machine Learning (cs.LG)

In this paper, we investigate how the order in which tokens are unmasked during masked diffusion models (MDMs) inference affects generative quality. We derive an expanded evidence lower bound (ELBO) that introduces a planner, responsible for selecting to unmask at each step. Our analysis suggests alternative unmasking strategies can improve performance. Based on these insights, propose Path Planning (P2), sampling framework leverages pre-trained BERT or denoiser itself guide decisions. P2 generalizes all known MDM and enables significant improvements across diverse domains including language generation (in-context learning, code generation, story infilling, mathematical reasoning, reverse curse correction) biological sequence (protein RNA sequences).