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 cannot reliably correct incorrectly denoised tokens, limiting generative quality. Path Planning (P2) uses an inference-time planner to select tokens for updating and resampling, achieving state-of-the-art improvements across language and biological sequence tasks while enabling MDMs to outperform larger autoregressive models.

  • Problem

    Masked diffusion inference cannot course-correct incorrectly denoised tokens, while random unmasking orders are suboptimal for imperfect denoisers.

    Source: §1 · p. 2
  • Method

    Path Planning expands the ELBO with an inference-time planner that selects tokens for updating, including resampling previously denoised tokens.

    Source: §1 · p. 2, §3.1 · p. 4
  • Results

    P2 generalizes major masked diffusion sampling strategies and delivers state-of-the-art improvements across diverse language, code, and biological sequence tasks.

    Source: §1 · p. 2, §4.4 · p. 10, §6 · p. 11
  • Takeaways & Limitations

    Inference-time token-selection strategies are important for improving discrete diffusion generation and can enable MDMs to outperform larger autoregressive models.

    Source: §6 · p. 11

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

    Source: Ethics Statement · p. 11

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

Path Planning (P2) improves masked diffusion model inference by selecting tokens to unmask or resample, enabling correction of errors that fixed uniform denoising cannot address. The framework generalizes existing MDM sampling strategies and delivers strong results across reasoning, code generation, protein design, and RNA design.

Motivation

  • P2 addresses the underexplored inference bottleneck in absorbing-state diffusion models, whose practical sampling has largely remained limited to simple uniform denoising despite advances in training.Discrete diffusion is especially appealing for domains such as biological sequences that lack a natural causal ordering, although autoregressive models remain more successful overall.Source: §1 · p. 1
  • Uniform unmasking is optimal only with a perfect denoiser; with imperfect denoisers, non-uniform planning can improve quality by correcting errors and limiting error propagation.Current MDM inference cannot revise incorrectly denoised tokens at later steps, making uniformly random unmasking suboptimal in many settings.Source: §1 · p. 2

Method

  • P2 introduces inference-time planners that select tokens for unmasking and optionally remask already unmasked tokens, thereby enabling refinement under an expanded ELBO.The framework offers self-planning with the denoiser, BERT-type planning, and lightweight planners trained offline, and generalizes known MDM sampling strategies.Source: §1 · p. 2, Figure 1 · p. 2

Empirical validation

  • P2 with a 1B-parameter MDM outperforms a 7B Llama model in math reasoning, surpasses same-sized autoregressive models for code, and achieves state-of-the-art protein and RNA design quality.For RNA design, generated sequences have higher structural plausibility than naturally occurring sequences.Source: §1 · p. 2

2 Background and preliminaries

The paper formulates discrete sequences over a finite vocabulary with a designated mask token and introduces masked discrete diffusion models that corrupt data into fully masked sequences before reverse denoising. Vanilla MDMs use irreversible unmasking, which fixes generated tokens and provides no control over generation order.

2 Background and preliminaries

  • A finite vocabulary V contains d categories, including a specialized mask token, and sequences are length-L elements of V^L represented by categorical distributions on the d-dimensional simplex.The data distribution is an empirical distribution over n training sequences.Source: §2 · p. 2

2.1 Masked Discrete Diffusion Models

  • Masked discrete diffusion models progressively corrupt data sequences into a fully masked prior and reverse this process through iterative denoising across discretized time steps.The forward process independently masks unmasked tokens according to a decreasing noise schedule, while the reverse process generates a clean sequence from the fully masked state.Source: §2.1 · p. 3
  • The reverse process is parameterized by a time-independent denoiser Dθ that predicts clean-token probabilities, with training induced by an ELBO and equivalent weighted cross-entropy objective.Reverse transitions factorize across tokens for the simplified masking process.Source: §2.1 · p. 3
  • Vanilla MDMs irreversibly fix unmasked tokens, and in the continuous-time limit independently updated denoising selects masked positions uniformly without controlling generation order.The paper motivates an inference scheme that can resample unmasked tokens to any vocabulary item.Source: §2.1 · p. 3

3 Discrete Diffusion with Path Planning

