PaperlayerView exact revision on arXiv

Source-linked AI summary

Path Planning for Masked Diffusion Model Sampling

Fred Zhangzhi Peng, Zachary Bezemek, Sawan Patel, Jarrid Rector-Brooks, Sherwood Yao, Avishek Joey Bose, Alexander Tong, Pranam Chatterjee

arXiv:2502.03540v5cs.LGcs.AI

TL;DR

Masked diffusion models provide order-flexible generation but cannot normally refine unmasked tokens, limiting inference when denoisers make mistakes. The paper introduces Path Planning (P2), which separates planning from denoising and uses planners to select tokens for unmasking or remasking. P2 generalizes existing MDM samplers and delivers state-of-the-art improvements across language and biological sequence tasks.

  • Problem

    Masked diffusion inference typically cannot iteratively refine unmasked tokens, even when they were incorrectly denoised.

  • Method

    Path Planning (P2) expands the ELBO and uses planner-guided inference to select tokens for updating, with configurations that generalize existing MDM sampling strategies.

  • Results

    P2 delivers state-of-the-art improvements across diverse language and biological sequence tasks and enables MDMs to outperform larger autoregressive models in selected settings.

  • Takeaways & Limitations

    Inference strategy and unmasking order are important determinants of discrete diffusion generation quality across text, code, and biological sequence design.

  • Takeaways & Limitations

    Biological experiments were purely computational, with no wet-lab synthesis or functional validation.

Abstract

from arXiv · show

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 that P2 establishes a (new) expanded evidence lower bound (ELBO) on the log marginal likelihood of data. We instantiate P2 with a family of planners including: 1.) Self-Planning, 2.) BERT-Planning, and 3.) Trained-Planning with a learned planner leading to SOTA generative performance for MDMs on a suite of domains. Specifically, solely using P2 inference, we observe relative improvements of 22% in protein sequence foldability, 8% in RNA sequence pLDDT, 4% in math reasoning, 68% in story generation (ROUGE score), and 33% in code generation for the challenging pass@1 metric.

1 Introduction

Masked diffusion models offer order-flexible generation, but fixed unmasking and irreversible denoising can propagate errors. Path Planning (P2) expands inference with planner-guided token updates and improves generation across diverse domains.

  • MDMs are attractive for discrete domains without a natural causal ordering, although autoregressive models remain the most successful discrete generative models.
  • Uniform random unmasking can be suboptimal with imperfect denoisers, while current inference cannot correct incorrectly denoised tokens later.
  • P2 expands the MDM ELBO with planner terms for selecting tokens to unmask and optionally resample during inference.
  • P2 planners can use the denoiser, pretrained BERT-type models, or a lightweight planner trained offline, and P2 generalizes existing MDM sampling strategies.
  • P2 achieves broad gains across language and biological sequence tasks, including outperforming larger autoregressive models in selected settings.
  • At each P2 step, the denoiser predicts z while the planner selects positions to unmask or remask.

2 Background and preliminaries

Masked discrete diffusion models define corruption toward a fully masked prior and learn reverse denoising with a tokenwise network. Standard inference fixes tokens once unmasked, while P2 motivates more flexible updating.

  • The model represents length-L sequences over a finite vocabulary V containing a specialized mask token and uses categorical distributions on the probability simplex.
  • Discrete diffusion constructs a forward noising process that progressively converts the data distribution into a structureless, fully masked prior.
  • The simplified masked process independently corrupts unmasked tokens, with α_t decreasing from 1 to 0; once masked, a token remains masked in the forward process.
  • The reverse process iteratively denoises across tokens, and a time-independent denoiser predicts clean-token probabilities from the partially masked sequence.
  • Training the vanilla MDM corresponds to optimizing a weighted cross-entropy loss, and its ELBO supports the denoiser objective.
  • In the continuous-time limit, vanilla MDM inference uniformly samples masked positions, providing no control over generation order and motivating schemes that can change unmasked tokens.

3 Discrete Diffusion with Path Planning

