Planner Aware Path Learning in Diffusion Language Models Training
Authors: Fred Zhangzhi Peng
Zachary Bezemek
Jarrid Rector-Brooks
Shuibai Zhang
Michael M. Bronstein
Anru Zhang
Joey Bose
Alexander Tong
Venue: ICLR 2026 Oral (2026)
Keywords: Diffusion Language Models
Discrete Diffusion
Diffusion Models
code generation
protein generation
text generation
Primary area: generative models
TL;DR: We propose Planner Aware Path Learning (PAPL), a simple planner-aligned training method for Diffusion Language Models that resolves the training–inference mismatch and consistently improves generation quality.

Abstract: 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 engineered sampling strategies, or *planners*, that select more favorable generation paths by iteratively planning---versus uniformly at random---where to denoise along the sequence. However, by modifying the reverse paths via planning, planners create an irrevocable mismatch between the uniformly random denoising paths during training and planning-based inference. In this paper, we systematically investigate the mismatch of discrete diffusion training and inference under planning and theoretically prove that the standard discrete diffusion training evidence lower bound (ELBO) does not accurately describe a denoiser that uses a non-uniform planner. To address this gap, we derive a new planned evidence lower bound (P-ELBO) that incorporates planner-based reverse dynamics directly into the training objective.
Using the P-ELBO, we introduce *Planner Aware Path Learning* (PAPL), a novel training scheme that aligns training and inference under a planned denoiser.
PAPL is implemented as a simple yet effective modification to the standard masked discrete diffusion loss, making it widely applicable and easy to adopt.
Empirically, we show PAPL delivers consistent gains across domains, including a 40\% relative improvement in protein sequences, improved text generation with up to a $4\times$ relative MAUVE gain, and 23\% relative improvement in code generation HumanEval pass@10.