P2 augments masked diffusion sampling with a planner that selects tokens for unmasking or remasking, enabling iterative refinement of partially denoised sequences. The framework supports multiple planner designs and a finer-grained ELBO for planner-based generation.

3 Discrete Diffusion with Path Planning

  • P2 generalizes existing MDM samplers and uniquely combines remasking, planning, and stochasticity control, with the preferred planner depending on the application.The framework therefore encompasses established sampling behavior while adding iterative refinement of previously unmasked tokens.Source: §3 · pp. 3–4, Table 1 · p. 4

3.1 The P2 Sampling Strategy

  • P2 introduces planning before denoising, allowing the sampler to select tokens for updating and resample incorrect unmasked tokens.The planner uses both the partially noised sequence and the denoiser’s predicted clean sequence to guide denoising paths.Source: §3 · p. 3, §3.1 · p. 4

3.2 Designing the Planner

  • The planner separates masked-token unmasking from unmasked-token retention, while a scheduler controls update counts and η trades efficiency against self-correction.This construction is designed to guarantee convergence to a fully unmasked sequence and permits multiple tokens to change per iteration.Source: §3.2 · pp. 4–5

3.3 A Family of Planners: Instantiations of P2

  • The planner family comprises Self-Planning, BERT-planning, and Trained-Planner, with the trained planner learning to prioritize positions where denoiser predictions match ground truth.Self-Planning uses denoiser probabilities, BERT-planning uses a pretrained BERT model, and Trained-Planner fine-tunes a planner while freezing the denoiser.Source: §3.3 · pp. 5–6
  • P2 establishes an evidence-based lower bound on log marginal likelihood that decomposes into denoiser, unmasking-planner, and masking-planner ELBO terms.The bound supports training and evaluating planners, and planners from 8M to 3B parameters achieve similar ELBOs and generation performance.Source: §3.3 · p. 6

4 Experiments

Experiments evaluate P2 across protein, language, and RNA generation, showing improved quality and task performance while generally preserving diversity. Ablations and scaling studies further assess planner variants, inference steps, and computational tradeoffs.

4.1 Protein Sequence Generation

  • 10.72-point foldability gain: DPLM with P2 increases foldability from 48.14% to 58.86% and pLDDT from 80.23 to 83.45 versus RDM, while maintaining comparable entropy and diversity.DPLM + P2 outperforms EvoDiff, ESM3, and 2.7B ProGen2 with fewer parameters and better foldability; foldability uses pLDDT > 80, pTM > 0.7, and pAE < 10.Source: Table 2 · p. 8, §4.1 · p. 8

4.2 Language Generation

  • 17.6% pass@1: DiffuLLaMA with P2 outperforms ancestral sampling at 13.2% and LLaMA2 at 1.7% on HumanEval, while improving all five language benchmarks.On GSM8K, P2 improves MDM from 58.5% to 60.9%, surpassing LLaMA2 at 58.6%; ROCStories ROUGE scores increase by more than 5 absolute points.Source: Table 3 · p. 8, §4.2 · p. 9

4.3 RNA Sequence Generation

  • 5.2-point pLDDT gain: BERT-Planning raises RNA pLDDT from 68.1 to 73.3, lowers MFE from −48.5 to −51.9, and increases GC content from 60.8% to 65.5%, with only a small entropy reduction.The 150M MDM already surpasses RiNALMo baselines in structural quality and energy before applying P2.Source: Table 4 · p. 9, §4.3 · p. 9

4.4 Ablation Studies

  • P2-Train achieves the highest protein pLDDT at 83.45 and foldability at 58.86%, while P2-Self and P2-BERT also consistently improve sampling strategies.P2-Self achieves the highest pass@1 and ROUGE scores in HumanEval and ROCStories, and P2-BERT improves RNA pLDDT and MFE while preserving GC content and entropy.Source: §4.4 · p. 9, Table 5 · p. 10, Table 6 · p. 10
  • P2 generalizes Vanilla, Greedy, RDM, and DFM through appropriate planner configurations, with P2-Self, P2-BERT, and P2-Train retaining diversity while improving structural and functional accuracy.Self-planning often outperforms BERT-based planning because it better fits the underlying denoiser.Source: §4.4 · p. 10