Path Planning (P2) modifies masked diffusion inference by using a planner to select tokens for unmasking or resampling before denoising. It supports self-, BERT-, and trained-planner variants, incorporates remasking and stochasticity control, and establishes an expanded ELBO for planner-based sampling.

  • P2 Sampling Strategy: P2 generalizes existing MDM sampling strategies and is the only compared framework combining remasking, planning, and stochasticity control.The optimal planner can depend on the application.
  • P2 Sampling Strategy: P2 decomposes inference into planning and denoising, selecting positions to update after sampling a denoised sequence from the current partially masked sequence.The planner uses both the partially noised sequence and the denoiser’s predicted clean sequence.
  • P2 Sampling Strategy: P2 can unmask masked tokens and remask selected unmasked tokens before resampling them, enabling refinement of previously generated tokens.This differs from vanilla MDM inference, which denoises coordinates independently without planner-based reassignment.
  • P2 Sampling Strategy: The planner separates masked-token unmasking from unmasked-token retention, while a scheduler controls how many positions are guaranteed to be unmasked at each reverse step.A stochasticity parameter η controls the frequency of remasking, exposing an efficiency–self-correction trade-off.
  • A Family of Planners: P2 includes self-planning, BERT-planning, and trained-planning implementations that use denoiser probabilities, a pretrained BERT model, or a planner trained against the denoiser’s correctness.The trained planner freezes the denoiser and learns to prioritize positions on an optimal decoding trajectory.
  • A Family of Planners: P2 yields an Evidence Based Lower Bound that decomposes into denoiser, unmasking-planner, and masking-planner ELBO terms.The bound supports both planner training and evaluation of off-the-shelf planners.

4 Experiments

The experiments evaluate P2 across protein, language, and RNA generation, testing functional quality, structural quality, diversity, task completion, planner choices, and inference-time scaling. Across these settings, P2 improves generation quality, preserves diversity in protein and RNA evaluations, and benefits from additional sampling steps.

  • Experimental scope: P2 is evaluated across protein sequence generation, natural language generation, and RNA sequence generation using functional, structural, diversity, and task-completion measures.The experiments also include planner-choice ablations and inference-time scaling.
  • Protein sequence generation: P2 improves protein folding metrics while maintaining comparable entropy and diversity, with pLDDT rising from 80.23 to 83.45 and foldability from 48.14% to 58.86%.DPLM + P2 also outperforms reported baselines including EvoDiff, ESM3, and ProGen2 with fewer parameters.
  • Language generation: P2 consistently improves generation quality across five language benchmarks; on GSM8K, MDM rises from 58.5% to 60.9%, while DiffuLLaMA reaches 17.6% pass@1 on HumanEval.The HumanEval result exceeds ancestral sampling at 13.2% and LLaMA2 at 1.7%; ROCStories ROUGE scores increase by more than 5 absolute points.
  • RNA sequence generation: P2-Bert improves RNA pLDDT from 68.1 to 73.3 and lowers MFE from −48.5 to −51.9 while preserving GC content and entropy.The RNA evaluation generates 100 sequences of 100 base pairs and uses an external folding model for structural assessment.
  • Planner ablations: P2-Train achieves the highest protein pLDDT at 83.45 and foldability at 58.86%, while P2-Self and P2-Bert also produce consistent gains over comparison strategies.The ablation compares P2 variants with RDM, Greedy Ancestral, MaskGIT, and Top-K Marginal.
  • Planner ablations: P2-Self achieves the highest pass@1 and ROUGE scores in code and story benchmarks, outperforming ancestral decoding and RDM.These results use a 7B DiffuLLaMA model on HumanEval and ROCStories.
  • Inference-time scaling: P2 consistently improves protein foldability as sampling steps increase from 50 to 300 and maintains its advantage beyond 200 steps, where other methods plateau.The experiment generates 300 protein sequences of length 200 for each sampling-step setting.
  • Computational complexity: P2 introduces a runtime trade-off: no-planner strategies achieve 673.16 tokens/sec, while external planners incur additional cost.The comparison uses a 150M denoiser on a single NVIDIA A100 GPU.

5 Related Work

Recent MDM sampling methods improve generation order using heuristics or latent-order modeling, but Path Planning provides a unified formulation that generalizes these strategies.

  • Heuristic MDM samplers include greedy unmasking, remasking, and informed correctors, but lack structured guidance.
  • ReMDM-conf schedules Gibbs-corrector temperature, while Top-K Marginal uses local confidence gaps to choose tokens.
  • Path Planning generalizes these sampling strategies within a unified, optimizable, and principled formulation.
  • LO-ARM is identified as the closest related approach because it treats generation order as a latent variable.

6 Conclusion

The paper shows that unmasking order materially affects MDM quality and introduces Path Planning as a general sampling framework with broad gains across language and biological sequence tasks.

  • Unmasking order significantly impacts the generative performance of masked diffusion language models.
  • Path Planning expands the ELBO to optimize token selection during inference and generalizes all existing MDM sampling strategies.
  • Path Planning delivers state-of-the-art improvements across language generation and biological sequence design, enabling MDMs to outperform larger autoregressive models.