PLANNER AWARE PATH LEARNING IN DIFFUSION LAN- GUAGE MODELS TRAINING
Fred Zhangzhi Peng -I;; Pengl.. Zachary Bezemek!..e Jarrid Rector-Brooks" Shuibai Zhang?,
Anru R. Zhang', Michael Bronstein Alexander Tong"t, Avishek Joey 2,688
California Institute of Technology 5 University of Wisconsin-Madison 6 University of Oxford AITHYRA
1 Duke University Mila 3 Université de Montréal
8 Imperial College London 1 Equal contribution
f Equal advising
ABSTRACT
Diffusion language models have emerged as a powerful alternative to autore- gressive models, enabling fast inference through more flexible and parallel generation paths. This flexibility of sampling is unlocked by new engineered sampling strategies, planners, that select more favorable generation paths by iteratively planning-versen uniformly at random-where to denoise along the sequence. However, by modifying the reverse paths via planning, planners create an irrevocable mismatch between the uniformly random denoising paths assumed during training and planning-basee inference. In this paper, we systematically investigate the mismatch of discrete diffusion training and inference under planning and theoretically prove that the standard discrete diffusion training evidence lower bound (ELBO) does not accurately describe a denoiser that uses non-uniform planner. To address this gap, we derive new planned evidence lower bound (P-ELBO) that incorporates planner-based reverse dynamics directly into the training objective. Using the P-ELBO, we introduce Planner Aware Path Learning (PAPL), novel training scheme that aligns training and inference under planned denoiser. PAPL is implemented as a simple yet effective modification to the stan- dard masked discrete diffusion loss, making it widely applicable and easy adopt. Empirically, we show PAPL delivers consistent gains across domains, including a 40% relative improvement in protein sequences, improved text generation with up to a 4x relative MAUVE gain, and 23% relative improvement in code generation HUMANEVAL pass 10. Code available at github. hattlubbcoom/enggggggg///////
INTRODUCTION
The landscape of generative modeling over discrete data has led to foundational breakthroughs in deep learning, with Large Language Models (LLMs) being exemplary technology that has transcended heyond natural language processing (Achiam et al., 2023). Until recently, the de facto gold standard for building [.LMs has been Autoregressive models (ARMs), which are highly scalable for pre-training LLMs-allowe rigid inference schemes due to the autoregression mechanism that generates samples causal order-e.g., left to right for natural language (Guo et al., 2025). In contrast to ARMs, recent advances in Diffusion Language Models (DLMs) have the potential to disrupt the current status quo for generative modeling of discrete data, as they natively support flexible generation orders and allow for fully parallel sampling of tokens inference time (Austin et al., 2021a; Lou et al., 2024; Sahoo et al., 2024; Shi et al., 2024). The increased flexibility of modeling discrete data in any order makes DLMS arguably a more natural tool than ARMs for tackling high-impact problem domains that lack a natural causal ordering, such as biological sequence design and code completion (Nic et al., 2025a; Wang et al., 2024; 2025b), spurring their rapid recent development and application (Song et al., 2025).
them to capture complex dependencies in data-but incur
Indeed, the most performant variant of DLMs, i.e., Masked Diffusion Models (MDMs) (Shi et al., 2024; Sahoo et al., 2024), approach generative modeling as a denoising task, wherein partially masked
"Correspondence 2p708duke. edu and zub@duke.edu

---

sequences are iteratively refined using learned denoiser that time-reverses the Markov transition dy- namics of the masked corruption process. However, a key assumption, and thus a central limitation, of DLMs is that following the reverse dynamics of uniformly denoising position at inference implicitly assumes denoising with perfect denoiser (Peng et al., 2025a). As a result, in practice, to fully take advantage of flexible generation paths and generate higher-quality samples beyond uniform decoding, the reverse process must be modified by a planner: a rule that selects which tokens to reveal next, such as greedy decoding (Chang et al., 2022), ancestral sampling (Shi et al., 2024; Schiff et al., 2025), or Path planning (P2) (Peng et al., 2025a). In fact, using a planning strategy is more than a humble artefact of optimizing for inference; it can be seen as avoiding denoising over exponential infilling training, in which the task is provably computationally intractable (Kim
problems at inference-unlin et al., 2025). More than just theory, employing a planner when denoising has been shown to sub- stantially improve sample quality across various application domains, including text, code, protein sequences, and discrete image modeling (Peng et al., 2025a; Nie et al., 2025a; Shi et al., 2024).
A fundamental aspect of DLM inference under planning is the fact that denoising is not necessarily conducted by uniformly picking a position to unmask. Consequently, this new reverse process creates irrevocable mismatch between the forward masking process, used during training, which comuppt sequences by masking positions uniformly at random. This mismatch of forward and reverse processes also suggests that, in effect, training denoisers in DLMS attempts solve harder problem than the one they are ultimately used for. This raises the following central question:
Q. How should we adapt the training of denoisers in diffusion language models when inference inevitably proceeds under planner?
Present work. In this paper, we seek to answer the question by introducing a new theoretical framework that aims to align the training of DLMs with pre-assumed knowledge of planner-based inference. Our framework is built using basic facts of Markov chains, which allows us to set up the training of DLMs as minimizing a path-wise KL divergence. More precisely, the path-wise KL is between the reverse dynamics of DLM using a planner and supervised ideal reverse dynamics, also under planning, that hit the data distribution.
Armed with this path-wise KL, we first theoretically prove in $3.1 how greedy ancestral sampling at inference of DLMs violates the standard DLM ELBO (Sahoo et al., 2024; Shi et al., 2024) the thesis of mismatch of forward and reverse dynamics under planning. We next derive new planned evidence lower bound (P-ELBO) in $3.2, of which the standard planning- aaannng-ggnstiic--- i.e., uniformly at random denoising-DLLM ELBO is a special case. Furthermore, we demonstrate recent heuristics that act as planners, such as MaskGIT (Chang et al., 2022), and P2 (Peng et al., 2025a), emerge as principled instances of our new P-ELBO, which unifies existing strategies under one unbrella. Given the insights in P-ELBO, we propose new loss function for training the denoiser that directly incorporates any choice of planner, and thereby allows training to match inference properly. We summarize our main contributions in this paper as follows:
Unifying framework. We derive a novel generalized planner-aware generalized lower bound (P-ELBO) that takes into account the use of planning in the reverse dynamics of a DLM.
Efficient implementation. Starting from the P-ELBO, we design new simplified loss termed Planner Aware Path Learning (PAPL) that amounts to a one-line code change and uses self- self-planning. Specifically, PAPL leverages the denoiser itself-ic... places where the denoiser is most confident- to compute a weighted loss more likely generation paths compared to standard DLMs.
Improved performance. Empirically, PAPL. consistently improves the quality of diffusion language models under identical configurations. We observe PAPL in protein sequence generation yields a 40% relative increase in foldability, surpassing larger diffusion and autoregressive baselines while preserving diversity. On code generation, PAPL improves HUMANEVAL pass@ from 18.5 to 20.8, pass@ @ 10 from 31.1 to 38.4, and HUMANEVAL-ING INFILL pass @ from 30.0 32.5. On text generation, PAPL achieves up to a 4x improvement in MAUVE and reduces generative perplexity by over 40% compared to prior diffusion models.
2 BACKGROUND AND PRELIMINARIES
Notation. Let V- = (1....,d d} denote a finite vocabulary. We reserve the final symbol, d = m, as a special mask token, while the remaining d symbols correspond to ordinary vocabulary items. We
2

---

(x² ......""") V The empirical data distribution Pdata is supported on a training set D C V We denote by A = {u E R u² 2 0, Σi=1 u' = 1} the probability simplex. Each u E A specifies a categori- cal distribution Cat(j: u) = u² over j E V. For particular token I E V, we write 5(z) E A for the one-hot distribution that places all its mass on z. To avoid ambiguity, we use superscripts (e.g.. x²) to index sequence positions, and subscripts (e.g., .Ft) to index time steps of stochastic process. For x, y E VL, we use diAMM(X)) y) to denote the Hamming distance hetween x and y. We also use NM to denote the number of coordinates in x which are equal to m. For a finite set S, we use Unif(S) to denote the uniform distribution on that set. For / < n E N, we denote by [[ n] = {1,7 + I,..., n}.
consider sequences of length L, so that data point is represented
2.1 MASKED DIFFUSION LANGUAGE MODELS
A masked diffusion language model (MDLM) generates samples from a data distribution Pdata E A 4 through an iterative sampling process which gradually denoises a full mask sequence [m] % to a sequence which does not contain any masks. This iterative sampling procedure makes use of denoiser Dg VL - (Ad)4, which outputs a distribution over the clean tokens at each position. In particular, for x E V and y E V, # m, Cat(y; D2 (x)) approximates the probability that the i'th token a sequence is y given the unmasked positions in the sequence match those of x (Hoogeboom et al., 2022; Sahoo al., 2024; Shi et al., 2024; Zheng et al., 2025; Ou et al., 2025).
Moreover, the Gillespic sampling scheme (Gillespie, 1977; 1976) allows us to establish an exact equivalence between DLM and any-order autoregressive model (AOARM) (Uria et al., 2014; Hoogeboom et al., 2022). More precisely, the Gillespie sampling scheme allows for a single coordinate be denoised each step. Concretely, starting from the fully masked sequence, the procedure iteratively (I) selects a position uniformly at random among the currently masked tokens, (2) samples a replacement token from the denoiser's predictive distribution that position, and (3) updates the sequence by filling in the chosen token while leaving all other positions unchanged.
Let umif be the distribution on V of x resulting from applying the above iterative sampling proce- dure in which the masked coordinate to denoise is chosen uniformly at random. This uniform unmask- ing process the main sampling strategy employed by MDLMs (Sahoo al., 2024; Shi et al., 2024). The connection between MDLMs and AOARMs (Ou et al., 2025) can be seen more explicitly through the MDLM evidence lower bound (ELBO), which lower bounds the log marginal log(p (xo)).
log(pö unif (xo)) 2 gd,u (),uni((xo) = Eo-UUit(()
log (Cat(zg(), ( D'g Dx(xx (xp (x8((5))
(I)
= (0(LL-)
Exe-Unl((Xe
Σ
I k log (Cat (a0: Dz Db(xx)))
where ΣL is the set of all permutations of length L. Additionally, for a E EL, we denote a a<< i) as the first 1 elements of o. Thus, for x V4, x<())) E y has coordinates in a<< i) set those of x and the rest set to m, and for x E V4 and k E [0: LJ. (x) c D is the set of sequences which are the same as but with exactly coordinates masked.
Existing DLM Sampling Strategies. While the vanilla masked diffusion sampler proceeds by unmasking one position chosen uniformly at random, a variety of alternative planners have been proposed to improve generation quality by biasing the unmasking order. A straightforward modification is greedy decoding as in MaskGIT (Chang et al., 2022), where each step the denoiser selects the position with the highest confidence to unmask next (Gong et al., 2025). Another line of work introduces remasking, where previously generated tokens may be reverted to the mask state and resampled. Resampling diffusion models (RDM) (Zheng et al., 2023) extend greedy planning with resampling. More recently, path planning (P2) (Peng et al., 2025a) has been proposed as a unifying framework that generalizes all of the above strategies. P2 decomposes each step into planning stage, where a planner chooses positions to update-including including both masked and already unmasked tokens-and denoising stage, where the selected positions are resampled with the denoiser.
3

---

Standard Uniform Training
Planner Aware Path Learning
[M]
[M]
P
P
P
[M]
[M]
P
P
P
[M]
[M]
[M]
[M]
A
[M]
[M]
[
[M]
[M]
P
P
P
P
[M]
P
P
P
P
[M]
[M]
[M]
L
MM]
[M]
[M]
Figure 1: Planner-Aware Path Learning (PAPL) resolves training-inference mismatch in DLMs.
Standard uniform training for DLMs (left) applies uniform loss across all masked positions, distributing capacity over regions that inference-time planners never traverse. PAPL (right) introduces planner-aware weights into the loss, aligning training with the planner's preferred trajectories (outlined arrows) and eliminating training-inie inference mismatch.
METHOD: PATH LEARNING
We now present Plameer-wware Path Learning (PAPL), aiming align training with planner-based inference. In 1, 1 we introduce a general formulation of sampling with a planner and derive the corresponding transition probabilities. In $3.2. we show that greedy sampling modifies these dynam- ics in a way that departs from the MDLM ELBO, motivating the need for new objective. We then ELBO in(P-ELBO) $3.3, from which different sampling schemes-including derive a planner- ancestral-emerge emerge special cases ($A.2). Finally, in $3.4, we simplify the uniform and greedy ancestral
P-ELBO into an efficient training objective, leading to our final efficient (PAPL) algorithm (Alg. 1).
Discrete-ting time Markov Chains. The starting point for our analysis is re-examine the ELBO in Equa- tion (1) through the lens a continuous time Markov chain (CTMC) (Campbell al., 2022; 2024; Lou al., 2024; Sun et al., 2023). From the CTMC perspective, the sampling path of a DI.M can be com- pared against family of reference chains parameterized by sample Xg. which generates that specific datum. For ease of presentation, the ELBO via a discrete-time Markov chain of the DLM X to that of family of reference discrete Markov chains Y, before taking appropriate limits to recover the continuous perspective. Specifically, for datum we have discrete Markov chains of length L which satisfy X Y and Y2 Xg-
forego lifting this continuous perspective and instead compute
We now recall that the transition matrix Q for discrete time Markov chain X encodes Q x) P(X++1 Xx = z). From this perspective, for vanilla DLMs we have (x) = P(X) = x) where X [m] and x dynamics are described the transition matrix, given by:
( {o (G(()':
dHAM(x. 1, m, y m
(y,,
otherwise
where dHAM is the hamming distance and x, y VL, We compare y with y define the transition matrices as follows:
[m]] and
{o
Cat(1' (Calt):d(aa))/NN dHAM(x, () = 1, m, y m otherwise
R(y, x; xo)
(2)
In this view, the second expression of equation is can be rewritten as:
(x)) R(y,
g).uni = L.pp-Unit((o.L--))
Expors(sxo)
Σ R(y,xxx0)) log
(3)
VEVL
where we set (x(x; xo) the (discrete) path wise KL divergence in our Proposition $A.7. Importantly, this ELBO enjoys a simple interpretation in how the denoised coordinates incorporate coordinates chosen under a planner in Proposition 3.2. Moreover, this perspective is more easily amenable to modifying the sampling dynamics for X and deducing the corresponding ELBO.
P(Y x). We rederive the ELBO from this perspective by stating
chosen fact which will later develop to
4

---

3.1 REVERSE DYNAMICS WITH A PLANNER
We begin by introducing a planner function G. V x V2 - A that modifies the sampling process by selecting reverse transition which coordinate in the sequence to be denoised next. For simplicity, we assume the planner only selects masked positions, i.e. Cat(i; Gz( x)) = 0 whenever x² of m matching greedy ancestral sampling (Nie al., 2025a). Under these dynamies, each backwards step can be decomposed into two substeps as follows: 1.) starting from the current sequence Xk. the denoiser produces candidate predictions z (x(xx) for all positions.
2.) The planner then samples an index i ~ (z X and the token at this index is updated by setting " = z' while all other positions remain unchanged. Thus, each transition from X₁ to X differs exactly one coordinate.
The probability of unmasking the i-th coordinate of X to token y, after marginalizing over z can be explicitly written as a transition kernel db (3) xx.) as follows:
(v) xx) = Cat (y: D) Fe, Fx v.
(4)
Fo. (x, y, (x()) [Cat (i: (z x*))]
(5)
where z - denotes the same sample except that the i-th component has been replaced with y-
3.2 GREEDY ANCESTRAL VIOLATES THE VANILLA DLM ELBO
Greedy ancestral sampling (Nie et al., 2025a), as widely used in the literature (Gong et al., 2025; Besnier al., 2025) such MaskGIT (Chang et al., 2022), employs a specific choice of planner which selects the most confident position according to the denoiser itself,
G (x) = arg max Cat(z";Do(xu))
(6)
j:cr/ =m
Unfortunately, when using greedy sampling, the standard DLM ELBO may not, in fact, satisfy the ELBO inequality. Using the transition probabilities equation 4, we prove the following:
Proposition 3.1. For pi greedy (xo) defined with Go in equation 6 and Da an imperfect denoiser, we may have
log(rg gredd (xo)) < Jm.i (x0),
where ((.) is in equation 1.
We prove Proposition in $A.1.7. The key takeaway is that the ELBO in equation only valid for the uniform unmasking process pgnif, As result, the reverse dynamics of greedy sampling longer match those assumed by the standard training loss. In a nutshell, the model is being trained for uniform unmasking, but inference follows different process. More broadly, this insight applies to any planner: whenever the sampling procedure deviates from uniform, the training objective no longer strictly reflects the quality of the generated samples.
We remark that the proof of Proposition 3.1 hinges on constructing an explicit counter-example in particular assuming the denoiser inconsistent along different paths: That is, we assume our construction that there are two different permutations a,, o of [1: L) and Xo E V4 a clean sequence in
the data distribution such that IT Cat
(1), D' X
+ I Cat ( ((1), DJ
(x
a(<))
Indeed, for perfect denoiser, in the above we would have equality for any Xo. o, and and sampling along any path a would always result in a sample from the data distribution. In this situation, there would be no point of planning generation path. However, in practice there no relationship being enforced between these quantities, and it has been observed repeatedly in the literature that the "rath"="denoisingg order" taken greatly influences sample quality see, e.g. Ou al. (2025) Appendix J.4. Shih et al. (2022) Section 4, and Li et al. (2021) Section 6.
3.3 PLANNER-AWARE EVIDENCE LOWER BOUND
The key mismatch that vanilla DLM training assumes uniform unmasking, while inference instead follows planner. To correct this, we next introduce a planner-aware ELBO (P-ELBO) that explicitly accounts for the planner's ' role the reverse dynamics.
5

---

Go
Proposition 3.2. For any planner Go, let Po denote the distribution of x1, obtuined via the planner-guideds sampling scheme. Then have the following ELBO:
log(Ps @ (xo)) 2 g0 Ex:0(x)) (xo) = (xx((x)) + Eh.0 (xo),
E (x0) = (u-L-
Cat(i; xx) (Car(ad: Dà (xxx))
Go
(s0ca) "i=aa
Carli; (Xo, xx]) Fo.o 2h:
E c0 (xo) = (0:
E (-000)
Catli; (x(o; xx)) log
Go is the distribution at time of a Markov chain with initial data (m,.... m) and where T'k
transition rates as in equation but with 1/N (x) replaced by Go (Xx, x).
A proof of this ELBO from a self-contained, discrete-time Markov chain perspective can be found in 4. and from a continuous time Markov chain perspective in $A.3.3. The first term Ef g0.e resembles the standard DLM ELBO: it the peg-robbability of predicting the correct token in each masked position, but now weighted by the probability that the planner would choose that position next. In other words, it a planner-weighted cross-entropy.
The second term E, new. It appears only when the planner's 's decision can depend on the full clean target sequence Xy- Intuitively, it measures the gap between (a) the "ideal" planner that has access to the ground truth, and (b) the "effective" planner that only relies on the denoiser's 's predictions, and mathematically is the negative KL divergence between these distributions. We highlight that for the uniform planner setting, this term vanishes, recovering the standard DLM ELBO. Putting this together, the planner- aware loss used for training is simply the negative ELBO from Proposition 3.2 [(0, Φ) = - Exx [xx)) which minimizes the KL divergence between the
Go and the data. We curate results on how several common existing instantiations model distribution Pe
of planned denoising fall into this framework and their corresponding ELBOs in $A.2. We also
prove result identifying the form of the optimal planner for a fixed denoiser under the trainling loss associated the ELBO from Proposition 3.2 in $A.1.5.
3.4 EFFICIENT IMPLEMENTATION OF PLANNER AWARE PATH LEARNING (PAPL)
Greedy decoding is widely used at infer- and often improves sample quality. A natural idea is therefore to train the
Algorithm 1 PAPL Training
1:
while not converged do
denoiser under the same greedy planner but with corrections such that we optimize the P-ELBO. Unfortunately, the exact greedy
2:
Xu Paua(2)
3:
le Unif( 0 L 1))
X2 ~ U
5:
w² + Cat(i; Go (Xo, xx)) for all masked
ELBO from Cor. A.9 is computationally infeasible: simulating the greedy path
// Setting a recovers standard DLM training
LPAPL +- T (1 + arw") " log Cat(sb; (xx)))
for each data point requires l denoiser
8: Update parameters using &oLexe
evaluations at step k, whereas vanilla
9: end while
DLM training needs only single forward pass. Consequently, we aim to design an efficient algorithm that still falls within our theoretical framework of training under planning.
10: return Trained denoiser Do
Soft greedy planner. We relax the deterministic argmax in Equation (6) into a softmax:
Cat(j; (Z x)) x exp log Cat (st:
where 7 is the softmax temperature. This assigns higher weight to positions where the denoiser is confident, and reduces to uniform sampling as T oo and greedy sampling as 7 4 0. Specializing Prop. 3.2 to GT yields planner-weighted cross-entropy plus complex correction term as outlined in Corollary 10. Detaching gradients through the planner removes the correction term, leaving only the simple weighted cross-entropy. This makes the objective stable and efficient.

---

Stabilization. Instead of simulating planner-driven paths to sample x;, reuse the vanilla DLM masking scheme: mask L k positions uniformly at random. This keeps training as cheap as standard DLM. In practice, the pure weighted loss have high variance, as shown in the training curves Fig. 5. We stabilize this by interpolating with the vanilla DL.M loss. The resultant effect is that the uniform weights 1/(L. - k) are replaced with planner-adjusted weights: z-x (1 + aw"), xx)) and a controls the strength of planner weighting.
where the weights are w' x Cat(i;
Final objective. The resulting Planner-Aware Path Learning (PAPL) loss is therefore just the standard masked diffusion cross-entropy, augmented with planner weights:
CPAPI. (0) = - Exxg, Ni,
+ + arw") log (Cat(ali D2(x)))) Σ Bm)
(7)
This amounts to a one-line modification of the vanilla DLM loss, making PAPL. easy to adopt. We detail the connection between the softmax regularized training objective from Corollary A. and the PAPL training objective equation in $A.2.3.
4 EXPERIMENTS
We evaluate PAPL on protein sequence generation, text and code generation, domains that demand non-trivial structure. Across domains, we compare against (i) the masked diffusion language model (DLM) baseline (same architecture, size, and training setup as PAPL), (ii) prior autoregressive and diffusion-based models (cf. $B for details).
4.1 PROTEIN SEQUENCE GENERATION
Setup. We evaluate PAPL. on the task of protein sequence generation. We train 150M DLM baseline and PAPL- augmented variant under identical configurations. During in- ference, both models use P2 sampling. For evaluation, each model generates 100 sequences at lengths 200, 300,.... 800, which folded into 3D structures using ESMFold (Lin et al., 2023). Structural quality is measured by pLDDD, D'TM, and PAE; diversity is measured by token-level entropy and sequence uniqueness. To obtain a single robust metric of functionality, we define a sequence as foldable if it simul- taneously satisfies pLDDT > 80, pTM > 0.7, and pAE < 10. Comparisons include diffusion-bart baselines (EvoDiff (Alamdari al., 2023), DPLM (Wang et al., 2024)) and autoregressive baselines (ESM3 (Hayes et al., 2025), ProGen2 (Nijkamp et al., 2023)). See $B.1 1 for further details.
pLDDDT:92
pL.DT Longth 389
pLDD0 Length 450
200
PLDOT 900
PIDDTT: Length 700
PLDDTT 700
Figure 2: Visualization of PAPL gen- erated proteins folded with ESMFold.
Table I: Protein sequence generation benchmark. We evaluate structure quality via PLDDT, pTM, and PAE, and diversity via token entropy and sequence uniqueness. Foldability is the percentage of sequences satisfying PLDDT 80, pTM 0.7, and pAE < 10.
Model
PLDDTT
PTMT
PΛEd
Foldability (%)
EntrupyT
Diversity (R)T
ESM3
34.13
0.23
24.65
1.50
3.99
93.44
Large
ProGen2-medium
$7.94
0.38
20.81
12.75
2.91
91.45
ProGen2-large
55.07
0.35
22.00
11.87
2.73
91.48
DPLM-65UM
79.53
0.66
11.85
49.14
3,18
92.22
150M 150Mascale sccae
EvoDiff
31.84
0.21
24.76
0.43
4.05
93.19
ProGen2-ssmall
49.38
0.28
23.38
4.48
2.55
89,31
DPLM.
0.65
12.07
DLM-I
81.32
0.65
12.00
42.43
3,21
92.45
DLM-150M PAPL (ours)
81.48
0.72
8.97
59,40
3.12
91.73
Results. Table 1 reports quantitative results. Compared to the DLM- baseline, PAPL achieves
higher PLDDT (81.48 vs. 81.32), stronger pTM (0.72 vs. 0.65), lower PAE (8.97 vs. 12.00), and yields a 40% relative increase in foldability (59.40% vs. 42.43%). Importantly, entropy (3.12) and diversity (91.73%) remain on par with the baseline, confirming that improved folding does not induce collapse. PAPL-trained models outperform Evo and ESM3 in all structural metrics, compared to DPLM-650M and ProGen2-2.7 2..6n2.22.B... Figure visualizes 3D folds of sequences generated by PAPL., which exhibit well-formed helices and sheets with coherent tertiary organization.
when
7

---

Takeaway. By aligning the training objective with the inference-times planner, PAPL significantly improves the structural plausibility of generated proteins. This is achieved under identical model configurations and without sacrificing diversity, outperforming all baselines.
4.2 TEXT GENERATION
Setup. We evaluate PAPL on unconditional text generation using the OPENWEBTEXT (Gokaslan & Cohen, 2019) (OWT) corpus, a large-scale collection of web pages designed to replicate the distribution of OpenAl's Web'Text. Text is tokenized with the GPT-2 byte-pair tokenizer, and sequences are wrapped to a maximum length of L = 1024 tokens. We compare against both autoregressive and diffusion-based baselines, including an autoregressive GPT--style language model, standard MDLMs, and MDLMs equipped with wrwaaddbackwwrd (FB) and discrete flow matching (DFM) correctors. All checkpoints are reused from prior work to ensure comparability. In inference, we generate 5,000 sequences using planner- planner-based decoding with P2 sampling (Peng et al., 2025a). We evaluate generation quality and diversity primarily with MAUVE (Pillutla al., 2021), which measures the divergence between generated and reference distributions. As secondary metrics, we also report generative perplexity (Gen PPL.) and entropy.
Table 2: Unconditional text generation. For each sampling step T, we report MAUVE (higher is better), generative perplexity (Gen PPL; lower is better), and entropy (higher is better). indicates nucleus sampling. For each T, the best diffusion scores are bolded.
T
32
T
64
T
128
Method
MAUVE
PPL
Eni.
MAUVE
PPL
Eni.
MAUVE
PPL
Ent.
(ref.)
1.00
14.8
5.44
1.00
14.8
ARI
0.760
1.21
5.22
0.760
1.21
5.22
0.760
1.21
5.22
MDLM1
0.006
100.45
5.60
0.011
72.08
5.55
0.015
61.5
5.52
MDLM+FB
0.007
95.76
5.56
0.016
59.05
5.49
0.064
42.8
5,44
M+DMM
0.004
303.8
5.31
0.007
108.8
5.33
0.041
37.9
5.31
ReMDMf
0.007
93.53
5,58
0.016
60.38
5,51
0.057
42.5
5.43
PAPL1 (ours)
0.013
40.19
5,32
29.98
24.33
Results. Table 2 reports unconditional text generation performance across sampling steps T E [32, 64, 128). Our method, PAPI., consistently and substantially improves diffusion-based generation. A 128, PAPL attains the strongest diffusion MAUVE (0.067) and lowest Gen PPL (24.3), outperforming ReMDM (0.057 MAUVE, 42.5 PPL) and MDLM+DFM (0.041 MAUVE, 37.9 PPL). entropy remains comparable across all methods (5.1-5.6), indicating that PAPL's gains in quality and perplexity are not driven by mode collapse.
Takeaways. PAPL delivers consistent and significant improvements models across all sampling budgets. These gains hold under fast sampling regimes (T < L.). While diffusion approaches still trail autoregressive models in absolute quality, PAPL markedly reduces this gap without sacrificing diversity.
prior discrete diffusion
4.3
CODE GENERATION
Setup. We evaluate PAPL on code generation, a domain requiring both syntactic validity and se- mantic correctness. Following the Open-dLLM framework (Peng et al., 2025b), we initialize from Swen2.5-Coder and adapt it to the diffu- sion setting with bidirectional attention. Models are trained on the FineCode corpus, which com- bines algorithmic and QA-style data, and are op- timized with masked cross-entroph loss as in the Open-dLLM recipe. Evaluation covers HU- MANEVAL (Chen et al., 2021), MBPP (Austin et al., 2021b), and their augmented variants, as well as SIMMANEAAL-INFLLL and the Python subset of SANTACODER FIM. We report pass and pass @ 10 for completion tasks, and exact match for infilling, using official protocols. See Section B.3.1 for more details.
Table 3: Code infilling performance HUMANEVAL- INFILL PASS @1 and SANTACODER-FIM EXACT MATCH. Large-scale models (27B) are shown as ref- erence, while the main comparison is among compact sub-billion models.
Model
HomanEval
SantaCroder
Reference Models (2:11)
LLaDA-8B
48.3
35.
Dream-78
39.4
40.7
DifliCoke-7B
54.8
38.8
Dream Coder-7B
55.3
40.0
Compact Model (S18) DLM (0.5B)
DLM (0.58) + PAPL, (Ours)
32.5
32_3
DLM (0.5B) PAPL (Oracle leugth)
77.4
56.1
8

---

Table 4: Code generation performance on HUMANEVAL, HUMANEVAL+, MBPP, and MBPP+. Large-scale models (27B) are shown as reference, while the main comparison is among compact sub-billion models. Results marked with adopted from prior work (Havasi et al., 2025).
HUMANBVAL
HUMANBVAL-
MBPP
MBPP+
Pass
Pass@
Reference Models (27B)
LLaDA (8R)
35.4
50.0
30.5
43.3
38.8
53.4
52.6
69.1
Dream (7B)
56.7
59.2
50.0
53.7
55.4
56.2
71.3
72.5
Compact Models (:S (B)
Autoregressive! (1.38)
17.0
34.7
14.0
28.6
25.6
45.4
Mask DFM (1.38)
9.1
17.6
7.9
13.4
6.2
25.0
Edit Flow (1.3B)
24.3
10.4
20.7
10.0
Flow (I.3R)
24.3
DLM (0,5B)
18.5
31.1
17.5
28.0
17.6
32.6
17.6
32.6
DLM (0.5B) + PAPL (Ours)
20.8
38.4
17.6
35.2
16.7
38.4
23,9
53.6
Results. Table 4 shows that PAPL consistently outperforms the baseline diffusion model across all completion benchmarks. On HUMANEVAL, pass improves from 18.5 to 20.8 and pass @ 10 from 31.1 to 38.4. Similar gains are observed on HUMANEVAL+ and MBPP+, where PAPL improves pass @ 10 by
than +10 points. This disproportionate improvement pass @ 10 suggests that PAPL not just improving the single best prediction, but is learning a more robust generative distribution the entire solution space, making highly effective generating a diverse candidates. Table reports results on infilling. Here too, PAPL improves over the baseline: pass @ increases from 30.0 to 32.5 on CEEAAAAEVAL-NNLLLL and exact-match accuracy rises from 30.7 to 32.3 on SANTACODER-FAEI settings, confirming that planner-aware training better aligns the denoiser with the planner-based reverse process.
of high-quality
These improvements hold under identical configurations and inference
Takeaway. Across both completion and infilling tasks, PAPL consistently improves the correctness of code generation. These results highlight the generality of our approach: learning better generative paths benefits tasks requiring strict logical structure just as it does those requiring biological fidelity.
18C
20
03
1E
65
10
12
(8) Perf. Training Steps
Perf. vs. Sampling Steps
(c) Perf. vs. Temperature (Training from scratch)
(d) Perf. vs. Temperature (Finetuning)
(ours)
Figure 3: PAPL consistently improves over DLM across training, sampling steps, and temperature. (a) Faster convergence in training steps. (b) Higher performance across sampling steps. (c) More robust temperature when training from scratch. (d) More robust temperature when fine-tuning.
4.4 ABLATION
Head-to-head comparison. We compare PAPL with the vanilla DLM baseline across training and inference conditions. Figure 3 shows that PAPL converges faster during training (panel a), maintains stronger performance across different sampling steps (panel b), and is substantially more robust to temperature variation, both when trained from scratch and when fine-tnedd (panels c-d). These results confirm that PAPL improves final quality while also stabilizing optimization and inference dynamics.
Hyperparameter tuning. We study the two key components introduced by PAPL: the softmax temperature 7 and the path learning weight a. As shown in Fig. 4, lowering below the default (7 1) consistently improves foldability, while larger values hurt performance, suggesting that sharper planner distributions provide more effective supervision. Increasing strengthens performance up to = 5, demonstrating that emphasizing seanner-wiighedd paths enhances stability and final quality. Increasing beyond this task lowers performance indicating the usefulness of interpolating

---

between MDLM and PAPL losses. Runs for these weights were stopped early to save compute and so do not appear in Fig. 4.
RELATED WORK
Masked diffusion language models (DLMs) have recently emerged as promising alternatives to autoregressive models for discrete sequence generation (Sahoo et al., 2024; Shi al., 2024; Nie et al., 2025a; Gong et al., 2025). A large body of work has focused on improving sam- pling efficiency and quality through heuristic strategies such as greedy unmasking (Gong et al., 2025), iterative remasking (Zheng et al., 2023; Wang et al., 2024), informed correc- tors (Zhao et al., 2025), and planner-guided approaches like P2 (Peng et al., 2025a) and confidence-planning O) While effective, these methods assume models trained under uniform denoising order and modify the sampling path only at inference time, creating a mismatch between training and generation.
Rolteaitty (4)
8
02 Tompastrnc (t)
Wright 60
Temporaurs Stuupy
WaytsSwap
Figure Effect of r and a on foldability. Lower (< 1) improves performance. Increasing steadily boosts foldability to = 5. The dashed line denotes the vanilla DLM baseline.
A related line of work investigates generation order in the setting of any-order autoregressive models (AOARMs). Shih et al. (2022) propose restricting the model to a fixed family of orders to reduce redundancy, whereas LO-ARM (Wang al., 2025c) and (Li et al., 2021) treat generation order as a latent variable and learn it with REINFORCE (Williams, 1992). However, the reliance on high-variance policy gradient estimators in LO-ARM limits its scalability large models and datasets. DDPD (Liu et al., 2025) trains for a planner-selected denoising order, where the planner effectively a uniform discrete diffusion model. This methodology requires planner of similar or greater size to the denoiser, and hence suffers from a similar pitfall to its AOARM counterparts.
CONCLUSION
In this work, we investigate the role of planning at inference time with respect to denoiser training for DLMs. Through this, we identify a mismatch in popular planner-guide guided inference paths and standard denoiser training that uses uniformly at random masking. We propose new Planner-Aware Evidence Lower Bound (P-ELBO) and developed a practical algorithm, Planner- Aware Path Learning (PAPL), to align the denoiser training with its intended use at inference. We demonstrate that the P-ELBO recovers all current planning strategies, and in particular enjoys straightforward implementation with PAPL being single line code change with additional computational overhead from standard DLM training. Empirically, PAPL achieves a 40% relative increase in protein foldability, 4x MAUVE gain in text generation, and 23% relative improvement in code generation HumanEval. demonstrating the benefits of reconciling the training and sampling processes. While PAPL uses the denoiser's own confidence plan, there are many other possible planning functions presenting ripe directions to extend and generalize the PAPL algorithm. We discuss how other unmasking schemes such as "top probability margin Kim al. (2025). RDM Zheng et al. (2023), Top-k "block-denoising Nie et al. (2025b), and "Confidence Thresholding" " Wu et al. (2025) can fit into a properly adapted version of our mathematical framework in A..2 and $A.5. Expanding this analysis to find computationally viable loss analogous to what the PAPL loss of equation 7 provides for greedy ancestral an interesting avenue for future work. Finally, we remark that our papers strength is also its limitation: while we show that one should modify the ELBO in order to account for alternative decoding strategies in MDMs, this means that some amount of post-training is necessary in order to test whether the decoding strategies performance be improved via using planner-aware loss. This training overhead makes our methodology expensive to implement with large planning models compared to just testing performance inference time.
10

---

ACKNOWLEDGMENTS
Fred extends sincere gratitude to Kaiwen Zheng for his invaluable insights. Zack extends his grat- itude Jim Nolen for his support and insightful discussions. The authors acknowledge funding from UNIQUE, CIFAR, NSERC, Intel, Samsung, as well as the Hartwell Foundation and CHDI Foundation. The research was enabled in part by computational resources provided by the Digital Re- search Alliance of Canada (https: / //aallaannccallaaaaaaaaaaaaaaaae and NVIDIA. This research partially supported by the EPSRC Turing AT World-Leading Research Fellowship No. EP/X040062/1 and EPSRC AI Hub No. EP/Y028872/1. Z.B. is partially supported by NSF-DMS award 2038056. F.Z.P. and A.R.Z. are partially supported by NIH ROTHI.169347.
ca). Mila (ht tnttps:: : /mi la quebec),
ETHICS STATEMENT
This work introduces Planner Aware Path Learning (PAPL), a training framework for discrete diffusion models. Our experiments limited to publicly available datascts in code and biological sequence generation. All results are computational; no biological synthesis or wet-lab experiments were performed. While improved generative models may be misused (e.g., generating harmful biological sequences or insecure code). we explicitly discourage such applications and release models only with documentation of intended use and limitations.
REPRODUCIBILITY STATEMENT
We provide detailed descriptions of model architectures, training objectives, datasets, and evaluation protocols in the main text and appendix. Hyperparameters, training schedules, and implementa- tion details are included enable replication. Code and pretrained models will be released upon publication to support full reproducibility.

---

REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, lige Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et Gpt-4 technical report. arXiv preprint arXiv:2303.087/4 2023. 1
Sarah Alamdari, Nitya Thakkar, Rianne van den Berg, Alex X. Lu, Nicolo Fusi, Ava P. Amini, and Kevin Kaichuang Yang. Protein generation with evolutionary diffusion: sequence is all you need. In NeurlPS 2023 Generative Al and Biology (GenBio) Workshop, 2023. 4.1, B.1
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in neural information processing systems, 34:1 19877-1999999 2021a.
Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang. Carrie Cai, Michael Terry. Quoe Le, et al. Program synthesis with large language models. arXiv preprint arXiv:2108.077: 2021b. 4.3, B.3.1
Mohammad Bavarian, Heewoo Jun, Nikolas Tezak, John Schulman, Christine McLeavey, Jerry Tworek, and Mark Chen. Efficient training of language models to fill in the middle. arXiv preprint arXiv:2207. 14255, 2022. B
Victor Besnier, Mickael Chen, David Hurych, Eduardo Valle, and Matthieu Cord. Halton scheduler for masked generative image transformer. In The Thirteenth International Conference on Learning Representations, 2025. 3.2
Amarjit Budhiraja and Paul Dupuis. Analysis and Approximation of Rare Events: Representations and Weak Convergence Methods, volume 94 of Probability Theory and Stochastic Modelling. Springer US New York, NY, 2019. A.3.3
Andrew Camphell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet. continuous time framework for discrete denoising models. In Advances in Neural Information Processing Systems, volume 35, pp. 28266-28279 Curran Associates, Inc., 2022. 3, A.3
Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Generative flows on discrete state-spaces: enabling multimodal flows with applications protein co-design. In Proceedings of the 41st International Conference on Machine Learning. ICML'24, 2024. 3, A.3, B.2.1
Huiwen Chang. Han Zhang, Lu Jiang, Ce Liu, and William T. Freeman. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 11315-11325, 2022. 1, 2.1, 3.2, B.2.2
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.0333 2021. 4.3, B.3.1
Daniel T Gillespie. A general method for numerically simulating the stochastic time evolution of coupled chemical reactions. Journal of Computational Physics, 22(4):403-434, 1976. ISSN 0021-9991. 2.1
Daniel Gillespie. Exact stochastic simulation of coupled chemical reactions. The Journal of Physical Chemistry, 81(25):2340-2360 23611, 1977. ISSN 0022-3654. 2.1
Aaron Gokaslan and Vanya Cohen. Openwebtext corpus. OpenWebTextCorpus. 2019. 4.2
io/
Shansan Gong, Shivam Agarwal, Yizhe Zhang, Jiacheng Ye, Lin Zheng, Mukai Li, Chenxin An, Peilin Zhao, Wei Bi, Jiawei Han. Hao Peng, and Lingpeng Kong. Scaling diffusion language models via adaptation from autoregressive models. In The Thirteenth International Conference on Learning Representations, 2025. 5.2,
12

---

Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Peiyi Wang, Qihao Zhu, Runxin Xu, Ruoyu Zhang, Shirong Ma, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in Ilms via reinforcement learning. arXiv preprint 2rXhh:2501.12944 2025.
Marton Havasi, Brian Karrer, Itai Gat, and Ricky TQ Chen. Edit flows: Flow matching with edit operations. arXiv preprint arXiv:2506.0900 2025.
Thomas Hayes, Roshan Rao, Halil Akin, Nicholas Sofroniew, Deniz Oktay, Zeming Lin, Robert Verkuil, Vincent Tran, Jonathan Deaton, Marius Wiggert, et al. Simulating 500 million years of cvolution with a language model. Science, 38766736;:50-8888 2025. 4.1, B.1
Emiel Hoogeboom, Alexey A. Gritsenko, Jasmijn Bastings, Ben Poole, Rianne van den Berg, and Tim Salimans. Autoregressive diffusion models. 10th International Conference on Learning Representations, 2022. 2.1
Jean Jacod and Albert Shiryaev. Limit theorems for stochastic processes, volume 288. Springer Science & Business Media, 2013. A.3.3
Jaeyeon Kim, Kulin Shah, Vasilis Kontonis, Sham M. Kakade, and Sitan Chen. Train for the worst, plan for the best: Understanding token ordering in masked diffusions. In Forty- International Conference on Machine Learning, 2025. 1, 6, A.2, B.2.2
Xuanlin Li, Brandon Trabucco, Dong Huk Park, Michael Luo, Sheng Shen, Trevor Darrell, and Yang Gao. Discovering non-monotonic autoregressive orderings with variational inference. In International Conference Learning Representations, 2021. 3.2,
Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmucli, al. Evolutionary-scale prediction of atomic-level protein structure with a language model. Science, 379(6637):1123-1 2023. 4.1, B.1
Sulin Liu, Juno Nam, Andrew Campbell, Hannes Stark, Yilun Xu, Tommi Jaakkola, and Rafael Think while you generate: Discrete diffusion with planned denoising, The Thirteenth International Conference Learning Representations, 2025. 5
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. In Proceedings of the 41st International Conference Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 32819-32844 PMLR, 21-27 Jul 2024. 1, 3, A.3, A.3.2
Shen Nie, Fengqi Zhu, Chao Du Tianyu Pang, Qian Liu, Guangtao Zeng, Min Lin, and Chongxuan Li. Scaling up masked diffusion models text. In The Thirteenth International Conference Learning Representations, 2025a. 1, 3.2,
Shen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, JUN ZHOU, Yankai Lin, Ji-Rong Wen, and Chongxuan Li. Large language diffusion models. In The Thirty-ninth Annual Conference Neural Information Processing Systems, 2025b. 6, A.2, A.5
Erik Nijkamp, Jeffrey A Ruffolo, Eli N Weinstein, Nikhil Naik, and Ali Madani. Progen2: exploring the boundaries of protein language models. Cell systems, 14(1 1):968 978, 2023. 4.1, B.I
Jingyang Ou, Shen Nie, Kaiwen Xue, Fengqi Zhu, Jiacheng Sun, Zhenguo Li, and Chongxuan Li. Your absorbing discrete diffusion secretly models the conditional distributions of clean data. In The Thirteenth International Conference on Learning Representations, 2025. 2.1,
Fred Zhangzhi Peng. Zachary Bezemek, Sawan Patel, Jarrid Rector-Brooks Sherwood Yao, Avishek Joey Bose, Alexander Tong, and Pranam Chatterjee. Path planning for masked dif- fusion model sampling. arXiv preprint arXiv:2502.0333 2025a. 1, 2.1, 4.2, 5, A.2. A.4, A.4.2, 1.2. B.2.2
Fred Zhangzhi Peng, Shuibai Zhang, Alex Tong, and contributors. Open-dllm: Open diffusion large language models. https: sttps://gitthn //tpthubbbbb/////// com/ comm/pengzhanggah//pee OLLM//Lnnnn.LLMMmn..mmm 2025b. 4.3, B.3.1
13

---

Krishna Pillutla, Swabha Swayamdipta, Rowan Zellers, John Thickstun, Sean Welleck, Yejin Choi, and Zaid Harchaoui. Mauve: measuring the gap between neural text and human text using divergence frontiers. In Proceedings of the 35th International Conference Neural Informa- fion Processing Systems, NIPS '21, Red Hook, NY, USA, 2021. Curran Associates Inc. ISBN 9781713845393. 4.2, B.2.1
Yinuo Ren, Haoxuan Chen, Grant M. Rotskoff, and Lexing Ying. How discrete and continuous diffusion meet: Comprehensive analysis of discrete diffusion models via stochastic integral framework. The Thirteenth International Conference on Learning Representations, 2025. 1.3.1, A.3.3, A.3.3
Hitesh Sagtani, Rishabh Mehrotra, and Beyang Liu. Improving fim code completions via context & curriculum based learning. In Proceedings of the Eighteenth ACM International Conference on Web Search and Data Mining, WSDM '25, pp. 801-810, New York, NY, USA, 2025. Association for Computing Machinery. ISBN 979840071 0008440071199..... 3293. B.3.1
Subham Sekhar Sahoo, Marianne Arriola, Aaron Gokaslan, Edgar Mariano Marroquin, Alexander M Rush, Yair Schiff. Justin T Chiu, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024. 1, 2.1, 5, A.3.2, B.2.1, B.2.1
Yair Schiff, Subham Sekhar Sahoo, Hao Phung, Guanghan Wang, Sam Boshar, Hugo Dalla-torre, Bernardo P de Almeida, Alexander M Rush, Thomas PIERROT, and Volodymyr Kuleshov. Simple guidance mechanisms for discrete diffusion models. In The Thirteenth International Conference Learming Representations, 2025.
Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, and Michalis Titsias. Simplified and generalized masked diffusion for discrete data. Advances in neural information processing systems, 37: 103131-103167, 2024. 2,2,1, 5, A.3.2
Andy Shih, Dorsa Sadigh, and Stefano Ermon. Training and inference on any-order autoregressive models the right way. Advances in Neural Information Processing Systems, 11-2762-2775.. 2022. 3.2, 5
Yuxuan Song, Zheng Zhang, Cheng Luo, Pengyang Gao, Fan Xia, Hao Luo, Zheng Li, Yuehang Yang, Hongli Yu, Xingwei Qu, et al. Seed diffusion: A large-scale diffusion language model with high-speed inference. arXiv preprint arXiv:2508.00 02193, 2025. 1
Haoran Sun, Lijun Yu, Bo Dai, Dale Schuurmans, and Hanjun Dai. Score-based continuous-time discrete diffusion models. In The Eleventh International Conference on Learning Representations, 2023. 3, A.3, A.3.2
Benigno Uria, lain Murray, and Hugo Larochelle. A deep and tractable density estimator. In Proceedings of the 31th International Conference on Machine Learning, 2014. 2.1
Guanghan Wang, Yair Schiff, Subham Sekhar Sahoo, and Volodymyr Kuleshov. Remasking discrete diffusion models with inference-time scaling. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025a. A.2, B.2.1
Xinyou Wang, Zaixiang Zheng, Fei Yc, Dongyu Xue, Shujian Huang, and Quanquan Gu. Diffusion language models are versatile protein learners. In Proceedings of the 41st International Conference Machine Learning, ICML 22,, 2024. 4.1, B.1
Xinyou Wang, Zaixiang Zheng, Fei YE, Dongyu Xue, Shujian Huang, and Quanquan Gu. DPLM-2: A multimodal diffiusion protein language model. In The Thirteenth International Conference on Learning Representations, 2025b.
Zhe Wang, Jiaxin Shi, Nicolas Heess, Arthur Gretton, and Michalis Titsias. Learning-order autore- gressive models with application to molecular graph generation. In Forty-second International Conference on Machine Learning, 2025c. 5
Ronald J Williams. Simple statistical
following algorithms for connectionist reinforcement
learning. Machine learning, 8(3):229-25 1992. 5
14

---

Chengyue Wu, Hao Zhang, Shuchen Xue, Zhijian Liu, Shizhe Diao, Ligeng Zhu, Ping Luo, Song Han, and Enze Xie. Fast-dllm: Training- free acceleration of diffusion Ilm by enabling kv cache and parallel decoding. arXiv preprint arXtv:2505.22611, 2025. 6, A.2,
G. George Yin and Qing Zhang. Continuous-Time Markov Chains and Applications, volume 37 of Stochastic Modelling and Applied Probability. Springer, New York, NY, 2013. A.3.1
Yang Zhang and Jeffrey Skolnick. Scoring function for automated assessment of protein structure template quality. Proteins: Structure, Function, and Bioinformatics, 57(4):702-714 2004. doi: B.1
Yixiu Zhao, Jiaxin Shi, Feng Chen, Shaul Druckmann, Lester Mackey, and Scott Linderman. Informed correctors for discrete diffusion models. In The Thirty-ninth Annual Conference on Neural
Information Processing Systems, 2025. 5
Kaiwen Zheng, Yongxin Chen, Hanzi Mao, Ming-Yo Yu Liu, Jun Zhu, and Qinsheng Zhang. Masked diffusion models are secretly time-agnostic masked models and exploit inaccurate categorical sampling. In The Thirteenth International Conference on Learning Representations, 2025. 2.1
Lin Zheng, Jianbo Yuan, Lei Yu, and Lingpeng Kong. reparameterized discrete diffusion model for text generation. arXiv preprint arXiv:2302.05777, 2023. 2.1, 5, 6, A.2, A.5
15

---

APPENDICES
Theoretical Derivations and Proofs
17
A.I Foundational Derivations
17
A.2 Instantiations
25
A.3 Altemative Proof of Proposition 3.2: Continuous Time Markov Chains Perspective
29
A.4 Generalization to Planners with Remasking (P2-style). .....y..........................
34
A.5
Other Sampling Methods with Multiple Denoising Positions and/or remasking
38
Experimental Details
38
B.1
Protein Sequence Generation
38
B.2 Text Generation
40
B.3 Code Generation
41
C Additional Results
42
C.1 Unstable Training with Pure PAPL Loss
42
C.2 Comparison of Training Curves with Vanilla MDM Loss
43
C.3 Empirical Estimation of the Effect of the Approximation Steps
................................
43
C,4 HumanEval Performance Analysis
44
Practitioner's Guide
46
16

---

A THEORETICAL DERIVATIONS AND PROOFS
A.1 FOUNDATIONAL DERIVATIONS
A.1.1 PROPERTIES OF KL DIVERGENCE
Recall for p. E AlR for X' some finite set, we define
p(æ) g(x)
DeL(Plla) = Σp(ar p(x) log
DEA
when p(x) 0 for every X' such that ((x) and to0 otherwise. Also recall that for such that p(æ) = 0, we interpret O log(O) = O. Here we will recall some basic properties of DEL that will aid in our proof.
Lemma A.1. Non-negativity of KL-Divergence: For any p. distributions a finite X,
DkL(pllg) 0,
with (pl(p) 0 if and only i q-
Proof. First observe that that g: 10, oc)
IR given by g(l) log(()) - -1+1 has g() log(!).
g((1) < 0 for E (0, 1) and g() > 0 for > 1, so y has global minimum of at 1 1. Then:
p(x)
Dιr(plla) = Σ p(x) log
((x))
ZEA
p(x)
p(x:)
Σ
(xx)
log
+1-1
(xx)
((x))
ZEA
p(x)
p(x:)
Σ
((x)
log
+ Σ ((x)
Σ p(z)
(xx)
((x)
FEX
xEX
ZEX
Σ
((x)
p(x)
log
(((:)
p(x)
+
((x)
<(x)
((x)
TE
p(z)
Σ
(())g
q(x)
ZEX
((x) 0,
xEX
with equality holding if and only q(x)= = p(x), Yx X.
Lemma A.2. Chain Rule for KL Divergence between Joint Law of Discrete Random Variables:
For p, q distributions on X Y, where X and y are finite sets, denoting by px, qx the marginals p(a.9)) and similarly for (y)xx)) of p and respectively and by Py) px(=]
DKL(pllg) = (pxllxxx + Exppx
Proof. By definition:
p(x,
(xı(plla)= = Σ
p(a,v) log
((r,v)
(x.))xxxy
(P)x(()())xx())
Σ
p(z, log
()1x((1z))xxx))
(xa))xxy
(xx(x)
Px)x(()2))
Σ
p(z,v) log
+
Σ
p( log
(x(x))
x (v(z))
(x.))XX
(x.))xx
17

---

= pxx(x) log (xx(a) gx(xx)
(2)x((12)
Σ px(x)
E ())((()))) log
+
4y1a ())x((/ax))
ZEN
nEX
DEY
DRi(pxllax) + Σ
DEN
DRI(pxllax) + Eurox (Dre(o)x((
Corollary A.3. Chain Rule for KL Divergence between Joint Law of N Discrete Random Variables:
For N E N and p,9 distributions No 21 .N.XN where Xo...., Ax are finite sets, denot- ing for E {0, 1, [,...,N] pask the No x ... x N3 marginal of and similarly for 4ozks and by Pk+ 10okk (#k+1 Io....,,2 mk) (e=s(epooMaddle)) for Is Nx 0,1 0.1..... fs and similarly for Pe (K
9k+1| we have:
DKI(Plla) = (Ru(polleo)
E Re (x) (Pe410010000 ................................ , (1)0000 (..
Proof. This follows from iteratively applying Lemma A.2.
Corollary A... Inequality for Marginalization of Discrete Distributions:
For p. distributions on X y, where X and y are finite sets, denoting by px, qx the marginals of and respectively:
D 2 DKI(Pxllqx
Proof. This follows from Lemma A.2 via noting that the term inside the expectation is non-negative for each via Lemma A.I.
A.1.2 DISCRETE TIME MARKOV CHAINS
Definition A.5. Discrete Time Markov Chains: A discrete time Markov chain on finite state space S a sequence of random variables (Xa JAE such that for any & E N and zo, T1, , It V, E S:
P(X; x|X = X FA
X1 X, v.) = P(X: = x 1 y).
(8)
The distribution of path of length (Xo, XI...., X₂) E Sk+1 of Markov chain {X*) at any time is entirely determined by its one step transition probabilities, which encode into its transition matrix:
Qx (z. :) = P = =) v). 2. Z.yEV.EN.
In the case where P(X; x x| y) transition matrix does not depend on the time k. we say the chain is time- nomo-homgeeneuss.
P(X = x|Xo = y) for all E N, ie. when the
Proposition A.6. KL Divergence Between Paths of Discrete Time Markov Chains:
R, P be probability distributions on SN+1 corresponding to the distribution of paths of length N of two discrete time Markov chains Y and X on S with transition matrices R and Q respectively. Also suppose that Yo ~ p and Xo ~ V for some μ, v E S1, Then:
N-1
Ru xx) log (R((y.
DRI(R[IP) = DKL(H(PP) +
Σ
E
Exg
Σ DES
where Px E A is given by:
rx(2)) PPYk x))
18

---

Proof. By Corollary A.3 (using the same notation as in the statement thereof):
DKI(RI()) = D
N-1
xx)] -100 [x))
(KL(H))
N-1
E ( [PKL(P(YR+I = + 1 = xx)i
by definition of v and the Markov property equation
N
= -Ri((ulpp)
Σ
Bawwrs [DRI(P((E++ =I = =2a)|P((x+11= = 1Xx =xx)]]
k=0
by definition of rhs and Ro.k
N-1
Rx, v.c)
= -KI(H((p)) +
Bag mork
Σ
Ru (v, my log
Qu
k=0
yES
by definition of the transition matrix equation 18.
A.1.3 APPLICATION TO THE ELBO
The loss corresponding to an ELBO 8 (xo) (i.e. a quantity satisfying log po (x0) 2 (x))) Vxo E VL and py the generated data distribution) always given by
[(0) Exa [E0 (xo)]
so that
Pdata (x)
DKL (psata (pp)
Σ
Pdota (xx) log
pa po(x)
Pdata (x) log Pdaco (x)
Palata (x) log(po(x))
XESL.
XESL
H (pdata) Exor-paara (og((o(xx)))
< H (pdata) Exor-paana
[x0(x))
H1 (pdata) + L(0).
Here (p) denotes the Shannon entropy of p. Note that, crucially, Pe must be the distribution on y which samples from at inference time, since this is what one wishes to make equal to Pdato via minimizing C(A) to H (pdata)).
In the following proposition. we show Proposition A.6 and the basic properties of KL divergence from Subsection A.1.1 can be used to derive an ELBO, and hence training loss, for any sampling procedure which can be described via discrete time Markov chain.
Proposition A.7. Application ELBO:
Suppose is a distribution on S the given by p(æ) = PPX = for some N E N and X a Markov chain on S with transition matrix Q. Then for xo S and Y any Markov chain with rate matrix R(: 2o) satisfying both
1. Y is equal in distribution to Xo
2. PP zo)
19

---

have:
N-1
R& (u,
log(p(zo)) 2
Σ VEES
Rs Re log
(y.
k=
where rk(:; (0) E 115 is given by:
(1 = P(Y =x).
Proof. First we observe that:
log(p(2o)) - (((x))l)) -DKL(P(YX' = PI =.))
by definition. Then, applying Corollary A.4 to R(;; P the distributions on S corresponding to paths of length N of y and X respectively, using X = S and y = SN:
- -DRi, = IIPNNN = =.) 2 (R(
Finally, by Proposition A.6:
- (Dkx(R(:;20))) - DKL = - -))
N-1
Ra (x
E (na
Ra (v, log
Qx(y,
k=
yES
N
(u((u Qu(y,
E Emgora((2))
Ru (15,, (U log WES
Σ
R=0
where in the last step we used Lemma A.I and that IP(Y, = =) P(X =) by assumption.
A.1.4 PROOF OF PROPOSITION 3.2: SIMPLE, DISCRETE TIME PERSPECTIVE
Here we provide a novel, self-contained proof of Proposition 3.2. particular, this encapsulates the proof of the standard DLM ELBO equation 3 by setting Cat (1; G2 ()) = Nwr(x] for all z and
i such that æ = m. This proof methodology helps elucidate the freedom of choice of reference dynamics, and does not require any of the prerequisite knowledge on continuous time Markov chains as other existing proofs in the discrete diffusion literature.
Ge denote the distribution of x1 obtained via the planner- Proposition 3.2. For any planner Go let Pg guided sampling scheme. Then we have the following ELBO:
log(po (xo)) 2 (xo) = & (xxo + (xo),
L
Ef = L kwear(u:L-11)
E
E
Σ
Carli; Co xx)) log (Cat(ab: Dā (xsx)))
Go Xw xurur
(-000)
Lim 1,
cg.0(x)) -L
E knn (0: 1])
Σ
Carli; Ga(Xa, xx)) log
Catli; Go xx.) Fx, 26. i)
(-18a) 111,f
Go is the distribution at time k of a Markov chain with initial data (m....,m m) and transition where Ts
rates as in equation but with 1/NM (x) replaced by Go (xo, x).
Proof. By equation 4, for x E VL, Pg (x(() = P(x = x) where X G is the time homogeneous
discrete time Markov chain on yi. with transition matrix given for x, y E 1 by:
CCat (v': (v': Dx)x) Fo yf f), otherwise
duAM(x,y) = 1, x² + yd xt m
(y,
10,
Go.0 =x) Cat(x; ((m,);;((mm........ m))). and P(x
20

---

So to obtain an ELBO for Po using Proposition A.7, select any family of transition matrices ; parameterized by Xo Vd determining family Markov chains y*o such that
Go
P(YZ = X 1 = (m...., m)) =
(10)
There are infinitely many such choices for the reference dynamics Y*0, but in order to make the paths of the reference dynamics apriori as close to those of XG as possible, we opt to keep the planner Go in the transition probabilities and simply replace Dx(x) by 6(xn) in the rate matrix Q6.0 Recalling the definition of Fa, from equation 5, this yields R(-, Xo) to be the time homogeneous rate matrix R (5 (x) given by, for x, E V4:
Cat(i; Go x)))Catyy';
dx,M(x,y) = 1, o Fy m otherwise
Observe that this is the same transition matrix from equation 2 but with 1/ (x) replaced by Go (x0, Also observe that indeed equation 10 satisfied, since at each step we simply choose a coordinate among masked positions of Y with probability Cat(i; (x x)) and flip from m to zy--
Then, inserting these choices into Proposition A.7 and using
Cat(i; Go x)) Cat (y': (xxx) (x, o², log (Cat (v': Db(x)) - log
(x, Cat(i; Gg(Xo.
log
the proof of Proposition 3.2 is complete.
A.1.5 ORM OF THE OPTIMAL PLANNER FOR A FIXED DENOISER: PROOF OF PROPOSITION A.8
Recall that the loss associated to the ELBO from Proposition A.7 is:
C(0, ④) Exxo-aaa (0 (xo)]
(1I)
A natural question is: for a fixed (imperfect) denoiser, what is the optimal form of G. which minimizes this objective? We will show here:
Proposition A.8. Consider the loss [(Ø) = - (xo)) where (x)) is as in
Proposition A.7 and Do is fixed. Then (()) is uniquely minimized over Go when, for No. N2. E VL.
with Xg containing no masked tokens and X2 equal to xa all unmasked positions:
Go(xo, xx) x Lo (x/)xx))
(12)
for Lo.s the transition probabilities of the data generating discrete time Markov chain's 's dynamics from equation 4.
That is, our loss finds a planner which mutually consistent with the denoiser in that picks each step a coordinate to unmask with probability proportional to the probability of denoising coordinate zá at the current generation step under the planned path. In short: the optimal planner tends to assign high mass to trajectories whose sequence of single-coordinate updates yields high likelihood of producing the observed zig, which during training is supervised by the data. Observe that equation 12 fixed-point equation relating values of G (Xo, xx) itself and the imperfect denoiser through Fo (Xx, an, i) of equation 5, so one can not simply choose to use this optimal planner and indeed needs to train towards optimality in practice.
Proof. To see this, we first observe that for fixed æg. no constraint need be enforced in the relationship between the distributions Gz xx),, Gă Xx) A when (z, xx) o (z, E D2L, This means that minimizing equation 11 is equivalent to minimizing the integrand:
Cat(i; (Xp, xx))
Cat(i; Gx xx)) log
Cat(æt; D2 Fe.s (Xk, 26
21

---

DKu(GG(Xx,
xx)) + log (x
(13)
for fixed To, Tke, and k; where ro, Xo, xx) A is given by:
Cat(i; ( xx)) Car( (z ox =0, i) (0,0()))x)))
and where Co. its normalizing constant.
Next observe that, although Co. appears depend implicitly Ge Gx. xx) through
Fo. (xx, xy. i) (recall equation 5) this does not affect the minimization problem. To see this, observe that
(x,, i) III Cat(zz (xxxx)) (x x+))
374
+ Ee-Do(x)
1 12-1 Cat (i; G i. xx))
where for x V4,x Vt-1 x but with its i'th component removed.
Then
(x, xx.
Σ Cat(sb:
Cat(ag: (x(x)))
Σ Cat(i; (xo (xx:
Σ Cat(alj
[2-1426 Cat Go fndd xx),
II
Cat(ag: (x())))
j=1
Σ
Cat(abi ()
12-1420 .Ctt Go xx))
That is, Co. xx) only depends (z xx) for + xy. Hence, minimizing equation 13 over Go(Xn, (x) is equivalent to minimizing the KI. divergence between Go Xx) and ro,o X By Lemma A.1, this occurs precisely when equation 12 holds.
A.1.6 DERIVATION OF GENERAL PLANNED TRANSITION PROBABILITIES (EQ. 4)
Recall that the sampling methodology discussed in $3 is as per Alg. 2.
Algorithm 2 Gillespie Sampler with a Planner
I: Initialize: 4- 0, Xg + (m,..., m), planner Gø. denoiser Do
2: for ls = 0: L- do
3: Plan Sample Do(xx)
4: Sample dimension (2,
5: Denoise
Xk+1 X
7: wi+1 <=
8: end for
9:
return x
10:
Let Pa A d denote the distribution on 1d of sample obtained via running Alg. 2. and let Go
Go
( (1xx)) E Ad denote the distribution of x given X3- With abuse of notation, we will also Go
let Po (1 (x1))))) E Aaa denote the joint distribution of x and the l + 1'st sample Note that
22

---

Xk+1 may only differ from X₁ in single coordinate i such that 24. m. So, letting for V
E [1: y v,xxi.
1.0 E Vih be equal to [a", ...., ....,.""_ y. x 2t x..:
Pa (xx
Go
i.ll [xk) Σ
Pa
Ge (x,
by 2/xx)
ZEVL
Σ
II Cat (z(x)()))) Cat (i; Go(2,
1.
Cat (v: D2 (xxx)
Σ II Cat (z(xx)))) Cat (i; Gø (2 xx,))) xx)) BEVE
Cat (y: (xx)) Fo. (x..
Go where Fa.s as in equation 5. Po.k+1 (x* xx.) is precisely what denote Lb (yr xx. in
equation 4.
A.1.7 PROOF OF PROPOSITION 3.1 (GREEDY ANCESTRAL VIOLATES THE VANILLA ELBO)
Go (x, σ) denote the probability Continuing with the notation from the previous subsection, letting of generating the sample x along the path I',, we have:
Go Po,2 (x, 0) II k=1
Go Pa.e (xx (x0x)) (x((=)))
II Cat (k)))) Di Fo (x 2(0)) (kk)), k=1
where here we use the same notation as in equation Thus, we arrive at:
Po (x) = Σ
Po.
(x,,)
= Σ
II Cat (x
Fe, Fo (xe(()) (0))) (ookk)o(k))
(14)
GENT 1m
Using equation 14 and specializing G& to the case of greedy ancestral sampling, we readily obtain proof of Proposition 3.1.
Proposition 3.1. For ph greedy (xo) defined with Gø in equation and D9 an imperfect denoiser, we may have
log(rg greedy (xo)) < C.,unif (xo);
where (xx)) is as in equation 1.
Proof. Since we just need a counterexample to the ELBO property, we may restrict to the case of I= 2 and V = m. To construct an example denoiser in this setting, we only need to define 6
terms:
(m, m)), ci = (1, D3 (m,
22 Cat (1, (m, m)),
c3 =Cat (1, (m,
6 = Cat (1, [ (m.2)) 2)),
c Cat (1, (1, m)).
c =Cat (1, (2, m)).
Then:
Cat (2, Dg (m, m))
=1 Ct
Cat (2, D3 m, m))
=1 - =1 cs.
Cat (2, D/ (m, 1))
=1 C3: Cat (2, D2 m)) =1 Co.
Cat (2, (m, 2))
=1 C4.
Cat (2, D 1, m))
Note that imperfect denoisers need not inconsistent, meaning that there is no reason to enforce any relationship between c Cs
(0, 1).
23

---

Let's take for example (1, Then, from equation 1:
1
xx) = L!
log
log log(c1020355)
OEE ==1
To find prooods (xx, we use equation 14:
grooty (x) Σ I Cat D2 (x Fe, F
20000, o(())
MEXS 1=1
Fo. ((m, m), 1. F), ((1, ((1,m), 1,2) + c2c3 (m, m), Fo Fo ((m, 1,))
c + 0203d2.
di, d2 will be found as functions of the c's, and we will find c's such that (c1 + 02cyd2)2 C2C Taking log of both sides and dividing by 2, the inequality will be shown.
Inserting the definition Pa.s from equation and specific choice of Go from equation we have
Fo.0( 1.11.1
1) = E, (m CCtt (1, Gb((1,)) (m, mn)]] Cat (1, (o((1, 1), (m, m))) + (1 02 Cat (L,,,)), (m, (m)
= -2lla, + (1 02) 1 >1-6
and
Fo, 1,2) (1. (2, (( (1, = ca Cat (2, G ((1,)) (1, m))) + - cs)-s)Caa (32, (1, m)))
Cs (5) 1,
so di c22o0 + (1 - 02) Here 1 denotes the indicator function.
Similarly,
Fom m), 1,2) 2 (m.m) CCt (2, G5(z, (m, m))]]
= c Cat (2, ((1, (m, (), (m, + + aa (2, Go((2, 1), m)))
c 1l 11 (1-a))ll-sss
and Fe,o((m, 1,, = so -., + (1-
Taking any Cı, C2 such that c2 > сi and
сi + c2 we get di and d2 = Cj. Then (c c20zd22 e7 < 01020305
C20203 <
There are many choices here that work. For instance, C3 = 1/4, = = 1/2, as
C,02 = 1/32 < 1/2 = C5.
Note that this means there data distributions and denoisers for which
uniil (0) - Exgr-pem [ unil (xo] < expm [og Pe (xo)
greedy
(recall here the discussion in Subsection A.1.3), so
DPi(P Al greedy) (paata) Exo-pan
[log pg
preedy (xo)
> H (paara) + mmask (0).
This means that training to make mask (8) small cannot provide any guarantee that p greedy is close to Pdata-
24

---

A.2 INSTANTIATIONS
Our general P-ELBO recovers familiar training objectives when we plug in specific planners.
Uniform planner. If G. selects uniformly among masked tokens- that is, Cat(i; Gz(z, x)) 1/NM(x) for masked and O otherwise-then planner-based sampling reduces to vanilla ancestral sampling. In this case G does not depend on which makes E. Ex.0(x)) (xo) 0. Substituting into Prop. 3.2, we exactly recover the standard DLM ELBO equation
Greedy planner. If Go always selects the most confident position according to the denoiser, as in equation 6, then the sampling path becomes deterministic. The associated ELBO then as follows: Corollary A.9. Let Yo (m...., m), and define recursively for k 1,..., L:
i jk: otherwise.
jk arg max Cat(ah; D)))= greedy
YE
4: Y =m
For Pe the distribution xt x under greedy ancestral sampling,
log(re
greedy (xo))
gnlen (xo) L Unj
Σ
log Cat(z)i (3(T))
Y -
Compared to the uniform case in equation 1, the greedy ELBO only accumulates logits along the greedy path defined by the denoiser. This highlights the mismatch: the standard DLM objective trains uniformly random paths, but greedy inference relies single deterministic path.
Soft greedy planner. We use the soft greedy planner:
Cat (j: G2,z x)) exp (Hlog (Ca(z); /C,
(15)
Σ exp (Hlog (Cat(z'; (Cat(z);
(² =
i=1,0'=m
as a regularized approximation to the greedy planner 6 in order to motivate, after performing the series of modifications discussed $3.4, the PAPL training algorithm 1. The ELBO associated to this choice of planner is:
Corollary A.10. For pá the distribution of x1 resulting from the planned sampling Algorithm of $3 with G = T as in equation 15, we have:
log ( (xo)) g8, (xo) + cg. (xo),
00.0.T (xo) eQ.0.T(x0)
E 10
F xx (imo)
Carli; G3 xx)) log (Car(ab: DJ (((()))))
i=1,
E.
R
E
Σ
Carli;
kn. 1.- 1]) Mx (:x00
pa (x<)
1=1 =m
C (x), C² (itb,, xx)
x log
where here we recall the notation z. iiab means the i'th coordinate of z is replaced by the i'th coordinate xo. and rī. (xx IPY for 8x the discrete time Markov chain with rate matrix equation 16
xo.
Carli;; Carl
dHAM(x, y) = 1, # 4. x² m (16)
Rd (y, Xn)
00
otherwise
and yx (m.... (m..... m).
25

---

We remark that while this is simply used approximation to greedy ancestral sampling for the purposes of this manuscript, soft greedy sampling is also used in practice in, e.g. Wang et al. (2025a)'s "Confidence Based Schedule." so this result is of independent interest as cotrected ELBO to these sampling schemes.
Other unmasking schemes. We remark there are other unmasking schemes in the literature for which one obtains an ELBO via our Proposition 3.2. For example, to obtain an ELBO for the "top probability margin" method of Kim al. (2025), one inserts the choice
G2(2xx) = arg max Cat(y; D2 Cat(ū: Ca(()))((x))1] D)(x))))
tarf=m
where y arg jax Cat(j; Cat(j:Dz(x) (x) and y arg max j Cat(j; (x)) As the focus of this work is obtaining a viable objective for use with greedy ancestral sampling, we do not provide expanded details on how train for this planner user our ELBO.
Extensions to remasking and denoising multiple positions simultaneously. So far we assumed that once unmasked, token remains fixed. practice, planners such RDM (Zheng et al., 2023) P2 (Peng et al., 2025a) allow remasking and resampling, in addition to denoising multiple tokens simultaneously. There are also methods which attempt to denoise multiple tokens simultaneously. but do not allow remasking, such top-k block denoising Nie et al. (2025b) and confidence thresholding Wu al. (2025). Our proof technique extends naturally to these cases, yielding planner-aware ELBOS of the same form as Prop. 3.2. For completeness, in $A.4 planners and show its specialization to P2-TopK, in addition to discussion how the generalized version of the ELBO could be used for finding training stragies for these other sampling methods.
provide generalization P2-style
A.2.1 PROOF OF COROLLARY A.9 (ELBO FOR GREEDY PLANNER)
Specializing the ELBO from Proposition 3.2 to the case of greedy sampling, we set Go to be in equation 6.
Proof. We first observe that inserting the choice of G& from equation 6 into equation 2 in the place of 1 ar (x), yx becomes deterministie, with dynamics Yo (m...., m), and
jk-1 == 1=,...,,44
otherwise
jk argmax s Catiz D)(Y))). k= O......
For ex.((xo), we have by definition Cat(i; (x xx)) = Cat(i; 8(jx)) so:
(x,((x))
Σ log (Cat(a)); D (va))
k=0
Similarly, for the term El (xo), we have, recalling the definition of Fø.o from equation
L-1
50.0 (xxo = Σ
log
(Y,, ..k)) (p))))
k=0
L-1
L II
log
Σ
Cat(2"; Cat(2);D((Yo))
<c Dit
log
II
Cat(ri; Dí D2 by definition of js
k=
V (E)1,)).
log (Cat(a); ()+)))))
k=
Summing this expression Ex: with this lower bound on c6.0 we have the result of Corollary A.9.
26

---

A.2.2 PROOF OF COROLLARY A.10 (ELBO FOR SOFTMAX PLANNER)
We now specialize the ELBO found Proposition 3.2 smooth approximation of the greedy ancestral planner from equation 6 - namely, we take G = G as in equation 15.
Proof. Ex Ef.00 is simply inserting G G into g from Proposition 3.2.
Now we make lower bound on ef (xo). With this choice Gg:
4-1
Cat(i; (x0, xR)) F (x* Ty i)
(x0) =
Σ
Exx
(1xx))
Cat(i; (x xx)) log
k=0
where Fo,s is as in equation 5 with G = Φ. So, by Jensen's 's inequality:
(xo)
L Σ
Ex Ex (1x0)
Cat(i; (x,,, xx)) log (Cat(i; (Xu: xxs)))
1=Lai=m
Cat(i; G (x0, xx)) log (F7 (Xi, ud 21]]
Bx (x
1,2g
(x
Σ
Cat(i; GT log (Cat(i; (2 xxt)])
Σ MMI
E 3x
Σ
Cat(i; xx) [log (Ctt (; G22(2 i,a Xx
i=
Ex (:200
EaxDa(
Σ
Cat(i; (x,, xx))x
1,
Cat 1; GT t, 1, Xx
log
Cat(i; G7 (xo, xx))
C (x),
BawDa(sa)
Σ
Cat(i; (x0, log
C
1, (i xx)
i=1,1==m
where in the last step we use that for any z, Cat (1; and Cat(i; (x xx)) have the same numerator in equation 15, just different normalizing constants. This lower bound is denoted as 00,0.5 in Corollary A.10.
Eg
A.2.3 CONNECTION BETWEEN COROLLARY 10 AND THE PAPL LOSS EQUATION 7
Here show how one formally arrives at the PAPL loss from the detach gradient and stabilization steps taken in $3.4. We begin with the loss corresponding to the ELBO from Corollary A.10. This is given by:
C(0,6) -ExorPaa EBB (x00)((x0==0))) + ER.0. (xo)
where g8.0.7 00.4. are in Corollary A.10. .E2
Next, we detach logits from the softmax weights G given by equation 15. Observing that E 0.r depends only on these weights (through C") and not on the logits from the denoiser, have
27

---

minimizing L( is equivalent to minimizing:
[())
Exx
(xo))
E
E XL (2x)
Cat(i; (Xp,
xx)) xx)) log (Cat(z)) (x
Pran
k=
After this, we replace sampling X r (; Xo) with sampling X ( ::; Xn) with r3 as in equation Indeed, this was the reason for using the softmax approximation of Corollary A. rather than the greedy ELBO of Corollary A,9 in the first place- which the deterministic paths from F may be very far from the uniformly random paths of least we have rí becomes:
The loss
L-1
((0) =
Σ
E XampPin
E (-ica)
wii log (Cat(ro: Dā (
Xe
k=0
1=1,
where w Cat(ē; CJ (x0, xx)) oc exp identical to the vanilla loss
log (((()))(((x))) Finally, we observe that this is
Cunit (0) = Exo
[ (xo)]
associated the vanilla ELBO equation 1, except that T has been replaced by whi as the weight in the sum. Thus, interpolating with constant which decreases linearly with the number of samples yields:
CPAPL(O) = -E -Exo
[ uunif (xo)]
L
E Diau
E
Σ
a
whe log (Cat(ah; D2 D((x))))
xumora (-200
L
k
0=
iim I,m
E (:x0)
Σ
L- k (1 ow!t) ) log (Cat(z)); D (xx))))
A==0
Suppressing the distributions of the random variables xo, k X2 in the notation, this is precisely equation
A.2.4
COMPARING RELATIVE SIZE THE APPROXIMATE LOSSES
Here we will see how the vanilla DLM loss (recalling here the discussion in A.1.3 and equation 3):
unif (
Exxo-paona
Unni (xx)
LLEEx Paata
Et (aa-Un(((:::L1110 (0:
E Exerra
Σ
log (Cat (2): (xx)))
L L- h
1=1.2f pm
compares with the surrogate corrected loss:
C (0) - -IE panna
BA-Umi((U:L-1)) Ex (xxwra (:20))
Cat(i; G xx))
x log (Cat(z); (iss)))
which is the PAPL loss before interpolation with the vanilla MDM loss as per the previous subsection. Here recall rx(::))) from equation and C", Gt from equation 15. Proposition A.11. For any Ti 72 0, Lumi((0) (0) > LT (0) (0)
28

---

Proof. As the expected values are over the same distributions, it suffices prove the result for the integrands. Let X6, X₁: E Vi. and M = E (l...., L} : x² = m) Note |M| = L by definition. Define:
= log (Cat (r/: D)(xx))) MMM
so that
Cat(ù; G% (xo, xx)) = exp(e wt
(² - E
exp(((/1)
iEM
Noting the minus sign in front of the losses, we simply need to establish that
Σ
Σ wh
(e)e S Σ w (E)e.
S
1EM
iEM
SEM
Observing that lim, w (f) - vi M and we simply show that
d
Eet <0,0>>
dr
REM
Letting F(r) EiEm (()e), We have
dr [({
d
(e)
Σ
wt(0)e -
(1( [F(T) e
72
72
JEM
SO
F(T)
Σ
w (4) (4)
[F(T) - e- 1 72 -
Σ [u((00933
dr
72
iEM
iEM
By Jensen's inequality, (F(T))² S EiEM wi ()) so we are done.
A.3 ALTERNATIVE PROOF OF PROPOSITION 3.2: CONTINUOUS TIME MARKOV CHAINS PERSPECTIVE
Here, for reference, show how Proposition 3.2 be derived from the continuous time Markov chains perspective taken in the discrete diffusion literature (Camphell et al., 2022; 2024; Lou et al., 2024: Sun al., 2023).
A.3.1 IIMO-ENHOMOGS CONTINUOUS TIME MARKOV CHAINS (CTMC)
A (time-inhomoge (nimmoinhomogeeeuus) continuous-time Markov chain {X,) finite set X' is stochastic process satisfying the Markov property, which can be formally summarized as P(X, y X
Ii,..., Xs, Tk, X, = x) (P, y Xx x), Vy, .i,.,,,,, T X 0 < 81 82 <
8k <<<<tt 1. One construct such process by specifying "rate matrix" Q. E RR x with Q((,.)) > and Qr x)) = Enst.r Q:(v, for all x + y E X and 2 0. Along with an initial
distribution μ A|x|, Q determines the 1-dimensional time marginals P(X, -) E A [X] via the Kolmogorov equation:
P((x =.P==
120
(17)
P(Xo =x)= I X.
When the above holds, will say Q "generates" X. Note that necessarily that if Q generates X.
Q+(y,)) = lim -
d
P(X, x/ =), E
(18)
afe ds
29

---

Knowing the entries of also provides means of generating samples from X₃ at any given time, since paths of {Xi can be realized via a sequence of jump times (TaINEN, with Tj = inf(t > T--1 X # X and the effective discrete-time jump process (X hie. Then
Qr(y, x) Q((x, x)
P(X y 7/11 = -,T =
(19)
and
(P > t|X,, = 2, = =)) S. Op(z,
(20)
For more background on time inhoomooooog continuous- time Markov chains, see e.g. Chapter 2 of Yin & Zhang (2013) or the appendix of Ren et al. (2025).
A.3.2 DLMS IN THE CTMC FRAMEWORK
In the original CTMC framework for DLMs (Lou et al., 2024; Shi al., 2024; Sahoo et al., 2024), one begins with a coordinate-wer forward corruption process:
p=(=//=a)) (aa( + (1 - (x))(m))
(21)
for a [0,1] [0, 1] a differentiable, monotone-decreasing function with ag 1 and ai 1. Using equation equation 17, one sees that noising each coordinate independently according to corresponds
to CTMC X, with state space VL, intial data Xo and rate matrix given by, for x, E V
(t),
dHAM(x,)) 1, # v x = y otherwise
Q (y,x)
(t)
0g
where (t) de log(ae).
One then uses classic tive-reversal formula (see, e.g. Sun et al. (2023) Proposition 3.2.) to obtain rate matrix generating X, so that X Xo
<X
= x|xx = xo) V4, This rate
matrix is given by, for x, y E VL:
(Cat((y'0(x1)))
duAM(x, = 1, * 3g x²
Q. xy,x)
(x),,
=
(22)
0,
otherwise
where
don-t
Br
d8
(23)
0:1
Letting Ts for l: N be the time of X, h'th jump, we have by equation 20:
log|P(TT+1 > t/To s. X, =x) = -NN((x) ( B
and by equation 19, for x y:
yl Xx
P(X
= = Tk+1 =t)
(x), dHAM(X,)) = 1, + 1 x¹ m, 3²
xg
(24)
0.
otherwise
That is, one waits for exponential clock to ring with the given speed, then regardless of how long took, chooses uniformly at random between masked positions of x to get some index i, and unmasks that token to 2
30

---

One then seeks to denoise from (m..... m) to Xu ~ Pdata using the CTMC X, O,umii yi which one obtains via replacing (())) in equation 22 with neural denoiser (x(x). 8.uiff
with state space
A rate matrix
generating X,
is given by, for x,y VL:
B:Catty';: D2(x))), (x),
dIAM((X)) = =1, # FV²
2 m.makk (y.)) =
x =
(25)
0,
otherwise
0.mmask
This means, letting 7 for k E N be the time of X,
's k jump, we have by equation 20:
- mank
= =x) = -N =-Nx ' B,dr
log P( > tlT = 8, X
(26)
and by equation 19, for # y:
0.mnkk
+- 0,musk
P(X
=y(X,, =x, Tk+1 t)
Cat(y'; (j(()',)((x)))NN((
dHAM(x,y) 1, # V = -m, 3" + m
(27)
0,
otherwise
That is, one waits for exponential clock with the same given speed for X to ring then regardless of how long it took, chooses uniformly at random between masked positions of x to get some index i, and unmasks that token to y' with probability Cat(y': D2 (z=).
This is summarized succinetly via the corresponding Gillespie sampling scheme for a standard masked diffusion model, which, defining
M(x) (j (L. L) x m) V
(28)
is given by Alg. 3.
Algorithm 3 Gillespie Sampler for Masked Diffusion Models
I: Initialize: Xo (m, m m...., m), denoiser Dg
for = 0: do
3: Choose Random Coordinate for Unmasking:
Sample dimension Unif (M(x+))
5: Denoise:
Sample x² (xx)
7: X++1 + X
8:
#4+1 2"
9: end for
10:
return IL
11:
A.3.3 SETUP AND ELBO IN THE CTMC FRAMEWORK
Now we show how to derive a corrected ELBO for the dynamics described by Alg. 2 as one would using the CTMC framework. First, we observe that for any x'x a CTMC on time interval [0, 1] and state space S such that P(P:x = Xo) = 1 and a distribution on S given by p(x) = P(X, x) for another CTMC on time interval [o, 1] and state space S: log(p(x.))
D (6 1p) 2 DKL (())
(29)
where we let Rx denote the distribution of ý xo (on the Skorokhod space D( 1]: S) of all cádlág paths from [0, 1] 1d and Q the same but for X, and to get the bound, use the data-procsssing generalization of Corollary A.4- see, e.g. Budhiraja & Dupuis
inequality (an (2019) Lemma 2.4 (f))
31

---

That is, in order to make the terminal distribution p of X close 6xx)). one can simply require that its entire path is close that of yxo
The benefit of using the L divergence between the paths is that via Girsanov's Theorem for Markov Jump processes (see e.g. Theorem III.5.34 in Jacod & Shiryaev (2013) for general result or Ren et al. (2025) Theorems 3.3/3.5 for the specific Markov Chain setting) it yields a simple expression in terms of the rate matrices generating the dynamics of Y*x and X.
This result states that, for CTMC Y with rate matrix R, and Yo / and CTMC X with rate matrix Qr and Xo ~ y on the same state space S, denoting by IR the distribution of on D( 1): and Q similarly but for X. the equality:
( (2) DRI(Hlli)
(30)
+ . Bxcor
Rs Re(y, Qr Itt
(x( (r(
E VES.Ufr,
24) log
de
rx(x) = P(Y, - x), ES
holds under mild assumptions I and Q (see remark 3,4 in Ren et al. (2025)). Note that equation 30 is simply a continuous time extension of Proposition A.6. Recalling equation 20, the term R₂ (1x, 2.) Qx xx4) measures the difference in jump times between X and Y, while the second term is essentially a KL divergence between the transition rates.
We proceed by identifying CTMC X with state space V such that F (x(*-xx) PP
Go (x),
so that we may apply equation 29 and equation 30 to P9 Go obtain ELBO. Denoting by O₂ the rate matrix for 0 via equation 4 and equation 19, must have for any t E [0, 1] and x + y E V4:
Q (y,x) [Cat (v': (x, 2. x
dHAM(X, y) = 1, 1 + V =m otherwise
[0,
As the transition probabilities do not depend on the transition rates, we simply need that the transition rates so that by time all tokens will become unmasked. We thus simply maintain those from vanilla DLMs, found in equation 26, and so, recalling equation 20, we set for x e VL.
(x,*(x,)) - BaNNas
where we recall Br from equation 23.
We then have our full definition Q F, E V4:
Nv Cx) (x()) (o. (x),,
dHAM y) = 1, x #',""" m
(y,
=
0,
otherwise
(31)
Letting (x))) be the rate matrix for our reference chain You to be used for inserting into
equation 29, since we don't want to worry about enforcing the jump times of 8. and yxa in the
8,6 form of the ELBO (as these have no bearing the sample generated by X and hence should not be trained for), we also set for x E V!:
(x Xx x: Xb) = (x).
Now, to choose the off diagonal entries of Re(: (x)), we seek to modify jump locations to be different than in the vanilla setting, where the coordinate to flip is chosen uniformly random (see equation 24).
Instead, we opt to choose R(-, (x))) R xx)) to select coordinates to denoise according the planner. In this sense, we will be learning both the forward and reverse process simultancously when using this ELBOO.
32