4.5 Inference-Time Scaling and Computational Complexity

  • 673.16 tokens/sec: no-planner strategies provide the highest throughput, while the 8M P2 planner achieves 509.55 tokens/sec with 24% overhead and maintains quality gains beyond 200 sampling steps.P2 foldability consistently improves as sampling steps increase from 50 to 300, whereas other methods plateau beyond 200 steps.Source: Figure 3 · p. 10, Figure 4 · p. 10, §4.5 · p. 11

5 Related Work

Prior MDM sampling methods use heuristics, order-based strategies, or costly planning, while P2 provides lightweight, dynamic, mask-aware planning compatible with frozen denoisers. Compared with related approaches, P2 offers a unified, differentiable ELBO formulation for efficient generation-policy learning.

5 Related Work

  • P2 introduces lightweight, modular, dynamic, mask-aware planning compatible with frozen denoisers, contrasting with heuristic, costly, fixed-schedule, or mask-unaware alternatives.Prior approaches include greedy unmasking, remasking, informed correctors, AOARMs, and DDPD, which respectively lack structured guidance, may require costly planners or fixed schedules, or operate without mask-awareness.Source: §5 · p. 11
  • P2 unifies temperature scheduling, Top-K confidence-gap selection, and related generation-order strategies within an optimizable, principled path-planning formulation.ReMDM-conf schedules Gibbs-distribution temperature, while Kim et al. (2025) use local confidence gaps for Top-K selection.Source: §5 · p. 11
  • Compared with LO-ARM’s high-variance REINFORCE learning, P2 uses a simple, differentiable ELBO objective to enable efficient and scalable generation-policy learning.LO-ARM (Wang et al., 2025c) treats generation order as a latent variable and learns it via REINFORCE.Source: §5 · p. 11

6 Conclusion

The conclusion shows that inference-time unmasking order significantly affects masked diffusion model performance. Path Planning (P2) expands the ELBO to optimize token selection, generalizes existing sampling strategies, and improves results across language and biological sequence tasks.

  • Unmasking order significantly impacts masked diffusion model generative performance, underscoring inference strategy as a key factor in discrete diffusion.The conclusion frames improved inference as a route toward more efficient and effective sequence generation.Source: §6 · p. 11
  • Path Planning (P2) introduces an expanded-ELBO planner that optimizes token selection during inference.Source: §6 · p. 11
  • P2 generalizes all existing masked diffusion model sampling strategies and delivers state-of-the-art improvements across language generation and biological sequence design.These improvements enable masked diffusion models to outperform larger autoregressive models.Source: §6 · p. 11

Ethics Statement

The work acknowledges that improved discrete generative modeling may advance science while creating risks of misuse, especially for biological sequence generation. It mitigates these risks through computational-only biological evaluation, public benchmarks, open-science documentation, and explicit discouragement of malicious use.

Potential impacts

  • P2 may benefit reasoning, programming, biomolecular design, scientific discovery, and therapeutic development, but could also enable harmful or dual-use biological sequence generation.The authors explicitly discourage and do not support malicious applications.Source: Ethics Statement · p. 11

Risk mitigation

  • All biological experiments were computational and used standard public benchmarks, with no wet-lab synthesis or functional validation.This limits direct experimental enablement while leaving biological misuse risks acknowledged.Source: Ethics Statement · p. 11

Data and release practices

  • The study uses publicly available datasets without introducing private or sensitive data, and releases models with documentation of intended use and limitations.These practices support open science while discouraging misuse.Source: Ethics Statement · p. 12

Reproducibility Statement

The authors provide a PyTorch implementation and reproduce experiments by integrating P2 into established open-source codebases for language, protein, and RNA sequence generation.

Implementation

Appendices · A Related Works: Extended Discussion

Prior MDM research improves training and language-modeling performance, but irreversible unmasking remains vulnerable to approximation, discretization, and numerical errors. Related methods enable correction or planning yet lack P2’s external-planner flexibility or comparable theoretical treatment.