Ethics Statement

The paper describes benefits and misuse risks of stronger discrete generative models, while restricting biological experiments to computational benchmark evaluation without wet-lab validation.

  • More effective discrete generative models can advance reasoning, programming, and biomolecular design, including protein and RNA applications.
  • Protein and RNA sequence modeling may accelerate scientific discovery and therapeutic design, but could also enable harmful or dual-use biological sequence generation.
  • All biological experiments are computational and evaluated against public benchmarks; no wet-lab synthesis or functional validation was performed.
  • The datasets are publicly available and widely adopted, with no newly introduced private or sensitive data.

Reproducibility Statement

The implementation and experiments build on publicly available codebases for SMDM, DiffuLLaMA, and DPLM, with RNA results obtained by adapting DPLM for MDM training.

  • The authors provide a PyTorch implementation and integrate the approach into the SMDM GitHub codebase for the MDM (1.1B) + P2 results.
  • The DiffuLLaMA GitHub codebase supplies the DiffuLLaMA (7B) + P2 results reported in Table 3.
  • Protein experiments use the DPLM open-source codebase, while RNA experiments adapt DPLM for MDM training with RiNALMo-related components.

Appendices

This section motivates Path Planning by identifying the limitations of fixed unmasking in masked diffusion models and positioning P2 among related approaches. It highlights the need for theoretically grounded control over denoising order and refinement.

  • MDMs progressively unmask sequences, but unmasked tokens remain unchanged despite approximation, discretization, and numerical errors.
  • Existing correctors and guided samplers revise predictions, but the cited works do not provide the same planner-based framework for selecting updates.
  • The literature offers relatively sparse theoretical treatment of denoising-order selection, with prior approaches using fixed order families or specially trained external planners.
  • DDPD also separates planning from denoising, but analyzes uniform discrete diffusion and depends on training a large external planner.

B Proofs of Propositions 1 and 2

The appendixed proofs establish the expanded ELBO underlying P2 and characterize how planners should select unmasking and remasking actions. They also derive training objectives for learned planners under masked diffusion dynamics.

  • Proposition 1: The proof constructs a time-discretized ELBO for coordinate-wise conditionally independent masked diffusion dynamics and decomposes it into planner, denoising, and related terms.
  • Proposition 1: The derivation uses discrete reverse-process transitions initialized from a fully masked state and establishes the lower-bound expression through limiting sums and integrals.
  • Proof techniques: The proof handles the ELBO terms using reconstruction limits, Jensen’s inequality, cross-entropy identities, and bounds on logarithmic terms.
  • Proposition 2: The learned planner objective trains a single network to handle both keeping previously unmasked tokens and selecting tokens to unmask.
  • Proposition 2: The optimal planner probabilities are tied to conditional probabilities under the data distribution for both unmasked-token retention and masked-token selection.

D.2 Proof of the ELBO Proposition 1: CTMC Version

The CTMC appendix extends the ELBO analysis to continuous time and compares P2 with existing sampling schemes. It emphasizes that P2 modifies trajectories while targeting data-distribution fidelity at the endpoint.

  • CTMC ELBO: The continuous-time argument defines an evidence-based lower bound for the Markov chain induced by P2 and relates it to the log marginal likelihood.
  • Scope of the bound: The ELBO analysis bounds the endpoint distribution’s divergence while explicitly noting that the entire CTMC trajectories need not be close.
  • Sampling dynamics: In the continuous-time limit, standard MDM sampling denoises one uniformly selected masked position, while P2 permits resampling previously unmasked tokens.
  • Comparison with existing methods: P2 subsumes ancestral, greedy ancestral, DFM, RDM, and DDPD-style sampling through different planner and stochasticity choices.
  • Comparison with DDPD: Unlike DDPD, P2 uses masked diffusion as its forward process and separates mask and unmask planners, while DDPD uses uniform discrete diffusion and a single external planner.

E Implementation Details

The implementation iteratively predicts masked positions, optionally re-scores unmasked positions with a planner, and re-masks low-scoring tokens. Experiments evaluate this procedure across protein-generation models and folding-based quality measures.

  • Sampling procedure: Path-planning sampling repeats model prediction, stochastic token sampling, planner-based rescoring, and selective re-masking across refinement iterations.
  • Sampling procedure: The stochasticity multiplier η controls how aggressively unmasked positions are updated during score-based remasking.
  • Configuration: The implementation exposes controls for refinement steps, temperature, masking schedules, optional planners, and confidence- versus random-based scoring.
  • Protein experiments: Protein experiments compare P2 with discrete diffusion, autoregressive, and masked-language models across multiple model sizes.
  • Evaluation: Foldability is the proportion of sequences satisfying pLDDT > 80, pTM > 0.7, and pAE < 10.
  • Evaluation: Sequence entropy and diversity are used to detect mode collapse, where generated sequences become highly repetitive.

