Path Planning for Masked Diffusion Model Sampling
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 which tokens to unmask at each step. Our analysis suggests that alternative unmasking strategies can improve generative performance. Based on these insights, we propose Path Planning (P2), a sampling framework that leverages pre-trained BERT or the denoiser itself to guide unmasking decisions. P2 generalizes all known MDM sampling strategies and enables significant improvements across diverse domains including language generation (in-context learning, code generation, story infilling, mathematical reasoning, reverse curse correction) and biological sequence generation (protein and RNA sequences).
Code (0)
등록된 구현이 없습니다.
Tasks
Code GenerationIn-Context LearningMathematical ReasoningmodelText GenerationMethods 이 논문이 사용한 방법론
Similar Papers 제목 키워드 기반
Planner Aware Path Learning in Diffusion Language Models Training
Diffusion language models have emerged as a powerful alternative to autoregressive models, enabling fast inference through more flexible and parallel generation paths. This flexibility of sampling is unlocked by new engi…
Code GenerationText GenerationThe Cosine Schedule is Fisher-Rao-Optimal for Masked Discrete Diffusion Models
In this work, we study the problem of choosing the discretisation schedule for sampling from masked discrete diffusion models in terms of the information geometry of the induced probability path. Specifically, we show th…
Lookahead Unmasking Elicits Accurate Decoding in Diffusion Language Models
Masked Diffusion Models (MDMs) as language models generate by iteratively unmasking tokens, yet their performance crucially depends on the inference time order of unmasking. Prevailing heuristics, such as confidence base…
Reinforcement LearningDiffusion-State Policy Optimization for Masked Diffusion Language Models
Masked diffusion language models generate text through iterative masked-token filling, but terminal-only rewards on final completions provide coarse credit assignment for the intermediate filling decisions that shape the…
Self-Rewarding Sequential Monte Carlo for Masked Diffusion Language Models
This work presents self-rewarding sequential Monte Carlo (SMC), an inference-time scaling algorithm enabling effective sampling of masked diffusion language models (MDLMs). Our algorithm stems from the observation that m…