A Related Works: Extended Discussion

  • MDMs use continuous-time variational objectives based on weighted cross-entropy integrals, enabling state-dependent masking schedules and improved training efficiency.At GPT-2 scale, recent models outperform prior diffusion language models and show stronger zero-shot language modeling (Nie et al., 2025; Gong et al., 2025).Source: A Related Works: Extended Discussion · p. 19
  • MDMs progressively unmask sequences, but fixed unmasked tokens can remain incorrect because of approximation, discretization, and numerical errors.Once unmasked, tokens stay unchanged despite imperfect data-distribution fitting and additional sampling errors.Source: A Related Works: Extended Discussion · p. 19
  • Prior MDMs begin from fully masked inputs and progressively reveal positions, simplifying generation but preventing iterative refinement of earlier predictions.This fixed-update behavior is the central limitation motivating methods that revise prior predictions or guide sampling trajectories.Source: A Related Works: Extended Discussion · p. 19
  • Correction methods include informed correctors (Zhao et al., 2024b), greedy ancestral methods (Gong et al., 2025), and RDM techniques (Zheng et al., 2023; Wang et al., 2024).These approaches use model scores to replace random masking with targeted corrections, but none supports an external planner; RDM uses top-k sampling without theoretical justification.Source: A Related Works: Extended Discussion · p. 19
  • Theoretically grounded denoising-order selection remains sparse: Li et al. (2021) trains a specialized external planner, while Shih et al. (2022) fixes a priori generation-order families.These works approach order selection through Any-Order Autoregressive models and differ in how they address architectural cost and redundancy.Source: A Related Works: Extended Discussion · p. 19
  • DDPD (Liu et al., 2024) separates planning from denoising but relies on uniform diffusion, a specialized planner, and lacks P2’s masked/unmasked planner separation.The paper compares DDPD with P2 in Algorithm 4, emphasizing differences in planner roles and greater flexibility.Source: A Related Works: Extended Discussion · p. 19

B Proofs of Propositions 1 and 2 … D.1 P2 Continuous Time Formulation

This section develops the discrete- and continuous-time ELBO and Markov-chain formulations underlying P2, specifying how the planner and denoiser govern unmasking and resampling dynamics. It also proves the continuous-time evidence lower bound, characterizes the optimal planner under its ELBO-informed loss, and analyzes the associated stochasticity and limiting transition rates.

B.1 Proof of Proposition 1: Time Discretization Approach

  • The time-discretized P2 dynamics converge as T→∞ to a continuous-time Markov chain with an evidence-based lower bound on log marginal likelihood.The proof starts from the standard discrete-time diffusion ELBO, shows the reconstruction term vanishes in the limit, and converts time sums into integrals.Source: §B.1 · pp. 19–21

B.2 Finding the Optimal Planner Under the ELBO-informed Loss: Proposition 2

  • The ELBO-optimal planner keeps previously unmasked tokens according to their conditional data probability and unmasks denoiser-suggested tokens according to their data probability.Using one network for both planner roles provides meaningful training signals for correctly and incorrectly denoised inputs.Source: §B.2 · p. 24

B.3 Form of the ELBO for Varying η

  • Increasing η shifts the ELBO and planner loss toward the masked-token role G_M and away from the previously unmasked-token role G_U.For arbitrary η≥0, the resulting bound remains decomposed into masked-planning, unmasked-planning, and denoising terms.Source: §B.3 · p. 26

B.4 Proof of Lemma 1

  • The continuous-time limits used in the proof follow from independent coordinate dynamics: transitions vanish when states differ in multiple coordinates and yield single-coordinate rates otherwise.The lemma derives these limits by analyzing coordinatewise transitions as T→∞.Source: §B.4 · p. 27

B.5 Deriving the Discrete Time ELBO Eq. 10

  • The appendix derives Eq. 10 by applying Jensen’s inequality to the reference path distribution and marginalizing variables from the resulting terms.The derivation identifies the discrete-time ELBO used for training and connects it to the reverse-process transitions.Source: §B.5 · p. 28