F.2 Language Generation Evaluation Details

The language-generation evaluation covers five tasks spanning comprehension, prediction, reasoning, story infilling, and code completion, using standard task-specific metrics and established baselines. Experiments apply P2 across MDM and autoregressive model settings with tuned stochasticity and fixed evaluation protocols.

  • Benchmarks: The benchmark evaluates TriviaQA, LAMBADA, GSM8K, ROCStories, and HumanEval using exact match, accuracy, accuracy, ROUGE-1/2/L, and pass@1, respectively.These tasks represent reading comprehension, last-token prediction, math reasoning, story infilling, and code completion.
  • Benchmarks: The evaluation includes five language-generation task examples collected in Table S1.
  • Setup: Experiments use 1.1B MDM and 7B DiffuLLaMA base models, apply P2 with η ∈[0, 2.0], and report best-performing settings.Decoding uses standard ancestral sampling unless otherwise noted, with oracle length truncation for autoregressive baselines lacking native infilling.
  • Baselines: Published comparisons include GPT2, DiffuGPT, SEDD, Plaid1B, LLaMA2, and TinyLlama autoregressive baselines.
  • Setup: Stochasticity η is tuned separately for each model-task pair using a grid sweep while keeping the evaluation set fixed.The experiments do not use instruction tuning or chain-of-thought prompting.

F.4 RNA Evaluation Details

The RNA evaluation trains and tests masked diffusion models on generated sequences, measuring structure confidence, folding energy, entropy, and GC content against pretrained and natural-sequence references. P2 improves structural quality beyond native and pretrained baselines while preserving nearly unchanged sequence diversity.

  • Training and evaluation: The RNA model is trained on 27M RNACentral sequences with a 150M-parameter MDM for 100K steps.Training uses 320K tokens per batch and follows the RiNALMo tokenizer and vocabulary.
  • Training and evaluation: Evaluation generates 100 RNA sequences of 100 base pairs and obtains predicted structures with an RNA folding model.
  • Metrics: The metrics are pLDDT for predicted local structure confidence, MFE for minimum free energy, entropy for mean token entropy, and GC content for guanine-cytosine proportion.
  • Baselines and configuration: Baselines include RiNALMo-150M, RiNALMo-650M, and 100 natural RNA sequences of matching length, with BERT-Planning applied to the RiNALMo-150M-derived planner.The P2 stochasticity parameter η is swept from 0 to 2 in increments of 0.02, reporting the best configuration.
  • Results: P2 improves RNA structural quality beyond native and pretrained baselines while keeping sequence diversity nearly unchanged.Structure visualizations are provided separately.

G.2.3 Ablation of Path Planning

The ablation studies examine how P2’s planning, remasking, stochasticity, planner choice, and inference-time scaling affect protein and RNA generation. Across these studies, P2 improves structural or folding quality while retaining diversity, and lightweight external planning can preserve efficiency.

  • Design space: P2 generalizes vanilla ancestral, greedy ancestral, RDM, and DFM sampling by combining controllable stochasticity with planning guidance.Protein foldability is used as the principal ablation metric.
  • Sampling strategies: Zero-stochasticity vanilla and greedy ancestral sampling disable remasking and perform poorly, whereas RDM enables remasking and uses denoiser confidence for self-planning.
  • Planner size: An 8M BERT planner guides a 150M MDM competitively with a 3B planner across a broad range of entropy values.The BERT planner also preserves foldability better than self-planning at extreme entropy settings.
  • Planner ablation: P2 outperforms existing sampling strategies across folding metrics while maintaining strong token entropy and sequence diversity.Against DDPD with the same BERT planner, P2 separates remasking from unmasking, reducing planner bias and using denoiser planning capabilities.
  • Inference efficiency: P2 with self-planning retains ancestral sampling’s inference cost, while an 8M BERT planner adds an extra evaluation but has negligible overhead relative to the 150M MDM.At higher sampling scales, their performance gap becomes indistinguishable.
  • Robustness: P2 maintains strong mean performance across 20 independent DPLM+P2 runs despite non-negligible variance, particularly for pLDDT.
Loading 2502.03540v5…