---

Recalling equation 19, thus want for # y Vi: R xox)
Cat(i; Cx(xo x))Cat(y'; ((z)))
d 1, 2, t # v2, otherwise
(x, x;Xo) Xo)
(0
yGs and its rate matrix :();;x0) (., (x)) have been Now the dynamics of yxu determined.
We have for x, E V!:
3 (xN Gx(xxxx
d= 1, + f -m, x/
(y,
(x)
x =
0.
otherwise
(32)
Note that, taking Yo GG
x8 = (m, m). indeed y. = Xg. so that equation 29 applies.
We arrive at the following proposition:
Proposition A.12. For any planner Go, we have the following ELBO:
log(pg Go (xo)) Z E (xu) + E. E (xo),
where, letting Tt Go (:xx) be the distribution of Go with rate matrix equation 32 and Yo Go.Xo (m..... ..)
p..0(x0) = ((¹ Be Go
(xx)
Σ
Cat(i; Go(xa, x:)) log (Car(zé; D& (x(x))))
x
(1220)
i=1=)==
(x0)) = f(..
Car(i; Go(xo,xx))) ell
}
8.
E (-100)
Nxr(xe)
Carli; Go(xo, x.)) log
Fe, (Xe,
Proof. Let Ro denote the distribution of paths of G and Q those of X0.0 from the preceed- ing discussion. Then, by equation 29, log(pe Go (xo)) Dr (R (R**))((..)). From equation 30, have, using x. and X have the same initial data that the first KI term O:
= =-10" Xerrr ( (x,,((x, + & X1;
(C xxi xx;x)) (0x oft (y, xi)
]
+ Σ
RCo X,; log
y#xi
{-¹
I.
3
xenti x (:200)
(x(ex)
Σ
Cat(i;
f=1,/==
Cat(i; Go (Z( xx))
)
x log
Cat(a);; (x:))) Fo,s)) 26.
=(.0((x)) + (xxu)
This form is seen to be equivalent to that in Proposition 3.2 in the next subsection.
A.3.4 TIME INDEPENDENT FORMULATION OF THE ELBO
Now observe that, expected, the ELBO is independent of the time schedule Or (and hence BB). We start by observing that
PY
(y. Go.xo) =
(L) kx
exp (-f(
(-f Beds)) )) k (1 (1-exp exp (-f²
[-f[Adds B,ds)
4-N
33

---

One can this via law of competing exponentials solving equation 17 for the pure death chain representing the number of mask states.
Then, recalling that the transition probabilities for Y, G.a..o are independent of time and using equa-
tion 19, we have Y2 X LN (Y, ( L k equal in distribution Y Ge.Xo where G is the effective discrete time Markov chain with rate matrix equation 16.
Denoting by Pk Go (::x0) the distribution of y G. we have, for example:
L
( Be (oxp
(-fo (,ds)) (1-cxpp exp (- B,ds)
((dds))) L-
(xo) (x/)(
Σ
dt
k
k=1
x
kE
Cat(i; Go(Xo, xx)) log (Cat(zoi Dÿ (x()))))
Pr
(:x0)
Using exp(-Jo 3,ds) = exp(- de log(a1-a)ds) lim, exp(log(a(1 s)) log(a(0))))
lim a(1 5) 0 since a(1) a() = have
S 8: cxp (-So 8,ds) )) ((1 cxp (-fo 8ods)
1.-
dt ('1 un-1(1 u) u du
= Blk, L - k +1)
k(W))
where B is the beta function.
Thus, changing k to L k in the sum, have:
(x0)
Σ k=
E
Cat(i; Gp(xo, xx)) log (Cat(z); Dj (2
(1x00
Recalling that Pk is what is denoted as (we only added the bar here to distinguish from the distribution of the continuous time chain), see E8.0 from Proposition A. is equal c from 3.2.
Applying the same manipulations to ex,), we arrive at the following:
(xo), Proposition A.13, For E E. as in Proposition 3.2 and Eq.0 E. as in Proposition A.12, we have:
cQ.0 (xo) = (xo), E0.0 (xo) = Eg El.0 (xp), VXp E VL
That is, rather than simulate the CTMC up to random time t Unif(0, 1) to obtain sample of y, y.Go.Xo one may instead either sample an entire trajectory of the discrete time chain y G and accumulate losses, or sample random number of jumps k ~ Unif([0: L 1J) and trajectory of G. k
up to time k.
A.4 GENERALIZZTION TO PLANNERS WITH REMASKING (P2-STYL.E)
P2 (Peng et al., 2025a) allows for the remasking of clean tokens while still requiring that the number of unmasked tokens in X₁ k. We replace the planner G VL x y + A with a sequence of planners Gh DL VL (t) b,...,A, I,, where Gk 0. outputs a distribution on subsets of size & of [1 LJ. We then do:
In P2, we in practice use "P2-Topk",, which corresponds to:
Gk (2, x) = 5 (Top-KrElLL) Cat Go (z, (z,
xi)))
(33)
34

---

Algorithm 4 Gillespie Sampler with a P2-Planner
1: Initialize: X (m,...,,)) m), P2 planner 2 1i denoiser Dg
for = 0: -
3: Plan Sample x ~ Do(xx)
4: Sample dimensions 1 = (f), (x41)) ~ G (2,22,,1)) Xt)
5: Denoise
6: xk+1 + 2² for E 7h+1
7: Øk41 + m for & |++
8: end for
9:
return xL
where for i E (1,.... ,.),,, I.) and v) 12220
Cat(i; (2, x)) & ryCat(r"; ((m)) Db(x)) + Cat(zf;ō(m)) Cat(z'; Bx))) (34)
in the case of P2-BERT, and
Cat(i; (Z, x)) & r/Cat(z'; ((m)) D( + -1--Can Cat((=;((mm))) Cat(z'; (x(x)) (35)
the case of P2-self. Here the output of the denoiser unmasked positions x is no longer assumed to be 8(x²) when we write as D, an B& denotes an external BERT model. 17 a "stochasticity parameter" which controls the frequency of remasking increasing boosts G in masked positions, so that unmasked positions are more likely to fall outside of the Top-k and be remasked.
A.4.1 DERIVING THE TRANSITION PROBARILITIES
Here we derive the one step transition probabilities for Alg. 4. This is the analogue equation 4 in the setting where generalize to allow for remasking.
Let Pe Go.2 E AL denote the distribution on V of sample obtained via running Alg. 4, and let Go.2
Pa,6+1 (-1xs) E A denote the distribution of x given Xę. With abuse of notation, we will also Go.2
let P8.+11 (., (xx)) Ad denote the joint distribution X++1 and the k + I'st sample z. Note that
NM = -- - (k + 1) with probability 1. So for E y with NM L (k + 1), we let
C(y)= = (i E [1: L] y² f m). and have:
Go.2 Pa,k+1
(y|x)) = Σ
Pg Pg.k+1 z|xx)
.a.2
Σ z zEV:''Ey MEC
II Cat (2); Cat (Cly); G (z, (2,
II tECly)
Cat (b(x))
Σ II Cat (ot, (((ss)))))) Cat (C(x):: G (z
ZEVL
II Cat (y': (x(x)) pdot1 (xx,y) REC(y)
(36)
where
F 1l+1 0.0.2 (xk,)) y) = Exn Dx(x) Cat (Cly): G(t): (2
(37)
and for z,y E V4 we denote by z E VL the sequence which is the same as except with z' replaced by y² for all i Cly). Note by our assumption Do that this is 0 y' o at, in position where y ot # Also note that the proof the exact in Subsection A.1.6.
35

---

A.4.2 MARKOV CHAIN SETUP AND ELBO FOR P2 PLANNER
Now we observe how apply Proposition A.7 to get an ELBO for Pa Go.2 We will then specialize this to the case of P2-TopK, and discuss the difficulties in obtaining "regularized" approximation similar Corollary A.10 which lends itself computationally viable approximation as in Subsection 3.4.
By equation 36, we have Pg Go.2 -000((x (x) = P(X = x), where GG.2..0 the time homogenous
Markov chain on Voh with transition matrix given for x, by:
Qu (y, x)
II Cat (v's (x))) Fo Nusly) (x iECly)
(38)
when (y) (x) and y = Vi
and otherwise.
Once again, to obtain an ELBO for Pg Go.2 using Proposition A.7, we select any family of transition
matrices : parameterized by X E y determining a family Markov chains Y*0 such that equation 10 holds.
We choose, as in Subsection A.1.4, Y* &. to be the Markov chain with rate matrix obtained from replacing (x) with ((r))) in equation 38. Recalling the definition of Fo. from equation 37, this yields :(:,;x0) be R (t, xxx)) given by, for x. E V!:
(y, II Cat (vt:)(a)))) Cat (Cly): G/22 (.N-N(y)) (x,))((x,xx)) (x)
(39)
sEC(y)
when (y) (x) - and 0 otherwise.
Note that indeed equation 10 holds, since the k'th step Y always has I. k masks, with the positions of the masks determined by sampling from the planner at each step.
Applying now Proposition A.7 with this choice of R (y, x), we obtain:
Proposition A.14. For any P2-style collection of planners 22 k E LI. let Pg Go,² denote the distribution of X obtained via the iterative sampling scheme Alg. 4. Then we have the following
ELBO:
log(po
Co.2 Go.2 (xx)) 2 00,0.2 (xo) + E2
cl.ib.2
(xx);
1.-1
29.0.2
Σ
Σ
Car(C(y); Gk+1 (xo: xx))
Go.
A=
(:x YEX' 1(x0)
Σ
log (Cat(ahi D3 (xx+))
vEC(y)
1--1
2..0.2(
E
Car(C(y); Gha (Xo: xx))x
C
Lx
(:x00
-EXI
Car(C(y); Gh (xco xx:)
log F (Xk,
Go.2 (x; xo) P(Y x) for Y the discrete time Markov chain with rate matrix where Tk equation 39 and Y (m...., m), and here we recall the notation X₃ Xx(Xo) from equation I.
Specializing to P2 Top-k, where G is in equation 33 the jump positions of x are determin- istic and defined recursively via (suppressing the dependence on Xo in the notation):
1 Te Cat(i; 6o Yo 1)),)), &o 1.... L, do 0 Y i o Th.
(40)
[nn-
E I
m,
36

---

We have Go22 5(1x+1). So E//0022 becomes:
Eq.0.up Σ
log (Cat(æ); (Y
k= iET411
(xo) similarly becomes:
Eh (x0)
= Elog log (Y,, (Ya
k=
Recalling the definition of Fk+1 9.0,2
we have here that:
Fk+1 (Yx, Ys+1)
0,0.2
Σ
II Cat(s ((x)))
xEVi.ca(i@ge Yet1, (2)))))))))) C.)) (a
II Cat(rl; (Yx)), igT111
by definition of T4+1- Thus we have:
cP 1oook (xo) + (3.0.10p((xo))
Σ
log(Cat(z D2(Yx)))
Σ
log(Cat(zz); D)(H))
k=0
iETRI1
Σ
Σ log(Cat(< Di(K))
Σ
Σ
log(Cat(z); Dí
k= 1=1,
where in the last step we recall that we impose (V+)) ((2)) Y * m, and for any such that
Y * m, Y = xd- The resulting ELBO specializing Proposition A.14 to P2-Topk is just the same that for greedy ancestral in Corollary A.9, except that the paths are determined by the P2-style sequence of planners:
Corollary A.15. For Pg port the distribution of XL from Alg. 4 with Gs in equation 33, we have:
log(pg top (xx)) 2 ((.xop-((x)))
to = L.
Σ
Carlat: D2(Yxx))
where Xx satisfies the recursion equation 40.
However, regularizing the in the definition the P2-Topk planner equation to make have full support, and hence be better approximated by uniformly random style of sampling, quickly reveals that using such the ELBO becomes computationally prohibitive, even using e.g. the soft-max gumbel noise trick. Recall this is precisely what did for the greedy ancestral planner using the softmax approximation done Corollary 10. This additional overhead arises because for each randomly sampled time step k, would need not only simulate Y to time k, which requires function evaluations of the planner, but also to compute and sum over all of the (E5) weights
corresponding the (x")) values of which possible locations for the next jump in the reference dynamics in lower bound Proposition A. 14. For this reason, even though use P2-Topk sampling in our experiments, we assume the role of remasking is relatively minimal compared to the unmasking selection process, that the series of approximations inspired by Corollary A.10 outlined in $3.4 and detailed in $A.2.3 still results in a training objective more reflective the sampling process than the vanilla loss equation 1.
37

---

Also note that, in P2 (Peng et al., 2025a), the ELBO result was proved assuming continuous- time dependent, randomly sampled unmasking process, although the top-k sampling procedure Alg. with 2 in equation 33 was used in practice. This creates similar mismatch between training and sampling that is present vanilla DLMs vs. greedy ancestral see Proposition 3.1. The results shown here are for the practically used P2-Topk sampling procedure, hence there is little to no overlap in the analysis performed between the manuscripts.
A.5 OTHER SAMPLING METHODS WITH MULTIPLE DENOISING POSITIONS AND/OR REMASKING
We show here how one can view other sampling strategies in the literature special cases of Alg. 4 and hence train also for these strategies using the result of Proposition 1.1
RDM sampling Zheng et al. (2023). is captured by taking
(2,) = 5 (Top-kse(:zz) (Cat (st:D((x)))
where Dg is as in equation 35. This strategy uses the denoisers logits on both masked and unmasked positions, and only keeps tokens which fall in the top-k k threshold.
Top-j Block denoising Nie et al. (2025b). is captured by taking
(z, = (Top-iiEB(s) (Cat (x)))) u (i: x' + m))
Here previously unmasked tokens are kept, (x) is a block of masked positions of x of predeter- mined fixed size, and j is fixed desired number of positions be denoised at each step. Note that time here should instead only run to [L/j] rather than L - 1, since at this point all tokens will be clean.
Confidence Thresholding Wu et al. (2025).: In this strategy, the number of tokens to be clean at given time step is adaptive to how many tokens exceed a certain confidence threshold. This means that one would need to allow & in Alg. 4 to change adaptively based on and x. While we don't provide details here for the sake of brevity, one can make such a modification to the sampling algorithm and follow the general principle of writing down the transitions of the discrete time Markov chain for the generation dynamics in terms of a planner and comparing with paths which use a supervised planner using Proposition A.7 obtain rigorous ELBO in such regimes.
B EXPERIMENTAL DETAILS
B. PROTEIN SEQUENCE GENERATION
Setup. We evaluate our approach against leading protein sequence generation models. The com- parison includes three discrete diffusion models-DPI.M (Wang et al., 2024), EvoDiff (Alamdari et al., 2023), and ESM3 (Hayes et al., 2025)-alon with the autoregressive baseline ProGen2 (Ni- jkamp et al., 2023). Each model generates 100 sequences for target lengths in 200, 300...., 800. DPLM follows its standard setting, using sequence length tied to the number of sampling steps and temperature 0.9, with rejection resampling disabled to ensure fairness. ESM3 is sampled with temperature 1.0, cosine noise schedule, top-p 1, and 500 denoising steps. Special tokens stripped to guarantee valid amino acid outputs.
Evaluation. Generated sequences are assessed using structure prediction as a proxy for functional plausibility. Specifically, we fold each sequence with ESMFolld (Lin et al., 2023) and extract three structural metries:
pLDDT (predicted Local Distance Difference Test): an estimate of local per-residue confidence, defined as the expected accuracy of predicted interatomic distances. Formally, for residue i,
pred
Arue ij
(L 100 x E
d
JEN(()
38

---

where dıj denotes pairwise distances and (i) indexes local neighbors. The reported score is the average over residues.
pTM (predicted Template Modeling score): measures global structural similarity between predicted and true structures, adapted from the TM-score (Zhang & Skolnick, 2004). Given length L and alignment u(t) between residues,
I. Σ 1 (dinat))/o(L))
рTM
max alignments
where di)( the distance deviation and dL( Singth-dependemt scaling factor.
pAE (predicted Alignment Error): estimates the expected positional error in aligning residue of the predicted structure to residue j of the true structure. Formally,
PAE(i, IE
||
true 112
averaged across all residue pairs. Lower values indicate better global alignment.
Since high local confidence can mask poor global geometry (e.g., high pLDDT but high pAE), we combine these into binary foldability criterion: the fraction of sequences with PLDDDT 80, PTM > 0.7, and pAE < 10. This composite measure penalizes degenerate patterns (e.g., repetitive "ABABAB" sequences) that often achieve mislcadingly high scores in isolation.
In addition to structural metrics, compute two distributional statistics that capture whether the model avoids mode collapse:
Token entropy: measures per-position variability of amino acid usage across generated sequences. Let A denote the set of amino acids observed in the generated set, and p(a) the empirical frequency of amino acid A. The token entropy is
H H= Σ logp(a))
aEA
with higher values indicating richer amino acid usage.
Sequence diversity: quantifies variability across full sequences. For batch (x(11 x of equal length L, define pairwise identity as
x(B)
Id(x(m), x(1)
(m) (n)
xx
L
The sequence diversity is then
Diversity
Σ
(d(x(m)) x(()),
B(B -
IEmCncB
which ranges from (identical sequences) to (completely dissimilar).
Training Details for the 150M DLM We train 150M-parameter masked diffusion model on protein data. Training follows the open-source DPLM implementation! using the same transformer backbone MM-MM100M and ESM2-150M. The model is trained from scratch for 500k steps with an effective batch size of 320k tokens per iteration, achieved via multi-GPU, multi-node training with gradient accumulation on L40, H100, and A100. The dataset UniRef50, which contains roughly 40M protein sequences clustered at 50% sequence identity, ensuring non-redundant coverage. UniRef50 is widely adopted resource for protein language modeling.
hhttps://gi htt.bb.com/btteddace///g//
39

---

B.2 TEXT GENERATION
B.2.1
SETUP
Dataset. We the OPENWEBTEXT (OWT) corpus", large-scale collection of English web pages curated to match the distribution of OpenAl's WebText. The dataset is preprocessed with the GPT-2 byte-pair tokenizer and sequences are wrapped or truncated to maximum length of 1024 tokens. For evaluation, held-out split is reserved to compute distributional metrics.
Baselines.
We compare against wide range of autoregressive and diffusion- baadd language
models:
AR (GPT-sttyle:: standard autoregressive language model trained on OWT.
MDLM (Sahoo et al., 2024): masked diffusion language model with uniform random masking.
MDLM FB DFM (Campbell et al., 2024): MDLM augmented with haward-backwart discrete flow matching correctors.
All checkpoints are reused from prior work for comparability.
Training. We follow the same training configurations as in MDLM (Sahoo et al., 2024). Unless otherwise noted, models are initialized with the same GPT-2 tokenizer and architecture. Training details are:
Optimizer: AdamW, with learning rate 3 x 10-4 and linear warmup of 2.5k steps.
Batch size: 32 sequences per GPU, 16 H100 GPUs in total.
Gradient clipping: 0.1.
Training steps: 228k steps, with checkpoints saved every 19k steps.
Sampling and Decoding. We use P2 self-planning (Peng et al., 2025a) in sampling. Unless otherwise specified, we use:
64-bit floating-point, as prior work showed 32-bit sampling led to reduced diversity.
nucleus sampling with = 0.9.
5,000 sequences per model-sample pair.
Evaluation. We evaluate generated text against the held-out OWT distribution. Metrics include:
MAUVE (Pillutla al., 2021): MAUVE directly compares the generated distribution from text generation model to a distribution of human-written text using divergence frontiers. Generated samples and ground-truth corpus data are embedded using an external language model. The two distributions are compared in the embedding space using an area under the divergence curve to summarize both Type and Ttype errors.
Generative perplexity (Gen PPL.): cross-entropy of generated samples under a pretrained language model. This measures the concordance of the model under consideration and strong pretrained language model on the generated text. This can be problematic as only considers generated text. Some distributions of text make it much simpler (and thereby better Gen PPI.) even when not satisfactory to human reader.
Entropy: average per-token entropy of the generated distribution. This measures the token diversity the generated text. Model collapse can be evaluated if entropy decreases substantially.
MAUVE a most robust indicator, while Gen PPL. can be gamed by overconfident sampling schedules. For all evaluation metries match the settings of Wang et al. (2025a).
htttp://Sky1 github. httpbb...t//peeeeeeeeettt..oooo
40

---

B.2.2 ABLATION OF SAMPLING METHODS
Table 5 compares inference-time time planners while holding the denoiser fixed. Across all step budgets, P2-Self sampling consistently yields the best quality-diversin tradeoff: it achieves the highest MAUVE and the lowest generative perplexity, with only a mild reduction in entropy relative to Greedy and Probability Margin. The gap most pronounced in the fast-sampling regime. At T 64, Greedy decoding as MaskGIT (Chang et al., 2022) substantially underperforms, suggesting that purely myopic confidence selection can lock the trajectory into suboptimal local choices when the model is imperfect. Probability Margin (Kim et al., 2025) improves over Greedy at intermediate and large T, but remains below P2-Self and does not scale monotonically with more steps, indicating that margin-based heuristics are less robust to sampling budget. Overall, these results reinforce that planner choice remains crucial degree freedom inference, and that the P2 framework (Peng et al., 2025a) provides more reliable path selection rule for converting denoiser confidence into high-quality generations.
Table 5: Comparison of Greedy, Probability Margin, and P2-Self Sampling across metrics and number of sampling steps.
Metric
Method
= 32
= 64
T 128
Greedy
0.011
0.021
0.056
MAUVE
Probability Margin
0.011
0.039
0.051
P2-Self
0.013
0.046
0.067
Greedy
44.34
34.18
29.38
Gen PPL
Probability Margin
44.27
34.37
29.39
P2-Self
40.19
29.98
24.33
Greedy
5,35
5.30
5.25
Entropy
Probability Margin
5.35
5.30
5.25
P2-Self
5.32
5.24
5.16
B.3 CODE GENERATION
B.3.1 SETUP
We evaluate PAPL on code generation. Training follows the Open-dI.LM framework (Peng et al., 2025b), where we initialize from Qwen2.5-Coder and adapt it to the diffusion setting with bidirectional attention. The model is trained on the FineCode corpus, which integrates the ppe-ammealing-corpp and Sorr--oooo-SSyheiiiee. providing both algorithmic and synthetic QA data. Following the Open-dLLM recipe, each training sequence is randomly masked with ratio uniformly sampled from [0, 1], and the model is optimized with cross-entrope loss over masked positions, weighted by the inverse of the mask ratio. This ensures compatibility with the released training pipeline, data recipes, and evaluation protocols.
Evaluation is conducted on the following benchmarks:
HUMANEVAL (Chen et al., 2021): 164 hand-written Python programming problems designed to functional correctness.
MBPP (Austin et al., 2021b): 974 crowd-sourced Python problems of varying difficulty.
HUMANEVAL+ and MBPP+: augmented variants that extend the original datasets with paraphrased prompts and additional test cases.
For code infilling, we use HUMANDVAL-INFFLL (Bavarian et al., 2022) and the Python subset of SANTACODER- (Sagtani et al., 2025). We report pass@) 1 and pass @ 10 for HUMANEVAL and MBPP, and exact match for SANTACODER- FIM, following the official evaluation protocols. To examine length control, we additionally test under different initial mask spans (4, 8, 16. 32, 64) and report their averaged results.
41

---

Training Loss
Validation Loss
Training Step
Training Step
Figure 5: Training with pure PAPL loss (r 1) leads to unstable behavior, with large fluctuations in training (left) and poor convergence on validation (right).
B.3.2 ABLATION OF SAMPLING METHODS
We ablate the role of the sampling strategy while keeping the denoiser and training setup fixed. Table 6 compares P2-self against several commonly used alternatives: vanilla ancestral sampling with uniform unmasking, (ii) greedy ancestral sampling that always selects the highest-confidence positions, (iii) entropy-based confidence ordering, and (iv) the TopK -mrginn heuristic. Across all six code benchmarks, P2-self is consistently the strongest method. The improvements are most pronounced Pass @ where P2-self exceeds the best baseline by large margins (c.g., 20.8 vs. 12.6 HumanEval and 16.7 vs. 9.2 MBPP), and it also yields steady gains in Pass @ 10 and in the infilling and SantaCoder evaluations. Among baselines, confidence-awam orderings (greedy ancestral and entropy-base baarep confidence) substantially outperform vanilla ancestral sampling, indicating that exploiting denoiser uncertainty during decoding is critical for code generation. However, these heuris- ties remain myopic: they make locally optimal unmasking decisions without explicitly reasoning about longer-range dependencies or future refinement steps, which limits their ability to traverse the high-pu brebb-bbbbbblitt decoding trajectories that matter at inference. JopK-margin is weaker than other confidence-basee methods, suggesting that hard- thesshodreh margins can be overly conservative and discard useful intermediate updates. Overall, this ablation shows that simply reordering unmasking by instantaneous confidence is insufficient; self-planning with path-level lookahead is necessary to gap and achieve robust gains.
fully close the
Method
HumanEval
HumannEval+
MBPP P8 1
MBPP
HumanEral Infill
SantaCoder
P@1
P@:10
P@1
P@10
P810
P@1
P@10
P@
P@l
Vinilla Ancostral
15.2
13.2
2.9
21.8
72.7
53.8
Greody Ancestral
9.3
31.1
8.1
28.7
5.3
20.0
8.7
41.5
75.1
53.7
Fntrogy-hased Co Confidence
12.6
35.4
10.9
29.9
9.2
36.8
15.2
50.7
75.1
53.2
Margnn
26.2
3.9
6.2
33.5
75.0
54.4
Table 6: Performance comparison across coding benchmarks for sampling different methods.
C ADDITIONAL RESULTS
C.1 UNSTABLE TRAINING WITH PURE PAPL LOsS
We investigate the effect of training solely with the PAPL loss under the default temperature setting (r = 1). As shown in Figure 5, training exhibits large fluctuations and fails to achieve stable convergence the validation loss. We hypothesize that this instability arises from the denoiser becoming overly confident: the path weights bias the model toward a narrow set of generation paths, which reduces diversity and leads to overfitting.
42

---

0.4
012
Thakiin Soeps
Figure 6: Vanilla MDM vs PAPL Training Curves. Early in training, the PAPL loss remains small because the denoiser has not yet formed meaningful beliefs about the correct token positions, causing the planner-dependent weights w; to be close to zero. As the model begins to identify correct positions with higher certainty, these weights increase, leading to a temporary rise in the PAPL loss. Once the denoiser becomes sufficiently confident, the loss decreases and eventually mirrors the behavior of the standard MDM loss. These dynamics provide additional evidence that PAPL naturally adapts its emphasis the model's s confidence improves, stabilizing precisely when confidence becomes reliable signal.
C.2 COMPARISON OF TRAINING CURVES WITH VANILLA MDM LOsS
In Figure 6 we observe that the PAPL, loss exhibits distinct training curves compared to the vanilla MDM loss. In particular, initially the PAPL loss is small, because the denoiser has not established any confidence about the correct token positions, and hence the weights w² are near-zero. As the denoiser gains confidence about the correct token positions during the initial training stages, the weights and hence loss increase, before decreasing with curve similar to that of the vanilla MDM loss.
C.3
EMPIRICAL ESTIMATION OF THE EFFECT OF THE APPROXIMATION STEPS
Loss
Greedy
EI- softmax
E2- softmax
Vanilla- softmax
PAPL-
E2-
PAPL
Vanilla
softmax
rand
rand
Value
23294
.002
3.364
17.897
17.897
.002
3.365
18.362
18.362
Table 7: Ablation of the approximation steps used to obtain the PAPL. loss.
Here we provide Table 7, where we empirically estimate the effect of the series of approximations used move from the true ELBO of Proposition 3.2 specialized to Greedy Ancestral (see Corollary A.9) the PAPL loss equation 7. Here, rather than sampling random timestep k, sum along the entire generation trajectory. Greedy the full greedy loss from Corollary A.9, calculated along the true greedy path from Corollary A.9, We remark that this term is extremely large due to the crude lower bound E6 used arrive at A.9 from Proposition 3.2. E1-softmax is the contribution of the term E 89.6.7 from specialization of Proposition 3.2 to the softmax planner from Corollary A. 10 to the true softmax loss, and E22-oofmax is the contribution of E c.. PAPL-softmax is the PAPL loss equation 7 but with x sampled from the softmax path from Corollary A. 10, and Vanilla-softmax is the vanilla loss from equation I but with X₁ sampled from the softmax path. Finally, E E1-rand and E2-rand are the contribution of g c8.05r gl,er Corollary 10 but taken along random path, and PAPL and Vanilla are the actual PAPL and Vanilla losses, which are computed along the random path by definition.
The losses are estimated via accumulating along the entire trajectorics for a batch-size 1024, using the PAPL-fine-tuned protein MDM with 5 and 7 1.
We remark that the computations of the E2 terms are highly unstable. However, we can observe still in Table the following: As predicted by Proposition 3.1, Greedy Loss much larger than Vanilla Loss Vanilla Loss does not provide an upper bound for greedy ancestral sampling, so we would expect the true upper bound be larger. Moreover, as predicted by Proposition A.11, when computed
43

---

along the same path, E₁ is dominated by the vanilla loss. Finally, observe that indeed the path taken in terms of unmasking order does have effect on the losses the losses increase while taking a random path. This is expected, as it is known that taking random path, is trained for in vanilla MDMs, yields significantly worse sample quality than a greedy or soft-greedy path.
C.4 HUMANEVAL PERFORMANCE ANALYSIS
We evaluated the model on 40 HumanEval tasks. While it handles straightforward problems well, performance declines sharply for tasks requiring careful constraint handling, multi-step logic, or less familiar algorithms.
Strengths The model performs best when tasks align with standard Python idioms or texthook solutions. In HumanEval/12; longest, for example, produced the compact and idiomatic implementation in Listing I, which is more direct than the canonical reference.
Listing Canonical solution.
Listing 1: Model's solution for 1umanEva1/12: longest.
def longest (strings: list) str | None: if not strings:
def longest (strings: list) -> str None: if not strings:
return None
maxlen max (len (x) for in strings) for in strings:
return None
return max (string,,,, key=len)
if len -- maxlen: return
The model also demonstrates competence in basic algorithmic tasks. For example, HumanEvall / factorize was solved with standard trial division approach (Listing 3), and string prefix genera- tion (rtmmnEvv1114) and set-based deduplication (HumanEval) 1344) were handled correctly.
def factorize (n: int) -> list [int]
factors = []
while $ == 0:
poctors.appendle
n //=
3
while + i <= n:
while & - 0:
pactors.appent (i)
n //=
1 += 2
if 2:
factors. aprend (n)
return factors
Listing 3: Model's correct solution to HumanEval/25: factorize.
Weaknesses The most common failures stem from flawed algorithmic reasoning. In HumanEval/9: rolling the model produced a redundant nested loop instead of the correct single-pass running maximum (Listing 4).
Another recurring issue misinterpretation of constraints. In HumanEval/3: Relow_zerr the model ignored the requirement to detect negative balances at any point, checking only the final state instead (Listing 5).
Finally, there are occasional catastrophic failures, where the generated code bears no relation to the task. In FumanEval /2: truncate number, for instance, the model produced irrelevant
44

---

Model (incorrect)
def rolling_max (numbers: list) -> 1ist [int]:
result []
max_val numbers
0
while len (numbers)
max_val max (max_val, numbers
while i < len (numbcrs): # Redundant nested doop max (max_val, numbers(i])
max i
desult.appente (max_vall
return result
Canonical
def rolling max (numbers: list) -> list/ints running_ma max = None
result [] for in numbers:
running_max if running_max is None else max (rxunning__axx, n)
result, append (running_max)
return result
Listing 4: Incorrect vs, canonical solutions for HumanEval/9: rolling max.
Model (incorrect)
def below_zero (operatioo:: list [int]) bool:
balance for op in operations:
Missing balance update logic
if balance 0:
return False return balance 0
Canonical
def below_zers (operations: List(int]] bool:
balance
for op in operations:
balance += op if balance 0:
return True return False
Listing 5: Misinterpretation of temporal constraint in HumanEval/3: below.zero.
variable assignments and returned the input unchanged, instead of applying simple modulo operation (Listing 6).
Model (incorrect)
def truncate. nunber (number: float) -> float:
a, b, return number
2, 3 /rrelevant ass inments
Canonical
def truncate. numberrumeerrrreeeeeeee (number: float) -> float:
return number 1.0
Listing 6: Catastrophic failure on HumanEval/2:
truncate.number.
45

---

Discussion Overall, the model succeeds when tasks resemble common idioms well-documented examples, but struggles as soon as prompts introduce layered constraints or require multi-step reasoning. The lack of self-checking mechanisms is evident in tautological or irrelevant code that would be immediately rejected by a human programmer. While the system can accelerate routine coding tasks, it cannot yet be relied upon for problems that demand algorithmic novelty striet logical consistency. Improving decomposition of complex prompts and incorporating verification steps remain key directions for future work.
D PRACTITIONER'S GUIDE
There are main parameters for PAPL, temperature and loss weight a. For practitioners, recommend initial settings of r = 1 and = 1. All experiments in this work leave r 1. We experimented with a values the range 1 ... 10. We found that higher values of a can be more effective in the range of R 5, especially for protein models, but not beyond that. If hyperparameter tuning with PAPI,, we therefore recommend starting with = and a = and doubling a to efficiently search the space.
We recommend this because there is no huge issue with setting a as low as 0, the model will still achieve good performance, as reminder 0 recovers standard DLM training. However, you may not get the benefit of PAPL training, and slow convergence on practical inference paths.
With too high, the training may become unstable, and may not be trained well on all paths. We hypothesis this may hinder performance on more out of distribution tasks, where the planner is not as confident, may hinder performance by learning sub-optimal paths for generation by latching to a specific path too quickly.
46


Reviews:
[{"review_id": "C5BvgZ6lxG", "reviewer": "Reviewer_2qD5", "summary": "The paper shows the mismatch of diffusion objective between training and testing under planning and address it by revising the ELBO to incorporate planner basics in training. The new ELBO allows to effectively exploit any planning strategies at test time and it is simple to implement. The method show performance improvements on various tasks, such as protein sequences, text generation, code generation.", "questions": "1. Is the model sensitive to the temperature and the weight alpha? Need more ablation on text generation and coding. \n\n2. What exactly P2 sampling is  used? Since the P2 sampling has different variants. \n\n3. Apart from P2 sampling, have the authors tried other planning strategies like greedy decoding like MaskGIT, top-k, top probability margin (from Kim et al, https://arxiv.org/abs/2502.06768).", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper tackle the relevant and important mismatch issue of training and testing of diffusion model under planning.\n2. The method is theoretically grounded though it is quite dense.", "weaknesses": "1. Using confidence as a heuristic indicator could be unreliable for training. This is exactly shown in Figure 5 in the appendix. Though I see the authors using the vanilla DLM loss as a rescue, I think this is not an efficient way since at earlier training iterations, the confidence from model is not reliable, inducing noise into training (this is bad). One alternative is to do annealing where earlier step, use uniform weight as a sole vanilla DLM and it is annealed gradually to planner-weighted DLM.\n2. Look at the Table 1 and Table 2, it seems that the method negatively affect the data entropy since both are worse than vanilla model. Why?\n3. What is limitation? It is worth to mention.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "P2nieZnLc7", "reviewer": "Reviewer_LjAj", "summary": "This paper proposed PAPL, a lightweight modification to the classic MDLM training objective that brings a performance boost. The authors first indicate a problem in the domain of diffusion language models, the mismatch between training and evaluation. Based on this insight, the authors derived an analytical form of the planner-aware ELBO (P-ELBO) and show that P-ELBO is theoretically general. Given that P-ELBO is computationally expensive, the authors proposed a cheap approximation, PAPL, and show that across protein modeling, language modeling, and coding benchmarks, adding PAPL to the MDLM loss can boost performance without extra computational expense.", "questions": "* Is the CTMC perspective necessary for the derivation of the ELBO, i.e., Proposition 3.2? If not, can the authors add the derivation from the discrete Markov chain perspective in the Appendix? \n* Can the authors add more explanation on the transition from Proposition 3.2 to Equation 7, at least in the Appendix?\n* The authors claim the new ELBO is planner-aware, so it should be better than the MDLM ELBO. But the final training objective is also a fixed objective; are there any justifications for why this objective is better for various commonly used planners? In other words, is this objective only better for the greedy-sampling planner, or is this objective better for all the commonly used samplers for masked diffusion language models?\n* I would like the authors to clarify the planners used in their experiments. Are the planners just greedy sampling? If so, can the authors provide results with other planners, such as block decoding topk confidence from Nie et al. [1], confidence thresholding from Wu et al. [2], adaptive MDM inference from Kim et al. [3], and remasking samplers from Wang et al. [4]?\n* From a theoretical perspective, the authors claim that P-ELBO is a unified framework for all possible planners. Can the authors provide an instantiation for the four planners mentioned above?\n* In the right subplot of Figure 4, the downstream performance monotonically grows with $\\alpha$; what if $\\alpha$ goes beyond 5? I am assuming that at some point the downstream performance will go down, otherwise we should not interpolate between the MDLM loss and PAPL loss, but use PAPL loss alone instead.\n\n---\n**References**\n\n[1] Nie, S., Zhu, F., You, Z., Zhang, X., Ou, J., Hu, J., Zhou, J., Lin, Y., Wen, J.R. and Li, C., 2025. Large language diffusion models. arXiv preprint arXiv:2502.09992.\n\n[2] Wu, C., Zhang, H., Xue, S., Liu, Z., Diao, S., Zhu, L., ... & Xie, E. (2025). Fast-dllm: Training-free acceleration of diffusion llm by enabling kv cache and parallel decoding. arXiv preprint arXiv:2505.22618.\n\n[3] Kim, J., Shah, K., Kontonis, V., Kakade, S., & Chen, S. (2025). Train for the worst, plan for the best: Understanding token ordering in masked diffusions. arXiv preprint arXiv:2502.06768.\n\n[4] Wang, G., Schiff, Y., Sahoo, S. S., & Kuleshov, V. (2025). Remasking discrete diffusion models with inference-time scaling. arXiv preprint arXiv:2503.00307.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "* The storyline of the paper is clear. The paper aims to address a significant problem in the field of diffusion language models, and its solution is convincing.\n* The mathematical foundation is abundant. I especially appreciate that the authors proved the P-ELBO in an analytical form. The theory part is valuable and useful for the community.\n* The experimental analysis is solid, and the empirical results are strong. I am particularly surprised by the increase in pass@10 in the coding part.\n* The method is simple and practical. With some minimal changes to the training objective, PAPL can improve downstream performance.", "weaknesses": "* As a person with proficiency in discrete diffusion, I have to say the math is a bit confusing. I hope the authors can add more explanation, at least in the appendix, to enhance comprehension (see Question No. 1 for details).\n* The transition from Proposition 3.2 to Equation 7 is not apparent, and more explanation might be needed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cZDShssAeJ", "reviewer": "Reviewer_jzen", "summary": "The paper proposes planner-aware path learning for masked diffusion language models. The authors learn a model  $G_\\phi(z, x_k)$, where $z$ is a sample from the denoiser probabilities and $x_k$ is the k-th state, which outputs the next index to unmask. \n\nThe authors derive an evidence lower bound to lower bound the denoiser + planner log-likelihood, however as the objective requires $L^2$ function calls for sequences of length $L$, the authors propose a different tractable objective. The tractable objective re-weights the denoiser log-likelihood with an extra planner term.", "questions": "1. The proof for the counter-example in proposition 3.1 requires that the denoisers are inconsistent. However, uniform un-masking with inconsistent denoisers can lead to inconsistent sampling as well. \n\nIt would be beneficial to a reader to see a re-written proposition 3.1, as the authors construct a counter-example rather than ***prove*** that greedy decoding does not sample from a unique learned joint distribution. Under the constructed counter-example, even uniform sampling would not sample from a single joint distribution as the denoisers are inconsistent.\n\n1. If the denoiser is consistent, then would greedy sampling work? As MDLMs learn an any-order autoregressive model, would any decoding order at inference would produce valid samples? \n2. If the denoisers are consistent, then is the following line correct:\n    1. *The key takeaway is that the ELBO in equation 1 is only valid for the uniform unmasking process (line 247)*. \n3. what does the planner learn at optimality\n\nExperiments\n\n1. Can the authors compare PAPL to greedy sampling in the experiments? \n2. In several experiments, the gap between the MDLM+planned and MDLM+baseline with  experiments are not substantial. Can the authors provide a confidence interval?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "See questions", "weaknesses": "See questions", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hbIL7I8ZLQ", "reviewer": "Reviewer_TRnB", "summary": "The paper addresses the training-inference mismatch in masked diffusion language models (DLMs). While DLMs are trained with uniform random masking/unmasking, they are deployed with planners that bias the denoising order (e.g., greedy, path planning). The authors prove that this mismatch causes the standard DLM ELBO to be violated at inference (Proposition 3.1). They derive a Planner-Aware ELBO (P-ELBO) that incorporates planner dynamics into the training objective (Proposition 3.2), then propose PAPL—a practical algorithm that simplifies P-ELBO into a planner-weighted cross-entropy requiring only a one-line code change. Experiments demonstrate consistent improvements: ~40% relative gain in protein foldability, up to 4× MAUVE improvement in text generation, and ~23% relative gain in code generation pass@10, all under identical model architectures and inference planners.", "questions": "1. Can you measure the correction empirically at different training stages across domains? Some figures may be helpful.\n2. How should practitioners choose the two hyperparameters?\n3. How tight is the P-ELBO bound in practice?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The core contribution is original and well-motivated. Proposition 3.1 formally proves that greedy planner-based sampling can violate the standard ELBO inequality (log p_θ < ELBO), meaning the training objective no longer lower-bounds the inference distribution. This is a concrete theoretical finding that justifies rethinking the training procedure, not just an empirical observation.\n2. The theoretical framework is rigorous and general. The P-ELBO formulation (Proposition 3.2) unifies multiple existing planning strategies: uniform (standard DLM), greedy (MaskGIT), and soft-greedy, as special cases. \n3. The implementation is practical and low-friction. PAPL reduces to standard MDLM loss augmented with planner weights and requires no architectural changes or additional forward passes during training. Algorithm 1 shows this is essentially a one-line modification, making adoption straightforward.\n4. Empirical validation is comprehensive across diverse domains. The paper demonstrates improvements in protein generation (Table 1: higher pLDDT, pTM, lower pAE, +40% foldability while maintaining diversity), text generation (Table 2: consistent MAUVE and perplexity improvements across sampling budgets), and code generation (Tables 3-4: gains in both completion and infilling). The ablations (Figure 2, Figure 4) support design choices and show robustness.\n5. The presentation is clear and well-structured. The method section logically progresses from dynamics formulation to ELBO violation to the new objective. Experimental setups, metrics, and baselines are clearly described with sufficient detail for reproduction.", "weaknesses": "1. The gap between P-ELBO and implemented objective is not quantified. Section 3.4 drops the correction term by detaching gradients through the planner, converting the full P-ELBO to a simpler weighted cross-entropy. The paper provides no analysis of the induced approximation error, theoretical bounds, or empirical measurements of this gap. \n2. Although the framework is generally applicable for any planner, the derived training objective and results are only for soft greedy decoding. For non-differentiable or hard planners, the correction term becomes non-trivial, and the paper does not provide a full derivation or empirical evidence beyond soft greedy-like planners. \n3. Text evaluation is limited to unconditional generation. While MAUVE and generative perplexity on OpenWebText are standard, they don't test whether planner-aware training improves conditional generation capability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)