C Additional Background: Continuous Time Perspective

  • Continuous-time discrete diffusion models use time-inhomogeneous Markov chains whose rate matrices determine marginals and sample paths through jump times.The framework interpolates data distributions with an easily sampled reference distribution and simulates a denoising chain toward pdata.Source: §C.2 · p. 29, §C.1 · p. 29

C.3 The Role of the Denoiser and the Approximate Backwards Process

  • The denoiser approximates unknown reverse-process rates induced by pdata, while conditional coordinate independence enables parallel simulation but introduces an approximation assumption.The conditional backward process is built from independently evolving coordinates, and the denoiser supplies the quantities needed to construct the approximate reverse rates.Source: §C.3 · p. 30, §C.4 · p. 30

C.5 Masked Diffusion Model: Continuous Time Formulation

  • Masked diffusion augments the token vocabulary with a mask state and uses an α_t-controlled masking transition, whose reverse rates are parameterized by denoiser predictions.The masked formulation is reported to outperform uniform diffusion and other transition choices across applications, and its discrete-time approximation follows from the continuous-time rates.Source: §C.5 · pp. 30–31

C.6 Role of the ELBO

  • The continuous-time ELBO minimizes path-space divergence so the approximate reverse process matches the exact process, yielding the masked-diffusion objective Eq. 4.Girsanov’s theorem gives an expression in terms of the denoiser, interpolation, and reverse rates; in the masked setting, this expression is exactly E_D from Proposition 1.Source: §C.6 · p. 32

D.1 P2 Continuous Time Formulation

  • P2 modifies the approximate reverse jump matrix with a planner G_ϕ that estimates whether each position should be unmasked or resampled, including already unmasked tokens.For masked positions, the planner estimates unmasking probabilities; for unmasked positions, it estimates probabilities of changing the current token based on the remaining sequence context.Source: §D.1 · pp. 32–33
  • The P2 formulation separates selecting which position should be unmasked or resampled from determining what it should be (re)sampled to.The planner selects the position, while ˆDθ determines the resampled value.Source: §D.1 · p. 33
  • The discrete-time sampling scheme approximates the CTMC Xθ,ϕ with rate matrix Qθ,ϕ, matching the continuous-time limit in Eq. 14.This approximation uses the same formal discrete-time procedure discussed above Eq. 29.Source: §D.1 · p. 33

D.2 Proof of the ELBO Proposition 1: CTMC Version … D.5 Deriving the P2 Gillespie Scheme Algorithm 5

The appendices establish the CTMC ELBO argument, connect masked diffusion sampling to any-order autoregressive generation, and position P2 against existing planners. They then derive an exact Gillespie formulation and its practical approximation for P2 sampling.

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

  • P2’s CTMC ELBO proof applies Girsanov’s theorem after reversing time, reducing the argument to four terms that match the time-discretized proof.The proof lower-bounds −DKL(Px0||Pθ,ϕ) and carries over the existing E(x0) bound after the time change s = 1 − t.Source: §D.2 · pp. 33–34

D.3 Equivalence of MDMs with AOARMs

  • Masked diffusion sampling is equivalent to any-order autoregressive sampling, with each unmasking revealing conditional information that can motivate backtracking and remasking earlier tokens.The corresponding standard masked-diffusion Gillespie sampler initializes all tokens masked and repeatedly chooses coordinates for unmasking.Source: §D.3 · pp. 34–35

D.4 Comparison with Other Sampling Methods

  • Existing MDM samplers fit P2 by varying remasking, stochasticity, and planner roles: ancestral sampling disables remasking, DFM enables it, RDM uses self-planning, and DDPD uses external planning.DDPD differs from P2 by disallowing mask-informed planning and decomposition of the planner into unmasked and masked components.Source: §D.4 · p. 36

P (x0) + EDDPD

  • P2 and DDPD assign different training responsibilities: P2 reuses the denoiser ELBO and can use lightweight or denoiser-based planners, whereas DDPD’s denoiser objective depends on its planner and requires a comparable-size external planner.The P2 planner can use partially masked inputs and denoiser outputs, while DDPD’s planner sees randomly flipped, unmasked sequences and cannot separate masked and unmasked planning or introduce η analogously.Source: P (x0) + EDDPD · p. 37
  • P2-BERT, P2-train, and P2-self are argued to outperform DDPD because distinguishing clean from noisy unmasked tokens is a bottleneck for uniform-diffusion planners.The cited comparison is reported in Table S5, while the appendix explains that uniform diffusion can reverse to masked diffusion when conditioning on whether tokens are clean or noise.Source: P (x0) + EDDPD · p. 38

D.5 Deriving the P2 Gillespie Scheme Algorithm 5

  • The exact P2 Gillespie sampler derives jump-chain transition probabilities from the CTMC rate matrix, then samples a coordinate to change followed by its replacement value.The derivation identifies transitions at Hamming distance one and yields Algorithm 5 after defining the relevant coordinate and token distributions.Source: §D.5 · pp. 38–40
  • Because exact planner expectations are impractical, Algorithm 5 approximates them by assuming that conditioning on one coordinate has little effect on the planner’s sequence-level expectation.The practical Algorithm 1 replaces the exact planner with η-controlled ˜Gη and top-k sampling, while dropping immediate denoising after remasking.Source: §D.5 · pp. 39–40

E Implementation Details … F.1.3 Computing the ELBO

The paper provides a PyTorch implementation of Path-Planning Sampling, details protein-generation evaluation and training setups, and specifies an ELBO computation protocol for comparing planners. The implementation iteratively samples, re-scores, and re-masks tokens, while experiments define structural-quality metrics and stochastic evaluation procedures.

E Implementation Details

  • Path-Planning Sampling iteratively predicts masked positions, optionally re-scores unmasked positions with a planner, and re-masks the lowest-scoring non-fixed positions.The procedure runs for num_steps, scales unmasked scores by η, and replaces remaining masks with the last sampled tokens.Source: E Implementation Details · pp. 40–41

F.1 Protein Generation Evaluation Details

  • Protein-generation evaluation compares P2 with DPLM, EvoDiff, ESM3, ProGen2, and ESM2 across multiple model scales and sequence lengths.P2 stochasticity is tuned from 1.0 to 2.0, while baselines use default or specified sampling settings; each model generates 100 sequences per length.Source: §F.1 · pp. 42–43
  • Foldability is the proportion of sequences satisfying pLDDT > 80, pTM > 0.7, and pAE < 10, complementing folding scores with entropy and sequence diversity to detect mode collapse.The combined thresholds address cases where high local accuracy coexists with poor global compactness.Source: §F.1 · p. 43

F.1.1 Training Details of the 150M MDM.

  • The self-planning ablation trains a 150M mask diffusion model from scratch for 500k steps on approximately 40M UniRef50 protein sequences using the DPLM-150M architecture.Training uses 320K tokens per iteration with multi-GPU and multi-node gradient accumulation.Source: §F.1.1 · p. 43

F.1.2 Training Details for P2 Train

  • For P2 train, ESM-8M is used for 100k steps with GU = GM = Tϕ in Alg. 2 and the same data and hyperparameter setup as for the 150M MDM.Source: §F.1.2 · p. 43

F.1.3 Computing the ELBO

  • ELBO evaluation compares five planners—ESM-8M, ESM-35M, ESM-150M, ESM-650M, and ESM-3B—with self-planning using weighted cross-entropy reconstruction loss.The negative ELBO is computed on protein sequence generation tasks to assess data fit.Source: §F.1.3 · p. 43
  • ELBO experiments use UniRef50 test sequences shorter than 300 residues, a 1022-token sequence length, a 60,000-token budget, and averages over 20 independent simulations.These restrictions align with Figure S3 experiments and mitigate memory constraints.Source: §F.1.3 · p. 43
  • For ELBO computation, masking ratios are sampled uniformly from [1/500, 1 −1/500], and masked and unmasked prediction losses are separately weighted inversely by the mask ratio.Non-maskable special tokens are preserved, and denoiser and planner predictions are evaluated across both position categories.Source: §F.1.3 · p. 44

F.2 Language Generation Evaluation Details · F.3 RNA Generation Details

The evaluation covers five language-generation tasks spanning comprehension, prediction, reasoning, story infilling, and code completion, using established protocols, base models, published baselines, and tuned P2 sampling. The supplied passages provide no substantive details for F.3 RNA Generation Details.

F.2 Language Generation Evaluation Details

  • Language evaluation spans TriviaQA for reading comprehension, LAMBADA for last-token prediction, and GSM8K for math reasoning.TriviaQA uses exact match, while LAMBADA and GSM8K use accuracy.Source: §F.2 · p. 44
  • ROCStories evaluates story infilling with ROUGE-1/2/L, while HumanEval evaluates code completion with pass@1.The cited benchmarks use the metrics specified for each task.Source: §F.2 · p. 44
  • Table S1 supplies examples for the five language-generation tasks used in the evaluation.The passage identifies the table’s scope but does not provide its example contents.Source: §F.2 · p. 44
  • Experiments follow SMDM (Gong et al., 2025) and DiffuLLaMA (Nie et al., 2025) protocols using MDM (1.1B) and DiffuLLaMA (7B) base models.P2 uses η ∈[0, 2.0], with best-performing settings reported; decoding is standard ancestral sampling unless noted.Source: §F.2 · p. 44
  • For autoregressive baselines without native infilling, evaluation uses oracle length truncation, and results are measured with the LM Harness (Biderman et al., 2024).This protocol addresses infilling incompatibility while retaining the stated evaluation framework.Source: §F.2 · p. 44
  • Reported baselines include GPT2-S/M, DiffuGPT, SEDD (Lou et al., 2023), Plaid1B (Gulrajani & Hashimoto, 2023), and LLaMA2 (Touvron et al., 2023).TinyLlama is additionally included as an open-source autoregressive baseline.Source: §F.2 · p. 44
  • P2 evaluation tunes η by grid sweep for each model-task pair while holding the evaluation set fixed and using no instruction tuning or CoT prompting.The implementation notes state that stochasticity is critical to P2 quality.Source: §F.2 · p. 44

F.4 RNA Evaluation Details … G.2.2 Ablation over model scale for ProGen2 and DPLM.

Across RNA, reasoning, code, and protein-generation evaluations, P2 improves quality by planning updates and selectively revising masked positions. Results span RNA structural quality, reverse-query performance, code completion and infilling, and protein-model scaling analyses.

F.4 RNA Evaluation Details

  • The RNA MDM is a 150M-parameter model trained on 27M RNACentral sequences for 100K steps with 320K-token batches, using the RiNALMo tokenizer and vocabulary.For P2, BERT-Planning derived from RiNALMo-150M sweeps η ∈[0, 2] in increments of 0.02 and reports the best configuration.Source: §F.4 · pp. 44–45
  • P2 improves MDM RNA structural quality beyond native and pretrained baselines while keeping sequence diversity nearly unchanged.Evaluation generates 100 sequences of 100 base pairs and measures pLDDT, MFE, entropy, and GC Content.Source: §F.4 · pp. 44–45

G.1.1 Breaking the Reverse Curse

  • P2 improves reverse-direction accuracy for Description To Name and Name To Description BLEU in both directions, where T5 and ARMs otherwise achieve zero accuracy and low BLEU.The evaluation uses exact-match accuracy and BLEU on fictitious bidirectional statements and reversals.Source: §G.1.1 · p. 46

G.1.2 Additional Comparison Among Sampling Methods

  • +17.5 Pass@1 on HumanEval and +14.9 on MBPP versus Vanilla Ancestral, while P2-self-plan consistently leads all six coding tasks across Pass@1 and Pass@10.P2 also improves infilling and FIM, whereas heuristic gains saturate because they cannot revise earlier decisions.Source: §G.1.2 · p. 47, Table S3 · p. 47
  • P2-self-plan jointly chooses unmasking and selective remasking through a lookahead objective, outperforming stochastic, greedy, entropy-based, and TopK-Margin sampling.All methods use the same Open-dCoder base model and inference budget; only masked-position update rules differ.Source: §G.1.2 · pp. 46–47

G.2.1 Performance Across Length Categories.

  • Protein-generation performance is analyzed across sequence-length categories spanning 200–800 base pairs because some models do not generate fixed lengths.Figure S1 compares model capabilities and limitations across these categories.Source: §G.2.1 · p. 47, Figure S1 · p. 47

G.2.2 Ablation over model scale for ProGen2 and DPLM.

  • P2 with a Trained Planner at 8M consistently improves DPLM variants, while increasing model scale alone does not ensure better performance.Table S4 reports the model-scale ablation for ProGen2 and DPLM.Source: Table S4 · p. 48

G.2.3 Ablation of Path Planning … G.3.2 Visualizing the Predicted Structures of Generated RNA Sequences.

Across protein and RNA generation, P2 improves sampling quality, efficiency, robustness, and structural plausibility by combining remasking control with planning guidance. Its benefits persist across model sizes, sequence lengths, alternative baselines, and generated structures, while RNA training follows a RiNALMo-based implementation.

G.2.3 Ablation of Path Planning

  • P2 consistently outperforms existing sampling strategies across protein folding metrics, with further gains from an external BERT planner and better foldability than DDPD through denoiser-guided unmasking.P2 separates remasking, delegated to the planner, from unmasking, guided by the denoiser, mitigating planner bias and avoiding DDPD’s low-entropy repetitive sequences.Source: Table S5 · p. 49, §G.2.3 · p. 50

G.2.6 Ablation Study on the Planner.

  • An 8M BERT planner guides a 150M MDM competitively with a 3B planner across entropy values, while BERT planning preserves foldability more robustly than self-planning at extreme temperatures.Planner size affects pLDDT and diversity, emphasizing its role in generated-sequence quality and consistency.Source: Figure S3 · p. 49, §G.2.3 · p. 50, §G.2.6 · p. 53

G.2.11 Variance Analysis of P2.; G.2.12 Generated Protein Sequences and Their Predicted Structures.

  • Self-planning uses informative unmasked-token logits, as randomizing them substantially degrades performance across MDM variants and self-planning attains an even better ELBO than BERT planners.Across 20 DPLM+P2 runs, performance means remain strong despite non-negligible variance, especially for pLDDT; generated protein structures also show diversity and consistency across lengths 200–800.Source: §G.2.3 · p. 50, Table S7 · p. 50, Table S6 · p. 50

G.2.7 Inference-Time Scaling: Performance vs. Sampling Time.

  • Increasing sampling steps generally improves quality but increases computation, and the available inference-time analysis identifies performance–time scaling as a key efficiency trade-off for P2.The detailed scaling results are deferred to future work.Source: §G.2.4 · p. 52, §G.2.7 · p. 53

G.2.4 Sampling Efficiency

  • P2 self-planning achieves higher foldability than ancestral sampling, while an 8M BERT planner scales effectively at higher sampling steps with negligible overhead relative to the 150M MDM.Self-planning has the same inference cost as ancestral sampling; the BERT planner adds an evaluation but remains lightweight, making the performance gap indistinguishable at larger scales.Source: Figure S4 · p. 51, §G.2.4 · p. 52

G.2.5 Design Space of P2.

  • P2 balances structural accuracy and diversity across pLDDT, pAE, pTM, entropy, and diversity metrics, demonstrating versatility in protein generation.The design space is characterized by these structural-quality and diversity measures.Source: §G.2.5 · p. 52, Figure S5 · p. 52

G.2.8 Performance on Short Protein Sequences (<200 residues).

  • P2 improves shorter protein sequences of 64–200 residues by consistently increasing pLDDT and pTM while reducing pAE, although gains are smaller because of training underrepresentation and less accurate ESMFold evaluation.The short-sequence results are reported in Table S8.Source: §G.2.8 · p. 53, Table S8 · p. 54

G.2.9 Comparison with Additional Baselines (ESM2); G.2.10 Comparison with Top-K Marginal.

  • P2 enhances generation quality across protein language models, yields particularly large gains with DPLM, and substantially outperforms Top-K Marginal on pLDDT, pTM, pAE, and entropy.These comparisons include ESM3, ProGen2, EvoDiff, and Top-K Marginal (Kim et al., 2025), with P2 subsuming Top-K Marginal as a special case.Source: §G.2.10 · p. 54, Table S10 · p. 54, §G.2.9 · p. 54
Loading 2502.03540v5…