wd1:  Weighted Policy Optimization for Reasoning in Diffusion Language Models
Authors: Xiaohang Tang
Rares Dolga
Sangwoong Yoon
Ilija Bogunovic
Venue: ICLR 2026 Poster (2026)
Keywords: Diffusion Language Models
Reinforcement Learning
Reasoning
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose a novel policy optimization method for dLLMs reasoning, reducing the error caused by log-likelihood approximation error.

Abstract: Improving the reasoning capabilities of diffusion-based large language models (dLLMs) through reinforcement learning (RL) remains an open problem. The intractability of dLLMs likelihood function necessitates approximating the current, old, and reference policy likelihoods at each policy optimization step. This reliance introduces additional computational overhead, and can lead to large variance and estimation error in RL objective -- particularly in computing the policy ratio for importance sampling. To mitigate these issues, we introduce wd1, a novel ratio-free policy optimization approach that reformulates the RL objective as a weighted log-likelihood, requiring only a single approximation for the current parametrized policy likelihood. We formally show that our proposed method can be interpreted as energy-guided discrete diffusion training combined with negative sample unlearning, thereby confirming its theoretical soundness. In experiments on LLaDA-8B model, wd1 outperforms diffusion-based GRPO (d1) while requiring lower computational cost, achieving up to a +59\% improvement in accuracy. Furthermore, we extend wd1 to denoising-stepwise weighted policy optimization (wd1++), achieving state-of-the-art math performance of 44.2\% on MATH500 and 84.5\% on GSM8K with only 20 RL training steps.

Published as a conference paper at ICLR 2026
wdl: WEIGHTED POLICY OPTIMIZATION FOR REASONING IN DIFFUSION LANGUAGE MODELS
Xiaohang Tang Rares 2 Sangwoong Yoon Ilija Bogunovic
Department of Statistical Science, University College London, United Kingdom
PCCL AI Centre, University College London, United Kingdom
3 Grraduaee School of AI, Ulsan National Institute of Science and Technology, South Korea DDppartment of Mathematics and Computer Science, Universität Basel, Switzerland
UlPat
ABSTRACT
Improving the reasoning capabilities of diffusion-based large language models (dLLMs) through reinforcement learning (RL) remains an open problem. The intractability of dLLMs likelihood function necessitates approximating the current, old, and reference policy likelihoods at each policy optimization step. This reliance introduces additional computational overhead, and can lead to large variance and estimation error RL objective particularly in computing the policy ratio for importance sampling. To mitigate these issues, we introduce wdl, a novel ratio-free policy optimization approach that reformulates the RL objective as a weighted log-likelihood, requiring only single approximation for the current parametrized policy likelihood. We formally show that our proposed method can be interpreted energy energy-guided discrete diffusion training combined with negative sample unlearning, thereby confirming its theoretical soundness. In experiments on LLaDA-8B model, wd/ outperforms diffusion- based GRPO (d/) while requiring lower computational cost, achieving up to a +59% improvement in accuracy. Furthermore, we extend wdl to denoising-stepwise weighted policy optimization (wdl++), achieving state-of-the-ar math performance of 44.2% on MATH500 and 84.5% on GSM8K with only 20 RL training steps.
INTRODUCTION
Diffusion-based large language models (dLLMs) have recently gained attention as promising alternatives to autoregressive (AR) models for language modelling tasks (Nie et al.] 2025b]) Ou cet al. 20255 Yang et al., (025). Unlike AR models, which generate tokens sequentially, dLLMs iteratively refine entire response sequences through denoising process. A primary advantage of this approach the significantly improved inference efficiency. Notably, recent closed models such order of magnitude speed-up up in generation compared to AR models, while maintaining comparable generation quality. Furthermore, open-weight dI.LMs demonstrate competitive performance on standard language benchmarks, with smaller models (Lou et al. 2024 Ou et al.) 2025 Nie et al. 2024) achieving parity with equivalently sized AR baselines, and larger-scale models such as LLaDA-8B (Zhu al., (025a) and Dream-7B (Ye et al. 2025) extending this trend at scale. While dLLMs demonstrate strong performance in text generation, remains an open and important question how best to fine-tune dLLMss using RL paradigm that has proven highly effective in alignment and improving reasoning capabilities of AR models (Ouyang et al. 2022: Shao et al.) 2024).
as Mercury (Labs et al. (025) and Gemini Diffusion achieve
A key challenge in applying reinforcement learning (RL) to dLLMMs is the intractability of their likelihood functions (Zhao et al. (0255 Yang al. (025). which necessitates approximation for policy optimization. Applying approximated log-likelihood for diffusion- based GRPO (Shao et al.) 2024 Zhao et al., 2025) can exponentially amplify the approximation error and lead to large variance
Equual contribution. Code: https//iitub..oo//xioaaaaa
Corresponding authors (i (i bogunovi cêun @uunibaa ch, swyoon duni ac. kc.kr)

---

Published as a conference paper at ICLR 2026
when computing the policy ratio for importance sampling. Moreover, GRPO requires the estimated likelihoods of the current, old, and reference policies at every training step. leading to significant computational overhead. These issues can be further exacerbated the completion length and the number of diffusion steps increase.
To address these challenges, we propose wdl, policy optimization approach with weighted log- likelihood objective for dLLMs. Crucially, this objective dispenses with explicit policy ratios and relies on a single likelihood approximation, thereby avoiding the potentially large bias and variance in policy ratio, and reducing the computational overhead. Our principal contributions are:
We propose novel reinforcement learning method for dLLMs, termed wdl, which formulates the objective as a weighted log- likelihood of outcome sequence, derived from reverse- KL regularized policy optimization. The weight, defined as (-w+ + w and dependent on the advantage A, balances two terms: wt X exp(A) increases the probability of higher-advantage completions, while w xx exp(-A) decreases the probability lower-advantage ones. Together, this mechanism amplifies beneficial outcomes meanwhile actively reducing detrimental ones.
We prove that our proposed RI. method for dL.LMs can interpreted as jointly training an energy- guided discrete diffuusin model-guided by the advantage function- and unlearning low-advantage data, thereby steering generation toward higher-ad advantag completions.
We conduct experiment with LLaDA-8BB -B-nstruttttt model (Nie et al. 2025a). Compared to the baseline method d/ (Zhao et al. 2025), our method wdl achieves 76.4% % on Sudoku (Arel, 2025) (+58.8% over df) and 51.2 % on Countdown (Pan et al. 2025) (+16% % over dl), without requiring supervised fine-tuning (SFT), and with significantly less computational burden in RL. training.
We further extend our method leverage intermediate completions generated in the decoding process, which we call wd/++. The extended method surpasses several concurrent RL. for dLLMS methods, achieving state-of-the-art performance 44.2% % on MATH500 and 84.5% % on GSM8K with only 20 training steps, and 10x fewer rollouts compared to the baseline methods.
PRELIMINARIES
We denote the generation policy of diffusion based Large Language Models (dLLMs) by ng. Denote prompt 9 E D, and completions E O. Notably, the reward function denoted by R(g. o) in this work is not limited to verifiers. We use superscript & to indicate the k-th token of completion: of or z
2.1 DIFFUSION LARGE LANGUAGE MODELS
The most promising discrete diffusion for language modeling is masked diffusion, which gradually corrupts the sequence with mask token (Sahoo al. 2024 Ou et al. [025] (hi et al. 2025 Lou et al. 2024). Let t E [0, 1] denote the diffusion timestep, and H; as the masked sequence at step The fully denoised sequence (i.e., the completion o) represented by wo, and the forward process (prjo 22o)) is formulated as a continuous-timg Markov chain. The transition kernel Qe
is absorbing (Austin et al. 2023 Camphell et al. 2022), meaning that at time I, Q Q absorb is a constant matrix (See Definition] 2
where a is decreasing scalar noise schedule and Q uuubbb
This work aims to apply reinforcement learning to fine-tune masked discrete diffusion models such as LLaDA (Ou et al. 2025| Zhu et al.) 2025a), which models the clean data distribution conditional on masked sequence (z 2.). The negative Evidence Lower Bound (ELBO) is reduced to
simple objective called Denoising Cross Entropy (DCE) (Ou et al. 2025 : Vzo ~ Pdata+
K
(DCE(20) -E =cmpg((r
M
(a [mask]] log (a(a 2x)
(I)
where K is the length of the sequence and z is the k-th token of zp. Specifically, intermediate steps t sampled from uniform distribution, and masked sequence is sampled following the predefined forward process [z, zn). DCE can be used to approximate the marginal likelihood log n (In)
for supervised fine-tuning and reinforcement learning (Nie et al. 2025al Yang et al. 2025).
2

---

Published as a conference paper at ICLR 2026
2.2 EXISTING POLICY OPTIMIZATION METHODS
The base method of current prevailing RL fine-tuning algorithms is Trust Region Policy Optimization (TRPO) (Schulman al. 2015), which forward KI. divergence applied to restrict the update:
max DDD (-e)) A (g. o) ADxL Kotd (=
|| no (-11)
where Ma is the advantage function, and o are denoted as the prompt and (clean) response, respectively. Proposition (Appendix A) demonstrates the monotonic policy improvement of TRPO.
PPO then extends the soft constraint (KL penalty) to clipping policy ratio na Ta(-lg)/Taia(- (g) and
employing pessimism for policy update, further employed in fine-tuning (Ouyang al. 2022
with additional reverse-Ki KL regularization w.r.t. the reference policy Tet Group Relative Policy
Optimization (GRPO) (Shao al. 2024) simplifies PPO by sampling group of completions {o; [=] and approximating their advantage with their normalized rewards. This advantage is corrected by subtracting the mean reward across the group (Liu al. 2025; A, Riq, o;) mean (R(g o1.))), which we refer to as the group-relative advantage.
2.3 POLICY OPTIMIZATION FOR DLLMS
Adapting GRPO diffusion blafe large language models (dLLMs) presents notable challenges, since dLLMs generate outputs via a -on-auutregressivee iterative denoising process, making the computation of log (olg)) intractable and necessitating approximation for policy optimization.
Existing works by Nie et al. 2025a); Yang et al.
2025) employ ELBO for per-token log- log-likelihood
approximation following DCE: (1x) E, Eecur o |
(2,, [ mask] log [=, 4)], where
1/t in DCE and a large sample size of t. resulting in inefficiency for online RL. A biased but efficient method is
1 in UniGRPO (Yang al. 2025). However, accurate estimation requires
introduced in dl (Zhao et al. 2025), requiring only sample at t 1: o= (=xxx) log x (1 q¹). where prompt q' is randomly masked, 21 is fully masked response.
In diffusion- based GRPO (Zhao et al.) 2025
Yang al. 2025), policy ratio is then computed
using the approximated ligliioelsoood: r (0)
[o(o J/ Toua (o)) = exp (o - x (o+)) for
importance sampling estimating the objective of GRPO:
G
D-D
GK 1=1
min Alip),
- 3D ( || mTee((1)
(3)
However, existing approaches are hampered by their re- liance on extensive likelihood approximation to compute the policy ratio. In current diffusion based GRPO methods, the ratio is computed as r se exp (Ø (of)) - Joa (o so approximation errors in likelihood can be exponentially amplified. As formally shown in Appendix [A. the re- sulting error in the estimated RL objective becomes more severe when less accurate log-likelihood approximations are used, such in d1, ELBO used in DCE and Uni- GRPO when the Monte Carlo sample size t is small.
Policy Ratio
ELBO
101
d1
100
20
100
Sample Size oft
Figure 1: Example policy ratio value r computed using ELBO and approxi- mated likelihood in dl on GSM8K after a policy update. Ratio's 's unclipped in- terval is [1 c, 1 c), where 0.5. ELBO-based likelihood approximation yields high-variance ratio estimates; dl induces a biased ratio that can deviate substantially from ELBO. Both meth- ods suffer from efficiently and accurately compute ratios.
Although increasing in the ELBO estimator can reduce approximation error, the induced ratio estimates can still exhibit high variance, as illustrated in Figure Although alternative approximator such as that in d/ can improve efficiency, but yields a biased ratio that can differ substan- tially from the ELBO-based ratio, thereby introducing systematic bias into the RL training objective. Finally, GRPO requires applying the approximation function ó separately to three policies-on -og Told- and गhe -hich further increases computational overhead.
3

---

Published as a conference paper at ICLR 2026
3 wdl: WEIGHTED POLICY OPTIMIZATION FOR DLLMS
In this section, we introduce wdl, novel RL algorithm that eliminates the need for approximating the likelihood (policy) ratios for importance sampling, aiming to reduce the computational burden, and the variance and approximation error in computing the RL objective. We further extend our method to wdl++ by applying denoising-ster stepwise policy optimization.
3.1 REINFORCEMENT LEARNING WEIGHTED LOG-LIKELIHOOD MAXIMIZATION
Prevailing RL methods are based on constrained policy optimization (Belousov & Peters 2017). penalizing the deviation of current policy no(-lq) from the old policy Tota-114) TRPO objective (Equation (2) applics a forward-KL penalty. We instead adopt reverse-KL penalty augmented with the reference policy regularization DKL (To [g] mer(-1)));
max (gED,oNTO((1) A The (q, (q.0) 0) ADki. ma (-14)
(ola
BKL (To(-lg) m
Note that the monotonic improvement guarantee still holds when using reverse-KI. penalty, as we show in Theorem [ From the method of Lagrange multipliers, the solution to Equation has the following form (Peng et al. 2019; Rafailov et al. 2023):
AAa (g. A + 3
n" (+19) x (oaa(-l)) (Tr((11))0(A exp
(5)
As the analytic form of the optimal policy 't known, we can train our policy by directly minimizing
DKL (x" (-14) Il m=(-)g)). This minimization can be expressed as the following weighted log-likelihood
(WLL) loss, where the weights xx exp (1AAT v 1 X+n and the samples are obtained from the
geometric mixture policy ( x Tau(-1g) (((+)) (eer(-l))
(See Proposition
2: Vg ~ D,
-exp exp (WA Tou (q, o)) log (ogmo(ola)
(0) = Bown x)
(6)
exp (vA
log No (o:lg)
T
Σラ=
G exp (vA,))
As shown in Equation 77. we approximate the advantage function using the group-relative advan- tage Å and normalize the weights, thereby limiting their magnitude and reducing variance in loss computation. Notably, dividing by the normalization constant does not affect the solution, since it independent of the completions. The resulting objective does not involve ratio (=(-)g)/ (oaa(1l for importance sampling or To(- (-14) /rr((-lq) for regularization, successfully avoiding the potential amplification of ligeliiiddood approximation error and large variance in diffusion GRPO.
Although the objective (0) in Equation 7 avoids the likelihood ratio estimation. it has two limitations. First, the algorithm is not fully utilizing all the completions. Due to the exponential form of the weighting, completions with relatively low advantage referred negative samples may receive vanishingly small weights, and do not contribute to learning. Second, due to the likelihood- maximization property of WLL, the likelihoods of all sampled completions are increased, even for negative samples. This issue is exacerbated in scenarios where all completions attain identical but low rewards (e.g. 0), thus all weights become equal and the likelihoods of these suboptimal samples are nonetheless reinforced.
3.2 wdl: FULLY UTILIZING COMPLETIONS
We propose wdl, an improved weighted log- likeliioddood objective that explicitly reinforces positive samples and penalizes negative samples:
G
Ludl (8) = (o, -1
Σ
(
w (g, + w (=( log ro(os11).
(8)
G
iml
4

---

Published as a conference paper at ICLR 2026
where the weights are based on group-relative (GRPO) advantage and are further normalized to avoid overly imbalanced weight A; R(q. 04) (ag,
exp (v)
exp ( vbA.))
G w w=(g,00 0.) Ej= exp (A),))
w (q,
Σj=1 exp (- A),
wdl objective balances positive and negative samples through a complementary penalty term, w (q. o:) log (T( 14), which minimizes the likelihood of low-advantage completions. This penalty induces negative gradients, thereby accelerating divergence from undesirable completions. Moreover, in the extreme case where all completions exhibit identical advantages, the optimization naturally halts since w+ w, thereby addressing the concern increasing likelihood of negative samples proposed in Sec 3.1 We demonstrate the effectiveness of this combination via ablations C.2
Our method wdl, simple ratio-free policy optimization based on weighted log- log-likelihood objective for dLLMs, is formally presented in Algorithm We first obtain G completions {o} sampled from geometric mixture (-1(()4)) x Tota( (g) A/Tx++) -Trer--lo [Trer(-lg) (g)))((++))) (line 5). Since the base model LLaDA parametrizes the clean token prediction ( [a,, q) for denoising, we approximate log ( Fr 9) R À log Tota (aff) x4, , q) + 3 log <==((=]]) Fe; 9) as the logits of the denoising distribu- tion each step t. We then use the samples to compute weights in Equation (9) (line 6). In weights computing, we leverage completions from all groups to estimate the normalization constant, order to restrict the the gradient norm and stabilize the training. Finally in line 8, we approximate the log- log-likelihood of completions, and compute objectives for policy update. Likelihood approximation in dl (Zhao al. 20225) is directly applicable wdl: log to (q) Ex log no (a (,,,")), where y' is randomly masked from prompt at every gradient step.
3.3 wd/++: STEPWISE WEIGHTED POLICY OPTIMIZATION
The decoding process in dLLMMs relies on confidence-ben based remasking (Wang et al.) 2025b). At each denoising step E {1.... ; I.] in decoding, clean data is predicted conditional on the masked sequence Ij and then tokens with low-confidence are re- re-masked for further denoising, which construct a refine- ment process. Since current diffusion RL methods only use the final predicted clean completion for training, there are bunch of clean completions in the intermediate denoising steps remaining unused.
To leverage intermediate clean completions, we extend our weighted log- liglliielhoodd objective to a step-wise formulation based on DCE, which we term wdl++. In wdl (as well as in GRPO), a group of completions for Jo is sampled for policy optimization. In wdl++.
expand this group to {O. where O, { | Foj Till( Zz, q), Tojı o,, / = L.. which contains all generated completions during the decoding process, including intermediate ones. The expanded group of completions is then used to estimate both the advantage function and the corresponding weights. The resulting loss objective is defined as:
Loddl+ ++ (0) = E
Σ (- w -0, (o))) w (g, vqp)) -log To (T (2, 4)]
(0.) Ca (l..
E0,
(10)
4 THEORETICAL INSIGHTS: ENERGY-GUIDER DIFFUSION SAMPLING
In this section, present a novel theoretical interpretation of policy optimization for dLLMs. We prove that the advantage-weighte log- liglliil objective (wdl) for dI..Mss can be viewed as energy-guided discrete diffusion training combined with negative sample unlearning.
Sampling from the solution policy of the reverse-KL policy optimization, as described in Equation (5) can be interpreted as energy-guides sampling. where the energy function E(g.-) A Tau (q, Equation (5) defines the marginal distribution of the clean data (xi = o) which we denote as pū
To obtain the guidance at intermediate time steps t O, we define the forward diffusion process for the target diffusion policy πT as following.
Definition 1. The forward diffusion process of the target policy (T*) satisfies Pijo(ael =
Pejo(z)/00), where Pejo is the forward process of old diffusion policy wold-
'To adapt the setting of diffusion, we use 3e denote the (masked) completions, and omit the prompt q-
5

---

Published as a conference paper at ICLR 2026
Algorithm 1 wdl: Weighted Policy Optimization for dLLMs
Require: Reference model Tief. prompt distribution group size G, reward function R, dLLM Te. regularization hyperparameters À and
1: Initialize Th <- Tie
2: while not converged do
3: Vold
4: Sample prompt ~ 'D and G completions ~ ( 9). Vi E [G]
5: Compute advantage A, = R(g. 0:) mean Rgg 01.G) ). vi E [G]
6:
Compute weights w and in Equation 9 V [G] for gradient update iterations = 1,2,.... ................................ do Compute approximated log-likelihong log (o 9: Compute objective fo (0) in Equation 8 or Equation 10 and update
7:
8:
10: end for
11:
end while
12:
return Te
Since the reference difffusion policy is the initial policy, three policies have identical forward diffusiio process, being (x+/20) = Pejo (zo) = Pijo (za) and thus, (xee = Pajo (1o()//))))) where p' is the geometric mixture diffusion p/ (3e (2o) prjo(= pft (x, (n)) We then obtain
the energy guidance at all time step
Lemma (Intermediate Energy Guidance Discrete Diffusion). The marginal probability distribu- tion of the masked responses (Ix) in the diffusion process satisfies pł (xx) = (x(x)))ex))) (A((z)) /Zz
which induces an pnergy-guided discrete diffusion:
(oje(rro(te) (oje(Fol exp(A(zn) As(x:)),
(11)
where -A((x)) = log Exo--p (-1,,) /exp (A(zo) is intermediate energy function for t > 0, and A()) is advantage function (Proof in Appendix A.3 A.3
The guidance provided in Lemma I demonstrates that it directs the sampling process toward generating completions that exhibit higher advantage values. However, conducting training-fres guided sampling following Equation (T]) requires estimating the posterior mean of the exponential of advantage (Lu et al. (023). Rather than relying on such estimation, we instead aim to find the training objective to directly approximate the target guided diffusion model.
Since existing masked dLLMs parametrize the concrete score (Meng et al. 2022). to apply the energy guidance, we aim directly approximate target guided concrete and é is identical to except the i-th token is unmasked (i.c. æf [M] and f * [M]. Concrete score is defined the marginal probability ratio between é; and 2y;
Denote 34 (xt .x)
t( def p( pla,.
(12)
We prove that the training objective to approximate the guided concrete score can be simplified as weighted Denoising Concrete Score Matching (D-CSM) (Meng et al. 2022):
Theorem 1. The model se approximates the concrete score of the energy-guided discrete diffusion
p" when the following loss objective is minimized. This objective is in form of advantage-weighted Denoising Conerete Score Matching, which we AW-D-CSM:
P6 ( (200)
LAW -Exo~p6() exp (A(zo))
[13
(13)
Advantage Weight
(occu((o))
We provide the proof in Appendix A.3 Additionally, D-CSM is approximation of CSM (Meng
et al. 2022), which is equivalent to Denoising score entropy (DSE) (Lou et al. 2024). For all Zo. is satisfied up multiplying constant that (D.) < (xs e LOSE (a) Lpcn(

---

Published as a conference paper at ICLR 2026
(Ou et al. 2025). Therefore, AW-D-WW--SS D--MMMMM can then be applied for both SEDD (ou et al. 2024) and RADD (Ou et al. 2025) model such as LLaDA. Denote Pe as the concrete score reparametrized model. AW-D-CSM can be converted to weighted denoising cross-entropy loss (AW-DCE):
LAW.DCE -Bzompb() exp (A(zo)) (1:)))
[
Σ x [=]mmss]
log po (cogle
(14)
DSE and DCE objectives both can be used for likelihood approximation in fine-tuning (Ou et al.
2025| Nie al. 2025a; Yang et al., 2025) since they can serve negative ELBO (Lou al. 2024 Shi et al., 2025). Thus, the advantage-weighted objective AW-DCE (or AW-DSE) used to learn energy-gum guided score is in weighted log- log-likelihood form as in wdl with only w (i.e. WLL loss in
Equation )6, which contributes to our main theoretical findings:
Remark 1. In the context of applying RL to masked discrete diffusion, the advantage-weighten log- likelihood (WLL) objective (Equation (()) induced by reverse-KL policy optimization, is equivalent to the objective of training energy-guided diffusion models, where the energy function is the negative advantage. Formally, Lw1J. &> LAW.DCE when DCE is used for likelihood approximation.
Remark 2. Additionally, based DCE likelihood, the additional penalty term on negative samples used to exlend WLL to wdl loss can be viewed as applying data unlearning by minimizing the ELBO (Alberti al. 2025), where the data { } (negative samples) has probability distribution Pduta ( ) Xx (o( exp(-A(g, 2y }), which corresponds to a Boltzmann distribution that places higher probability mass on regions with lower advantage values (more details in Appendix D.I).
5 EXPERIMENTS
In this section, we empirically validate the following key advantages of our approach:
i) Improved reasoning capabilities than existing methods on popular reasoning benchmarks:
ii) reduced computational burden, as reflected by decreased runtime, lower FLOPss and numbers of function evaluations (NFEs) per training step, number of training steps and rollouts; and
iii marked performance gains attributable to the incorporation of samples with low-advantage.
To evaluate our approach, we next detail the experimental setup and implementation.
Experimental Setup. We perform reinforcement learning (RL. fine-tuning on the LLaDA-8B- ------B--nstruct
model (Nie et al. 2025a) with Low-Rank Adaptation (LoRA) on: GSM8k (Cobbe et al. 202T). MATH (Lightman et al. 2023). Sudoku (Arel, 2025). ) and Countdown (Pan al. 2025)) ). As for decoding,
we follow the default strategy Mounier& & Idchpour (2025): Arriola et al. (2025); Wang et al. (2025b). Our main baseline is d/ (Zhao et al. (025), the first RL method developed for masked diffusion L.MMs (dLLMs). We reproduce the baseline methods Diffu-GRPO, which applies diffusion- based GRPO training directly the LLaDA base model, and d/, which performs SFT before applying Diffu-GRPO. We use sIK (Muennighoff et al. 2025) data for SFT in dl. We also compare with SDPO (Han et al. 2025), TCR (Wang al. 2025d), and MDPO (He et al. (025) on benchmarks GSM8K and MATH500.
MDPO is reproduced based on the official implementation and the training dataset (He et al. 2025).
Implementation. As for wdl, we conduct training on the same dataset as in dl (Zhao al. 2025):
training splits on GSMSk and MATH, and the dataset splits provided by Zhao et al. (2025) on Sudoku and Countdown.
our implementation of wdl, we apply the same likelihood approximation method as dl. The hyperparameters used in our method and our reproduction of d/ are listed in Table o and Table 5 As for wd/++, we train on dataset provided by (He et al., 2025), which is sampled from OpenR1 dataset (Face) (025)). Since previous works (Yu et al.) 2025) have demonstrated that the reference policy is empirically unnecessary, we set 3 and
to eliminate Net in practice. We report results using zero-shot and pass@1 evaluation sequence lengths of 256 and 5 512 tokens.
5.1 MAIN RESULTS
Superior Reasoning Ability. In Table we observe that wdl, even without supervised fine-tuning using any supervised data, consistently outperforms our reproduced implementation of dl. Notably,
In the technical report version of this work, our method achieved scores of 25.2 and 24.2 on Sudoku after 5K training steps. In this paper, extend the training to 12.5K steps, and wd/ results in improved performance.
7

---

Published as a conference paper at ICLR 2026
Table I: Test Accuracy (%) of wdl and dl. We reproduce dl and vary completion length. Our
approach without SFT, demonstrates particularly higher accuracy on Sudoku Countdown.
Sudoku
Countdown
GSM8K
MATH500
Model / Gen Len
256
512
256
512
256
512
256
512
LLaDA--B-Instruue
6.7
5.5
19.5
16.0
76.7
78.2
32.4
36.2
+ diffu-GRPP
16.1
11.7
27.0
34.0
80.7
79.1
34.4
39.0
+ SFT + diffu-GRPO (dI)
17.6
16.2
25.8
35.2
78.2
82.0
34.4
38.0
+ wdl
76.4
62.8
51.2
46.1
80.8
82.3
34.4
39.0
Table 2: Comparison of Training Cost on 4 A100. We show SFT cost, average training time, FLOPs evaluated by DeepSpeed Flops Profiler, and theoretical NFEs per training step which includes μ gradient steps. wd/ removes SFT and has less cost per-step in RL than dl.
SFT
RL Training
Method
Time Cost
Time Cost
FLOPs
NFEs for Likelihood
dl wd1
2.01 hrs
103.5 sec/step
9.922 x 10t5pstep
(n + 2)/step justep
0 hrs
81.16 sec/step
8.887 x /step
wdll surpasses d/ by 43% test accuracy on the Sudoku task, and achieves up to a 25% improvement on Countdown with maximum length 256. Relative to the base LLaDA model, the performance gain reaches as high as 54% on Sudoku and 42% on Countdown. On math problem-soll solving benchmarks GSM8K and MATH500, wdl attains slightly higher accuracy. Nevertheless, the extended method wdl++ obtains significantly better accuracy. Table 3 (left), we further compare with concurrent baselines released in recent months. wdl++ outperforms the baselines including strong MDPO.
Reduced Training Cost. Table demonstrates that the training cost required by wd/ substantially lower than that of dl. Unlike dl, wdl does not require SFT stage, which alone accounts for approximately two hours of training in d/. wd/ achieves additional speedup during the RI. phase, where runtime is measured by averaging over μ4 inner gradient steps per global step. Notably, the time efficiency gap is expected to widen further under settings with larger maximum sequence lengths and more diffusion steps. This efficiency gain is further supported by a reduced FLOPs and number of function evaluations (NFEs) per step, as wd/ bypasses the need to approximate the likelihood of the old policy. We exclude NFEs associated with sampling, since both methods share identical sampling costs as wdl removes the reference policy regularization.
In Table (right), we report the training cost required to obtain the best post-trained models on GSM8K and MATH500, measured in terms of the number of training steps and rollouts. wd/++ requires 10x fewer rollouts to achieve superior performance, clearly demonstrating the efficiency of our method. This rapid convergence arises primarily from the exponential ad advantage weights applied to the log-likelihood in wdl, In contrast, standard RL methods such as GRPO and PPO weight the log- likelihoooood (or policy ratio) terms directly by the advantage function.
5.2 ABLATION STUDY
We present an ablation study in Figure Notably, we observe that supervised fine-tuning (SFT) yields only marginal improvements within our approach, with slight gain in the Sudoku task. This contrasts with d1, where SFT plays significant role in improving performance. These findings indicate that wdl can eliminate the need for an SFT phase, thereby simplifying the training pipeline and substantially reducing computational cost. Additionally, we evaluate the impact of removing the negative- weighted term by setting w = O, thus relying solely on the positive advantage weights w + We provide further ablation on the combined method between w and w in Table g The results highlight the importance of explicitly penalizing the likelihood of low-advantage completions, thereby reinforcing the role of negative samples, and emphasize the critical balance between the two weights.
TTo facilitate efficient ablation studies, we restrict our comparisons to checkpoints saved prior 5000 steps.
8

---

Published as a conference paper at ICLR 2026
Table Left:
Extended method ++ compared concurrent RL methods fine-tune LLaDA- 8B-Instruce Methods denoted by "(full)" perform full fine-tuning Right: Training cost to obtain the best model on GSM8K and MATH500. We count the total number of steps policy iteration (model weights update), and the number of rollouts used for training (see Table 8 for details on counting).
Training Cost
Model
GSM8K
MATH500
LLaDA-8B Instruct
-nnssrructttttttttt
78.2
36.2
# of Steps
# Rollouts
30000
diffi-GRPO Zhao al. 2025
80.7
39.0
dd (Zhao al. 2025
82.0
38.0
-SDPO [Han al. 2025
(full)
81.2
TCR (Wang et al.
83.0
41.4
MDPO (He
2025
(full)
83.4
43.4
wdl
82.3
39.0
7500
wdl (full)
82.7
43.6
1280
wdl++ (full)
84.5
44.2
wd1+
MDPO
d1
Table Ablation on SFT and Negative Samples Weight (w7). We conduct wdl training after SFT (wdI-SFF) and with only wt (namely wdl-P or WLL defined in Equation 66 that well performs better without SFT on planning and math tasks. Removing negative sample reinforcement (w: significantly hurts performance, highlighting its importance.
Results show
Sudoku
Countdown
GSM8K
MATH500
Model Gen Len
256
512
256
512
256
512
256
512
wdl-P (WLL)
6.69
6.84
13.67
4.69
65.66
78.17 17
29.40
22.80
wdI-SFT
26.5
24.2
43.4
43.4
80.7
82.0
36.4
39.0
wdl
25.2
24.2
51.2
46.1
80.8
82.3
34.4
39.0
We further assess sensitivity to the relative weight- ing of positive and negative samples. The com-
cw=0.0
Training Reward
bined weight (cw) corresponds to in the mix- ture + (1 A)w which scales the log-
0.2 0.1
cw=0.2
cw=0.5
likelihood term in wdl. Training on negative sam- ples alone (cw= 0.0) yields pronounced dete- rioration in performance relative to our default setting (cw= 0.5). The results reinforce our argu- ment that a balanced contribution of positive and negative weights is most effective. In the absence of positive samples, the reinforcement-leuuee learennn signal collapses and optimisation becomes largely ineffective. A large emphasis positive samples (cw= 0.8) causes performance to deteriorate more rapidly, highlighting the critical role of negative samples in weighted ligelioed methods.
cw=0.8
0.0
1000
2000
3000
4000
Training Steps
Figure Training rewards of wd/ under differ- ent combined weights on Sudoku.
6 RELATED WORK
RL for Diffusion-baseed LLM. RL for discrete diffusion models has been explored through several approaches. One line of work, exemplified by DRAKES (Wang et al. 2024), leverages reward back- propagation along the denoising trajectory. This approach requires computing critic and propagating gradients through each denoising step, which is computationally intensive and prone vanishing gradients. Alternatively, methods such as MMaDA (Yang et al. (2025) and dl (Zhao et al. 2025) adopt direct RL formulations like GRPO, approximating missing diffusion components such per-token likelihoods- for policy optimization. Zhu et al. (2025a) applies Direct Preference Optimization (DPO) to fine-tune the LLaDA base model (Nie al. | (022a)), achieving notable gains in reasoning tasks. However, these approaches all depend on likelihood ratios, which can introduce bias and instability due to likelihood approximation erros. In contrast, method derives weighted policy optimiza-

---

Published as a conference paper at ICLR 2026
tion approach that eliminates the need for explicit policy ratios. Importantly, similar to prior works, our method directly optimizes the predictive distribution over clean data. A complementary line of research formulates policy optimization in terms of concrete scores (Lou et al.. 2024 Meng et al.) 2022). SEPO (Zekri & Boullé 2025), for instance, introduces a policy optimization objective that only depends on concrete score estimation, thereby circumventing likelihood approximation altogether.
RL for AR Models. The connection between GRPO and weighted regression has recently been explored in the context of RL with verifier reward (Mroueh) 2025), where binary rewards simplify policy optimization into likelihood-basee
objectives. Other closely related approaches are Rejection
Sampling Fine-Tuning (RAFT), which maximizes the likelihood of positive-reward samples (Xiong et al. 2025). Extensions of this idea incorporate negative samples to actively penalize the likelihood of negativerreward completions while enhancing that of high-reward ones (Zhu et al., 2025b Chen et al., (025). Other works introduce negative penalization through contrastive methods, such as Noise Contrastive Estimation (NCE) (Gutmann & Hyvärinen 2012 van den Oord et al.) 2019 Chen
ct al. 20244. Beyond binary rewards, preference-based learning has been widely studied using the
Bradley-Trryy model (Bradley & Terry] 1952; Ouyang et al. 2022 Rafailov et al., 2024; [Azar et al.)
2023; Ethayarajh et al. [024]] Wang al. 2023 Hong al., 2024p. In contrast to these approaches, our method accommodates general reward signals and can be interpreted as a form of soft rejection sampling, enabling efficient and stable policy optimization for dLLLMs.
RL via Weighted Regression. RL via weighted regression has been explored in earlier works advantagg--eiihtter regression (AWR) (Peng et al.] [019 Peters et al. al. 20TO)) and more recently in the
context of continuous control with diffusion policies (Ding et al. 2024 Zhang et al. (025). Weighted
likelihood-basse approaches have also been proposed for tine-tuning autoregressive (AR) language
models using general reward functions (Du et al. 2025 Baheti et al. 2024 Zhu et al.) 2023). However,
for AR models. where likelihoods are tractable. the necessity of such approaches remains unclear. In contrast, dLLMs suffer from intractable likelihoods, making weighted likelihood formulations particularly advantageous by reducing the number of required likelihood approximations. As such, RL via weighted likelihood provides a natural and efficient fit for optimizing dLLMMs In addition, we demonstrate in ablation study that merely optimizing policy with AWR (wdl-P) is ineffective.
FReee Policy Optimization. If policy optimization objective requires neither importance sampling nor regularization with respect to a reference model, then the objective ratio-free. Consequently, on-policy algorithms such as vanilla policy gradient methods (e.g., REINFORCE (Williams 1992)) and their variants (e.g.. RLOO (Kool et al.) (019)) are inherently ratio-free fatioffree This property is particularly valuable for dILMss, where errors in likeliielhood approximation can accumulate and propagate through ratio- computations. Concurrent work, such as SPG (Wang) et al. 2025a), adopts policy-gradien formulation and develops objective tailored specifically for diffusion language models. Another on on-policy optimization approach, d2 (Wang et al.] (022c)). removes both the ratios and the likelihood terms from the RL objective for dLLMs, offering more fundamental solution. However, our method wdl, similar to AWR (Peng et al. 2019). is inherently an off- policy loss, which is more general.
7 CONCLUSION
We introduce wdl, a weighted policy optimization method for reasoning with dLLMs. wdl is designed to minimize reliance on likelihood approximation, thereby mitigating the potentially substantial bias that can arise from approximation errors in policy ratios. Our method is grounded in weighted log- likelihood objective, derived to approximate the closed-form solution to the reverse-KL Woorss-KL-coostrainne policy optimization. Empirically, we show that wdl, even without supervised fine-tuning, surpasses the existing method d/ by up to 16% in accuracy on reasoning benchmarks, while also delivering notable improvements in computational efficiency during RL training. These results highlight the effectiveness of wdl and establish it as a more scalable and efficient approach for fine-tuning tuning dLLMs.
8 ETHICS AND REPRODUCIBILITY STATEMENT
This work raises no question or concern regarding the Code of Ethics. As for reproducibility of our results, we provide details of implementations in Section[5] in Experimental Setup and Implementation subsections. Additional details including dataset, reward functions, and hyperparameters are provided in Appendix B All the theoretical results are proved in Appendix A
10

---

Published as a conference paper at ICLR 2026
9 ACKNOWLEDGMENTS
Iijaa Bogunovic was supported by the EPSRC New Investigator Award PPPX0391XX/L the Engi- neering and Physical Sciences Research Council GP/S021566/1. Sangwoong Yoon was supported by the Institute of Information & Communications Technology Planning & Evaluation (IIP) grant I1201366, Artificial Intelligence Graduate
funded by the Korea government (MSIT) (No. School Program (UNIST)), the National Research Foundation of Korea(NRF) grant funded by the Korea MoV)) (No. RS-2024-004080033, and the Center for Advanced Computation at Korea Institute for Advanced Study. Xiaohang Tang was supported by the Engineering and Physical Sciences Research Council [grant number EP/T517793/1, EP/W524335/1J. Rares Dolga is supported by EPSRC, grant reference number EP/S021566/1.
The authors would like to thank UlPath and Che Liu (Imperial College London) for providing computing resources that supported our experiments, and Prof. David Barber, Yiming Yang, Xiaoyuan Cheng, and Keyue Jiang from University College London for valuable discussions during the early stages of this work.
REFERENCES
Silas Alberti, Kenan Hasanaliyev, Manav Shah, and Stefano Ermon. Data unlearning in diffusion models. arXiv preprint arXiv:2503.01034, 2025.
Arel. Arel's sudoku generator. html 2025. Accessed: 2025-04-08.
Marianne Arriola, Aaron Gokaslan, Justin T Chiu, Zhihan Yang, Zhixuan Qi, Jiaqi Han, Sub- ham Sekhar Sahoo, and Volodymyr Kuleshov. Block diffusion: Interpolating between autoregres- sive and diffusion language models. arXiv preprint arXiv:2503.09573. 2025.
Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured Denoising Diffusion Models in Discrete State-Spaces, 2023. URL https:://axii.crrr
abs/2100 03006
Mohammad Gheshlaghi Azar, Mark Rowland, Bilal Piot, Daniel Guo, Daniele Calandriello, Michal Valko, and Rémi Munos. A General Theoretical Paradigm to Understand Learning from Human Preferences, 2023. URL httppp///.a.x....rrr///o//////// abtp////ab//331100 12036
Ashutosh Baheti, Ximing Lu, Faeze Brahman, Ronan Le Bras, Maarten Sap, and Mark Riedl. Leftover Based Lunch: Advantage Based Offline Reinforcement Learning for Language Models, 2024. URL
https://arxiv http://abbb..../..//22/
Boris Belousov and Jan Peters. f-divergence constrained policy improvement. arXiv preprint
arXiv:1801.00056, 2017.
Ralph Allan Bradlcy and Milton E Terry. Rank Analysis of Incomplete Block Designs: I. The Method of Paired Comparisons. Biometrika, 39(314):324-3455 1952.
Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet. A Continuous Time Framework for Discrete Denoising Models. Advances in Neural Information Processing Systems, 35:28266-2822 2022.
Huayu Chen, Guande He, Lifan Yuan, Ganqu Cui, Hang Su, and Jun Zhu. Noise Contrastive Alignment of Language Models with Explicit Rewards, 2024. URL htttp:://axii..ort
2ttp:/242.0055666
Huayu Chen, Kaiwen Zheng, Qinsheng Zhang, Ganqu Cui, Yin Cui, Haotian Ye, Tsung- Yi Lin, Ming-Yu Yu Liu, Jun Zhu, and Haoxiang Wang. Bridging Supervised Learning and Reinforcement Learning in Math Reasoning. arXiv preprint arXiv:2505. 18//6, 2025.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reñchiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 10. 14168, 2021.
11

---

Published as a conference paper at ICLR 2026
Shutong Ding, Ke Hu, Zhenhao Zhang, Kan Ren, Weinan Zhang, Jingyi Yu, Jingya Wang, and Ye Shi. Diffusion-Bian Based Reinforcement Learning via Q- Weighted Variational Policy Optimization. arXiv preprint 4rXiv:2405.16173, 2024.
Yuhao Du, Zhuo Li, Pengyu Cheng, Zhihong Chen, Yuejiao Xie, Xiang Wan, and Anningzhe Gao. Simplify RLHF as Reward-Weighted SFT: A Variational Method, 2025. URL https
//axxiv.rgg/abs//// 11026
Kawin Ethayarajh, Winnie Xu, Niklas Muennighoff, Dan Jurafsky, and Douwe Kiela. KTO: Model Alignment as Prospect Theoretic Optimization, 2024. URL https: org. abs/
2402. 01306
Hugging Face. Open rl: A fully open reproduction of deepseek-rl, 2025.
Aditya Golatkar, Alessandro Achille, and Stefano Soatto. Eternal sunshine of the spotless net: Selective forgetting in deep networks. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 2304-9312, 2020.
Michael U. Gutmann and Aapo Hyvärinen. Noise-Contrand Contrastive Estimation of Unnormalized Statistical Models, with Applications Natural Image Statistics. Journal of Machine Learning Research, 13 (11):307-361, 2012. URL http://mmwrrrrrrrreeeeeeeeeeeeee
3ttt///gggttnnnnnnee 2a 2t.ttml...mm.mmmmmmm.mm.mm.mm...
Jiaqi Han, Austin Wang, Minkai Xu, Wenda Chu, Meihua Dang, Yisong Yue, and Stefano Ermon. Dis- crete diffusion trajectory alignment via stepwise decomposition. arXiv preprint arXiv:2507.04892. 2025.
Haoyu He, Katrin Renz, Yong Cao, and Andreas Geiger. Mdpo: Overcoming the divide of masked diffusion language models. arXiv preprint arXiv:2508. 13148, 2025.
Jiwoo Hong, Noah Lee, and James Thome. ORPO: Monolithic Preference Optimization without
Reference Model, 2024.
07691
Hugging Face. Open rl: A fully open reproduction of deepscek-rl, January 2025. URL https wetthub.com//ugggga.i.
open. open-rll
Sham Kakade and John Langford. Approximately Optimal Approximate Reinforcement Learning. In
Proceedings of the nineteenth international conference on machine learning, pp. 267-274, 2002.
Wouter Kool, Herke van Hoof, and Max Welling. Buy 4 reinforce samples, get baseline for free! 2019.
Hynek Kydlíček. Math-Verify: Math Verification Library. URL https://githube com/ hhggiingface//natt-vvrrri
Inception Labs, Samar Khanna, Siddhant Kharbanda, Shufan Li, Harshit Varma, Eric Wang, Sawyer Birnaaum, Ziyang Luo, Yanis Miraoui, Akash Palrecha, Stefano Ermon, Aditya Grover, and Volodymyr Kuleshov. Mercury: Ultra-Fast I Language Models Based on Diffusion, 2025. URL
https: Settps:/eerxvv. wtt.://ggg/ab.a.oooooo 17298
Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. In The Twelfth International Conference on Learning Representations, 2023.
Zichen Liu, Changyu Chen, Wenjun Li, Penghui Qi, Tianyu Pang, Chao Du Wee Sun Lee, and Min Lin. Understanding R1-Zero-Like Training: A Critical Perspective. arXiv preprint arXiv:2503.20783, 2025.
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution, 2024.
/ttp:////b///31...aa 2abbb//333118.4 16834
Cheng Lu. Huayu Chen, Jianfei Chen, Hang Su, Chongxuan Li, and Jun Zhu. Contrastive Energy Prediction for Exact Energy-Gute Guided Diffusion Sampling in Offline Reinforcement Learning. In International Conference Machine Learning, pp. 28825-22855. PMLR, 2023.
12

---

Published as a conference paper at ICLR 2026
Chenlin Meng, Kristy Choi, Jiaming Song, and Stefano Ermon. Concrete Score Matching: General- ized Score Matching for Discrete Data. Advances in Neural Information Processing Systems, 35: 14532-34545, 2022.
Nikita Mounier and Parsa Idchpour. Review, remask, refine (r3): Process-guided block diffusion for text generation. arXiv preprint arXiv:2507.08011 2025.
Youssef Mroueh. Reinforcement Learning with Verifiable Rewards: GRPO's Effective Loss, Dynam- ics, and Success Amplification. arXiv preprint arXiv:2503.06633 2025.
Niklas Muennighoff, Zitong Yang. Weijia Shi, Xiang Lisa Li, Li Fei-Fei, Hannanch Hajishirzi, Luke Zettlemoyer, Perey Liang, Emmanuel Candès, and Tatsunori Hashimoto. s1: Simple test-time scaling, 2025. URL https:
haaprrg/bb///00000 2501. 19393
Shen Nie, Fengqi Zhu, Chao Du, Tianyu Pang, Qian Liu, Guangtao Zeng, Min Lin, and Chongxuan Li. Scaling Up Masked Diffusion Models on Text. arXiv preprint arXiv:2410. 100::2400./81444 /8514, 2024.
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. arXiv preprint arXiv:2502.0999 2025a.
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, 2025b. URL https
//arxiv. http://////bb//////// 09992
Jingyang Ou, Shen Nic. Kaiwen Xue, Fengqi Zhu, Jiacheng Sun. Zhenguo Li, and Chongxuan Li. Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data, 2025. URL https: /ettps://azxe //tp:vv...ivvvv.aa org. abs /2066 03736
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, al. Training Language Models Follow Instructions with Human Feedback. Advances in neural information processing systems, 35: 27730 27744, 2022.
Jiay Pan, Junjie Zhang, Xingyao Wang, Lifan Yuan, Hao Peng, and Alane Suhr. Tinyzero. http://ggttbbbcom/lllaaaaaaa/// Pan/TinyZero, 2025. Accessed: 2025-01-24.
Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage-Weigee Regression: Simple and Scalable Off-Policy Reinforcement Learning. arXiv preprint arXiv:1910.00177,, 2019.
Jan Peters, Katharina Mulling, and Yasemin Altun. Relative Entropy Policy Search. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 24, pp. 1607-1612, 2010.
David Pollard. Asymptopia: An Exposition of Statistical Asymptotic Theory. URL http://www. stat. yale. edurpollard/Bookk/AAyyppppop 2000.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct Preference Optimization: Your Language Model is Secretly a Reward Model. Advances in Neural Information Processing Systems, (1.53728-537411 2023.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Stefano Ermon, Christopher D. Manning, and Chelsea Finn. Direct Preference Optimization: Your Language Model is Secretly Reward Model, 2024.
URL httlltppss//axxi.orr////////////
18290
Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. Deepspeed: System opti- mizations enable training deep learning models with over 100 billion parameters. In Proceedings of the 26th ACM SIGKDD international conference knowledge discovery & data mining, pp. 3505-3506, 2020.
Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin T Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and Effective Masked Diffusion Lan- guage Models, 2024. URL RRl..httss///rrctt http://awwww..oor/////////////
07524
13

---

Published as a conference paper at ICLR 2026
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust Region Policy Optimization. In International conference on machine learning, pp. (1889-197) PMLR, 2015.
Zhihong Shao, Peiyi Wang. Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Y Wu, et al. Deepseekmath: Pushing the Limits of Mathematical Reasoning in Open Language Models. arXiv preprint arXiv:2402.03300. 2024.
Jiaxin Shi, Kehang Han, Zhe Wang, Amaud Doucet, and Michalis K. Titsias. Simplified and Generalized Masked Diffusion for Discrete Data, 2025. URL htttp:://axiv..rrrr wwtpaawwww..i../////op 2406. 04329
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation Learning with Contrastive Predictive Coding, 2019.
03748
Chaoqi Wang, Yibo Jiang, Chenghao Yang. Han Liu, and Yuxin Chen. Beyond Reverse KL: Generalizing Direct Preference Optimization with Diverse Divergence Constraints, 2023. URL
pictps://arxiae org http///bs///ee
httrg/ab//2309... 2309. 16240
Chengyu Wang. Paria Rashidinejad, DiJia Su, Song Jiang, Sid Wang. Siyan Zhao, Cai Zhou, Shan- non Zejiang Shen, Feiyu Chen, Tommi Jaakkola, et al. Spg: Sandwiched policy gradient for masked diffusion language models. arXiv preprint arXiv:2510.09441 2025a.
Chenyu Wang. Masatoshi Uehara, Yichun He, Amy Wang, Tommaso Biancalani, Avantika Lal, Tommi Jaakkola, Sergey Levine, Hanchen Wang, and Aviv Regev. Fine-Tuning Discrete Diffusion Models via Reward Optimization with Applications to DNA and Protein Design. arXiv preprint arXiv:2410.13643 2024.
Guanghan Wang, Yair Schiff, Subham Sekhar Sahoo, and Volodymyr Kuleshov. Remasking discrete diffusion models with inference-time scaling. arXiv preprint arXiv:2503 2500000000000007 2025b.
Guanghan Wang, Yair Schiff, Gilad Turok, and Volodymyr Kuleshov. d2: Improved techniques for training reasoning diffusion language models. arXiv preprint arXiv:2509.21474 2025c.
Wen Wang, Bozhen Fang, Chenchen Jing, Yongliang Shen, Yangyi Shen, Qiuyu Wang. Hao Ouyang. Hao Chen, and Chunhua Shen. Time is feature: Exploiting temporal dynamics in diffusion language models. arXiv preprint MxXiv:2508.09138 2025d.
Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 8(3):229-256, 1992.
Wei Xiong, Jiarui Yao, Yuhui Xu, Bo Pang, Lei Wang, Doyen Sahoo, Junnan Li, Nan Jiang, Tong Zhang, Caiming Xiong, and Hanze Dong. A Minimalist Approach to LLM Reasoning: From Rejection Sampling to Reinforce, 2025. URL httpttt:///aii..oottt org abs /tt..//50......
Ling Yang, Ye Tian, Bowen Li, Xinchen Zhang, Ke Shen, Yunhai Tong, and Mengdi Wang. Mmada: Multimodal Large Diffusion Language Models. arXiv preprint arXiv:2505.. 15809, 2025.
Jiacheng Ye, Zhihui Xie, Lin Zheng, Jiahui Gao, Zirui Wu, Xin Jiang, Zhenguo Li, and Lingpeng htt.:/uuuun........s
Kong, Dream 7B, 2025. URL
gwt.ubtthhu.... wttpio////////oo/... 2025 /deen
Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yu Yue, Weinan Dai, Tiantian Fan, Gaohong Liu, Lingjun Liu, et al. Dapo: An open-source IIm reinforcement learning system at scale. arXiv preprint arXiv:2503. 14476, 2025.
Oussama Zekri and Nicolas Boullé. Fine-Tuning Discrete Diffusion Models with Policy Gradient Methods. arXiv preprint arXiv:2502.01384 2025.
Huaye Zeng, Dongfu Jiang, Haozhe Wang, Ping Nie, Xiaotong Chen, and Wenhu Chen. Acecoder: Acing coder rl via automated test-case synthesis. arXiv preprint arXiv:2502.01718 2025.
Shiyuan Zhang, Weitong Zhang, and Quanquan Gu. Energy-Weighted Flow Matching for Offline Reinforcement Learning. arXiv preprint 2503 2025.
14

---

Published as a conference paper at ICLR 2026
Siyan Zhao, Devaansh Gupta, Qinqing Zheng, and Aditya Grover. d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learning. arXiv preprint arXiv:2504. 2025.
Banghua Zhu, Hiteshi Sharma, Felipe Vieira Frujeri, Shi Dong, Chenguang Zhu, Michael I Jordan, and Jiantao Jiao. Fine-Tuning Language Models with Advantage-Induced Policy Alignment. arXiv preprint Sttv:2330.0223:: 2023.
Fengqi Zhu, Rongzhen Wang, Shen Nie, Xiaolu Zhang, Chunwei Wu, Jun Hu, Jun Zhou, Jianfei Chen, Yankai Lin, Ji-Rong Wen, al. LLaDA 1.5: Variance- Reduued Preference Optimization for Large Language Diffusion Models. arXiv preprint arXiv:2505. 19223, 2025a.
Xinyu Zhu, Mengzhou Xia, Zhepei Wei, Wei-Lin Chen, Danqi Chen, and Yu Meng. The Surprising Effectiveness of Negative Reinforcement in LLM Reasoning. arXiv preprint 276. 2025b.
15

---

Published as a conference paper at ICLR 2026
CONTENTS
Introduction
2
Preliminaries
2.1
Diffusion Large Language Models
2.2 Existing Policy Optimization Methods
2.3
Policy Optimization for dLLMss
3 wdl: Weighted Policy Optimization for dLLMs
3.1
Reinforcement Learning as Weighted Log-Likeliben Maximization
3.2 wdl: Fully Utilizing Completions
3.3 wd/++: Stepwise Weighted Policy Optimization
4 Theoretical Insights: Energy-Guided Diffusion Sampling
5 Experiments
5.1 Main Results
5.2 Ablation Study
Related Work
7
Conclusion
10
8 Ethics and Reproducibility Statement
10
9
Acknowledgments
11
Proofs and Additional Theory
18
A.I Objective Estimation Error due to Likelihood Approximation
18
A.2 Reinforcement Learning
18
A.3 Masked Discrete Diffusion
20
Additional Experiment Setup Details
23
B.1
Dataset, Training and Evaluation Protocol
23
B.2
Reward Function
24
B.3
Sampling from Geometric Mixture
24
B.4
Hyperparameters
24
B.5
Training Cost Estimation
25
B.6
Computing Resources
25
Additional Experiments
26
C.1
Summary of wd/ Results
26
16

---

Published as a conference paper at ICLR 2026
C.2 Additional Ablation Study
26
C,3 Coding Benchmarks
28
C,4 Training Dynamics
28
Limitations
28
D.I Additional Analysis on Unlearning
29
17

---

Published as a conference paper at ICLR 2026
A PROOFS AND ADDITIONAL THEORY
A. OBJECTIVE ESTIMATION ERROR DUE TO LIKELIHOOD APPROXIMATION
In this section, we aim show that diffu-GRPO amplity the log- likeliihohood approximation error. De- note the approximator by e such that | (g. o) (e q) S and | (g, o (olg() s c.
Then the objective diffu-GRPO in the worst case suffers from exponential error. We discuss the case without ratio clipping and omit the regularization for convenience. Denote LORFO the ground truth objective without likelihood approximation:
[ainiGGRPO LGRPOIl
[oc|
exp 6 (ot)/ exp o am (o.)) A.]
= D... a...D.a.-nmutlll
10:
- ED on ama
(To (TO(O))) Told (o)) 11
[o.1
E BgDD o~
exp (0 6: (of Ail
6
- -Ey on (0))
1
Too (oaa)
matal (o,)) As]
1=1
[oul
B om
[o.]
exp
llog no (o,)) - log man(())) + + ())
1=1
G
(/ooa// A.] II
-
0
0
1=1
lo.1
-1 -I costa
exp (e + e e Cexp (E (e+ (e)))).
(15)
0
where C is constant independent and e'. In contrast wdl has only linear approximation error. Denote the objective computed using approximated log- likelihood as Ls
[| - Lundl
Pot 2)
Σ
(g,0;)) + w -0 o:)) (ø(9. o.) log no (o: (4) II
SC'c.
(16)
A.2 REINFORCEMENT LEARNING
Reinforcement Learning Formulation. We first introduce the reinforcement learning notations and then extend it to the setting of LLM post-training, Denote 7 as a trajectory (F = (so. ao. $1....) ) ~ T) sampled following policy л. Specifically, so ~ pl a, (--(s)). st+1 P(-lqe at). The objective of Reinforcement Learning aims to find policy π,, which maximizes discounted total return,
(())) Erax
(c,
t=0
Let the discounted return of a trajectory be R(T) = Σ r 0z, 81+1). The advantage function is defined as AT a) = (=(ssa)-V== (V*()))))(()).. where V* Ey [ = s] is state value function, s, 00 = al is state-action value function. Denote Pruu as the
and Q* a) E,~- [R{T)]s0 marginal state distribution. Denote the total variation of two discrete probability distributions a. by Dry (a, b) := Σ, la; b, and Dry(a, b)2 (a || b) (Pollard 2000 Schulman et al. 2015 When and b are conditional probability distribution, denote D a, max, Drv(a ( and Dmax (a b) max, (a. (((()))
18

---

Published as a conference paper at ICLR 2026
We then extend RL for LLM post-training. In this paper only consider the sequence-level reward and loss objective, so we directly replace s with and a with completion o. Then the horizon of the RL for post-training becomes only 1. The following theorem provides a monotonic (non-decreasing) guarantee of existing prevailing RI methods.
Proposition (Policy Improvement Bound (Kakade & Langford 2002
Schulman al. A (s, a) and C'
2015)). Let surrogate objective L (T)) = (T(oll)) + Esmppa ()) a~x(11a))
4 amx (A* -7)2 7)2, then vk N:
nT(=") 2 Lixas (T=" - CDDVV (Tolds x x²)2.
Remark 3. Based Proposition [ due Div (a] b)2 I DRt (Pollard 2000 Schul
man et al. 2015). TRPO and PPO with fixed forward KL. regularization have the monotonic im-
provement guarantees. In other words, n(xx)))))))))))( (n*) CDD.pxxTTaa x )² z I (x²)
CTE(DaL (TTiiaaTTllll [x]] 2 (Fou)) = (TToi))
Proposition 2. Minimizing (x (+1g) (-(-1))) w.r.t. 0 is equivalent to optimize the following loss objective:
LWLL =B q~D Do
exp (NAAT o)) log mo
(17)
= - E (o: (o:)
exp
(v Tal (g. 0.))
log To (a(()))).
(18)
eexp ( Tu (q, o;))]
Proof. To obtain the practical objective in Equation and obtain the following. Vgg E D:
18). we first start from the cross-entropy loss,
DKL ( (:12) xo(1lg))
Fown (:19) log to (olg)
(19)
Σ x (olg) log (o
(20)
Tod exp (viAT
Σ Eor (o' lllll))(eeee (exp ( (q, o')]
log (
(21)
exp A Fmu (q, 0))
-E Eawrr (olg)
log [o()))
(22)
E wyr [exp ( Ame Taue (g. o'))l
Since the normalization constant Egg-ern (exp A STal (g, o') d')] is independent to o, we convert the objective to a weighted log- likelihooddd and approximate with samples from the group and weight normalization to obtain:
(0) = E. Eadry((o))) exp ( Rou (q, o)) log To((
(23)
G exp (1 AT (g.
~ - (
G E lexp (A) (g, o;))] log To To(0)1)).
(24)
We derive Equation (1 from Equation (0) by simply using the known form of the optimal policy
Tt (-1q) x ( exp (vhA(g, We derive Equation (22) from Equation (21) by using the
definition of expectation and from Equation (22) to Equation (44 by approximating through G
samples 1G ~ mi (o)q).
Theorem 2. Keverse-KL-reguuande Policy Optimization defined in the following objective has monotonic improvement guarantees. Specifically, denote regularized objective T Tς())
EgeD [3Dxa (T(:1q) [1)((19))] and denote
M(T) L(T) EgED ( (oa 3Dx (7((19) (DDx((((()))))) (m)((
(25)
19

---

Published as a conference paper at ICLR 2026
where L(T) (TTToa)) + Br own(110) (A (q. Let H* be the solution the objective maxo (TTo):
0" arg max EEED, (omDDoo AA 0) ADRI TT(-14) Tata(-1a) -3Drr, (Ta(-11) u) (Tnr((19))
(26)
then 1/ (T*) 1/ (Tola).
Proof. Based on Proposition ! we have
11(x") =n(=") - EyED 5DRL ( (-1()) -1a(
21(x') - C CD (Tola,
(Tola, )²²² RyED [DK.. (+1(-) mer((1a))
(27)
21(x') CD (TT (Toaa)))a) EDED [DKL (+19) (
(28)
21(x") - EgED AD ((10)((Taa(((1)))) 3DKL (+1g) m
(29)
=M(x")
(30)
(Tooia
(31)
=((rol) - E
(-14) + 3KI.. (Toa(-1l)) mee((11a)
(32)
|3DRL (-19 || Teer( (4)) }]
2((ool)
EyEDD
(33)
(m)(aaaa))
E
3KK..
(Toua (:4) mer((11)
(34)
=1y (Tota)
(35)
Equation 27 holds due to Proposition] Equation 28 holds due to (pllg) < DRi. (pll4) (Pol
lard 2000). Equation (29) holds due to the definition Dmax Equation (30p is according to the defini-
tion of M.)) The key iequality Equation (1) holds since xt" is the maximizer of function L(-). Equa-
(4) Equation (34 holds since L( ()(Tola) + E.yDD a~) [ATH( o)] (TTTlla). Equation 35
tion (32h holds due the definition of M(). Equation (33) holds since DRI.
from the definition of 1/'.
A.3 MASKED DISCRETE DIFFUSION
In this section, we show how our objective learns distribution for which all marginals at time t satisfy intermediate energy guidance as per Lu al. (2023).
Definition 2. The absorbing transition kernel is defined as Qt = (1) Qaaoobo where
0
0
1
Qaburrb
-1
0
0
Definition 3 (Concrete Score). Denote 2t (x),... x4) and &, is identical to æı except the i-th
token is unmasked (i.e. x² = [M] and 27 + [M]. Concrete score is defined as the marginal probability ratio between it and zn:
s(xe, 1) def p(
(36)
p(
Proposition (Marginal Distribution (Ou al. 2025 Denote {ar} as a continuous time Markov chain with transition matrix Qu (l)) Quosoro Assume d tokens in x. = (at. x44) are masked
tokens [M), and d2 d, tokens are unmasked, the marginal distribution pe (x:) satisfies
(((x)) -at))- (e
do -ot)) o 2 po (M)
(37)
where ā (()) So (s)dss and UM is the set of unmasked tokens in Tt-
20

---

Published as a conference paper at ICLR 2026
The following theorem provides the foundation of directly modeling the clean data distribution.
Proposition 4 (Analytic Concrete Score (Ou al. 2025 ) Denote 3tz = (z),..., xt) and & is identical to except the i-th token is unmasked (i.e. x [M] and Éi [M]). Then the concrete score at time can be expressed by the conditional probability of predicting this unmasked token.
E [(1)
1 - -0) Po x(
Pr a
The marginal probability distribution of the masked responses (2x) in the diffusion process satisfies (xx)) = pl((t)) exp (Ae (z<()))) /Z, which induces an energy-guide discrete diffusion:
Lemma [1.
(oje(Fo((e)) xx)) (xp(A( -((z)))))A(()))))
(38)
where intermediate energy function is defined as A((2o) log Bugrrro (1zz) (lx))eepp (A(zo)) for and Ao(zo) = A(zo), A(-) is advantage function, Ze is the normalization constant.
Proof. The theorem and proof mainly extend from theory developed in continuous setting (Lu et al. 2023 According to the marginal likelihood of clean data distribution (, = (6()))) 2 and
identical forward process, can rewrite the marginal likelihood of masked data:
(xx) f ( (2006(0))) dro /
/
e e*A((=))
dæo
7
eVA((2)) Z dzg.
Applying Bayesian rule we know that péjo (zz) = (00 hence we can further rewrite
pt (xx) = / (20
(vA(m))
dr = p1(2)) / Poicloolze)
cuA((0))
Z dro
Z
(2,, Egg, (mo) [en)((oo)
(xx) ettre(z))
Z
Z,
Therefore, the marginal likelihood of masked sequence satisfies: Př (zx) px())) exp (At(zz)) ZZe. Since Přjo = Pojo- based on the marginal likelihood clean data distribution satisfies Pī =
Pb(xo) gd(so) we can further applying Bayesian rule to obtain the energy-guided discrete diffusion model:
Přjo (1o (z0)
Pojz [x))
(39)
pi (xt)
A(zo) Pho 1o (0)o)(((())))) (2() 7
(40)
pf (Ite
P)oo(x:220)0)) (x()) p/((2e) (v()((2))
A(oo)
(41)
x Po)o(2o(11) exp(A(xo) As(x:)),
(42)
Lemma 2. According to Definition 1 due to the identical forward process
() = Plu(xxl x Paju (x)((=)))) based Lemmal 7 and Proposition we have the marginal probability of the unmasked tokens satisfies that for all step t,
(43)
P6(+L/19) (9) (2 LM (q)* >. ( 9) (eg (A(q,xo))1/4
(44)
where 7 is the normalization constant.
21

---

(g) Published as a conference paper at ICLR 2026
Proof. According to the identical forward distribution of three diffusion process (new, old, and reference), based on Equation 37,, we have Vt:
(xa q) [1 - [ [0-0 po( UM
(45)
(r (4) = e - di [e- da púl ( 1
(46)
(ma = - e -- di (e- da po (xt po
(47)
Then rewrite Equation 46 in the residual energy- based form defined in Equation 38
we have
--e ( di
(e (0-0 p pl (#rlg) (x((1)))) exp (Ae /Z. /Z.
(48)
By plugging (xr (g) pe (x) A (xrlg) and Equation 45 and Equation (47) into Equa-
tion (48), we have that the clean data distribution of the unmask tokens at diffusion time / satisfies:
p ( (q) =po (p > pb LM 1)) exp aA)
(
(49)
=pu( (M) 1g) Po (z ()) Ee (A(q, xo))l..
(50)
Proposition 5.
The marginal likelihood of the target diffusion model p" satisfies Equation 38
Consequently, the concrete score of the target diffusion model, denoted by s", can be expressed by the score of the mixture diffusion p' and the posterior
of the advantage:
(o (da) lexp (A(x)))//2
80x,,, 1) = (2te, 1)
(51)
E Eng (m)la)))) eep (A(zo) A(xx)))//Z
and equivalently
po(i)|2t UM ,9)² po (a:laa (g)3 Pegcoole (A(g, mo))
(52)
B lexp (A(g, zo 20)
Proof. According to Lemmal 2
p #t1g)
( #;1g)>
lexp xo))11
(53)
( q)) Po(xtM (q)a
pe (x((x UM 14)
Eng mx) lexp (A(g, zo))]/2
E (z lexp (A(g,2o))1/
1t .g) =pa(ii b UM ,4)*
po (#t at UM .)
(54)
E (mol) eepp (A(g,ro))]1
Both sides in Equation 52 multiply C(t)
I-gnte and based the analytic form of concrete UM
score introduced in Proposition C(f) Po (6f (ell .g) s( t). Thus, we have
Eppc (.) jexp (A(g, zo)) zo))
C(I) (0 UM .q) C(I) po .g)> p (1/ .
Ey (no (exp (A(g. mo))
(55)
s" (II, = =s((t,)) (pg() /exp (A(q, 2o))
(56)
(oo( /exp (A(g. zu zu))
Lemma 3. The normalization constant Z, = En, D',1((z (x+lg) Beole, (oxp (o)) is independent to the masked response 2t, In other words, Z == Ez, (( Enola, [exp Alg, æo)] for any St It-
22

---

Published as a conference paper at ICLR 2026
Proof.
Z
Σ pl (xelg) Ezol A(g, zo)] = Σ (x(()))) Ex'(zolz)) exp A(q.
(57)
Σ Σ p' (2o. xzl4) exp A(q. xo) = Σ> Σ p' (2o, 2z q) exp x()
(58)
2g
2o
Σ (2oog) exp A(q. 2o)
(59)
Thus Z becomes independent to It, leading to that
Z = 2 = Σ
(felg) Bzol2, lexp a0)
(60)
21
Theorem [ The score model se S defined in Equation (5) is satisfied when the following loss objective is minimized. This objective a form of advantage-weighted Denoising Concrete Score Matching (D-CSM), which we call AW-D-CSM:
p6 (ro))
CAW.D.CSM EBp)(10) (exp (A(q, zo)) teo
1211.
(61)
p/ (rra (2))
Advuntage Weight
(occs((7o)
Proof. Denote so (Ftt 1) I-c aपy Po
[ (M is the concrete score model induced by pa- Accord-
ing Lemma 3 2 = Z. Then according to Proposition 5 Equation
52 is equivalent to
p ( UM .9) (x66( A(qg,2)))
( .g)> po (# .9)8. (xg(tols:))exp (A(g, 20)]
(62)
We aim to update ( 9) pū LM .4) to satisfy Equation 62 thus construet (a
loss function objective by replacing pd with po and construct a L norm loss
p6( E (xl( (A(g.2o)) (1 .9) (xoola)))
po((i a)>ppp ( 91)
UM
(63)
ph(on)(tt)
= =R te (A(g,
Ilpo (atla MMM .g)
16
P6(a/l=! UM .0)1
(64)
P6( (xn)
-I ( (Alc:pp(AAaa)))0 (A(g, mo)) Il (at .9)
(nellooo)p p6 (6 19911)
(65)
p
=E =Ep (A(q, 20)) (x..)) pb(11(20)) [22].
(66)
B ADDITIONAL EXPERIMENT SETUP DETAILS
B. DATASET, TRAINING AND EVALUATION PROTOCOL
As for well and d1, we reproduce dl running the official code without and change, and train
our method wdl/ evaluated for accuracy of the test datasets at steps 1000, 2500, 5000, 7500 both GSM8k and MATH; steps 1000, 2500, 4000, 5000, 12500 Sudoku: and 1000, 2500, 4000 in Countdown. We evaluate less checkpoints compared to dI. On the GSM8K, we train models on
reasoonning/// https::/gi o
23

---

Published as a conference paper at ICLR 2026
the train splif and evaluate on the test split. On Countdown, train the 3-number subset of the datase from TinyZero (Pan et al. 2025), and evaluate on 256 synthetic 3-number questions provided by Zhao et al. (2025). On Sudoku unique puzzles and evaluate on 256 synthetic ones provided by ZZha et al. (2025). On MATH500, we train models the train spli
the 4x4 datase generated by Arel (2025). We train 1M
To train wd/++ for evaluating on MATH500, we use dataset provided by (He et al. 2023). which is subsampled from OpenRl dataset Face (2025). To evaluate on GSM8K, leverage train split to conduct wd/++ training. Notably, we leverage a more effective system prompt and Math- Math-Verify (Kydliček) to parse the answers for full parameter fine-tuning of wdl, wdl++ and MDPO.
B.2 REWARD FUNCTION
To train wdl and reproduce di, we use the reward function defined (Zhao et al. 2025). For completion, we provide the details as following.
GSM8K. Following the Unsloth reward setur we apply five addtive components: XML Structure Reward: +0.125 per correct tag: small penalties for extra content post-tags. Soft Format Reward: +0.5 for matching the pattern <reasoning>,
< snswer>
Strict Format Reward: +0.5 for exact formatting with correct line breaks. Integer Answer Reward: +0.5 if the answer is a valid integer. Correctness Reward: +2.0 the answer matches ground truth.
Countdown. We include three cases: +1.0 if the expression reaches the target using the exact numbers. +0.1 if numbers are correct but target is missed. otherwise.
Sudoku. The reward is the fraction of correctly filled empty cells, focusing solving rather than copying.
MATH500. We include two additive subrewards. Format Reward is +1.00 for <answer> with \boxed inside; +0.75 for <answer>> without \boxed; +0.50 for \boxed only. +0.25 for neither. Correctness Reward: +2.0 if the correct answer is in Aboxed ().
To train wd/++, we leverage Math- Math-Verify (Kydlíček). constructing a simple verifier reward function evalaute on GSM8K and MATH500.
B.3 SAMPLING FROM GEOMETRIC MIXTURE
Although the sampling strategy eliminates the need to approximate the reference policy's likelihood, it incurs computational overhead, as generating full completion requires multiple forward passes through the dLLM- compard to a single pass for likelihood estimation. An alternative is to sample from Vold and shift the advantage to Å, = TM (q, o,) + log Tíe/ (A + 3), which reintroduces the need for reference policy likelihood approximation. However, policy ratio has been removed, and the reference model can be reused when conducting multiple gradient updates with the same batch of rollouts (off- coffcpolicy). The increased computational burden is slight.
B.4 HYPERPARAMETERS
We provide the hyperparameters of SFT in Table and for wdl in Table
bacth size
max length
learning rate
grad accum steps
Value
4096
le-5
Table 5: Hyperparameters of SFT in d/ reproduction.
24

---

Published as a conference paper at ICLR 2026
Parameter
wdl
Model and Precision
use _e__peft
true
true
torch torch_dlype
blloat16
bfloat16
load in 4bit
true
true
iit__implementaiinn
flash Baish_atteniin_22
bash_attention
lorar
128
128
lora_alpha
lora_drip lora_dropout
0.05
0.05
CAUSAL_LM
CAUSAL_
Training Configuration
seed
42
42
bf16
true
true
sync_ref_mode
True
True
[e__model_syne_eee
64
64
adam_betai
0.9
0.9
adam beta2
0.99
0.99
weight_deese decay
0.1
0.1
< (Equation 9
1.0
max_grad_norr norm
0.2
0.2
warmup_ratio
0.0001
0.0001
learning_rait rate
3e-6
3e-6
constant_ with- warmup
constant_ waamup
Batching and Evaluation
per ser_ddvice__reee train batch size
6
6
device eval batch, size
1
1
gradient_ hetpiiitnacccmmmattiiinssssss
2
2
RL
generations
6
6
max completion length
256
max prompt_lengee
200
200
block_length
32
32
diffusion_ster ceptratton_baahhhh.ee size
128
128
6
remasking
low oow_confidence
www_confieence
random_ maaioong
True
True
p_ p_maak_prompt
0.15
0.15
beta
0.00
0.04
epsilon
0.5
num iterations
12
12
Table 6: Comparison of hyperparameters between wdl and dl.
B.5 TRAINING COST ESTIMATION
For the runtime measurements reported in Table 2 we set μ and train for total of 6 global
steps, corresponding to 48 gradient update steps. We use a batch size of and the rest of the hyperparameters the same as in Table To estimate the number of function evaluations (NFEs) involved computing likelihood approximations, we count only the forward passes, as the number of backward passes remains consistent across methods. The additional NFEs observed in the d/ model arise from evaluating the likelihood under both the old and reference models, which are used for regularization. These extra evaluations are required only when new samples drawn, their outputs can be cached and reused across all gradient updates for μ. We additionally report the number of floating point operations (FLOPs) per global training step, measured using the Flops Profiler from Rasley et al. (2020).
B.6 COMPUTING RESOURCES
For both wdl and dl, RL training is conducted four NVIDIA A100 GPUs (80GB), and SFT performed on four A6000 GPUs (48GB). For wd/++ and MDPO, RL training is conducted on 8x A800 (80GB).
25

---

Published as a conference paper at ICLR 2026
C ADDITIONAL EXPERIMENTS
We additionally report results for comparison the results of the baseline d/ reported in the paper (Zhao et al. 2023). As shown in Table 7 our method wd/ evaluated and selected from less checkpoints, can outperform d/ with large margin in Sudoku and Countdown, achieving comparable performance in math problem-solving tasks.
C.I SUMMARY OF wdl RESULTS
Table 7: Test accuracy across different tasks. Our method demonstrates higher accuracy, especially significant in Sudoku and Countdown. The shaded area indicates where our method outperforms.
Sudoku
Countdown
GSM8K
MATH500
Model
256
512
256
512
256
512
256
512
LLLDA-SB-Instruct
6.7
5.5
19.5
16.0
76.7
78.2
32.4
36.2
+ diffu-GRPO (reported)
12.9
11.0
31.3
37.1
79.8
81.9
37.2
39.2
+ diffu-GRPPO (reproduced)
16.1
11.7
27.0
34.0
80.7
79.1
34.4
39.0
dl (reported)
16.7
9.5
32.0
42.2
81.1
82.1
38.6
40.2
dl (reproduced)
17.6
16.2
25.8
35.2
78.2
82.0
34.4
38.0
wdl
76.4
62.8
51.2
46.1
80.8
82.3
34.4
39.0
Table 8: Training cost. The training steps to obtained the best gost-trained model of three methods are 20, 150, and 7500. To compute the total rollouts, need compute the average rollouts single training step. Gradient steps per rollout batch represents the number of gradient descent conducted with single batch of rollouts. In other words, represents it pure on-policy RL training, and for any value >> 12, off-policy RL is executed. Total Batch Size is computed by multiplying per-device batch size, gradient accumulation and the number of gpus. Therefore, the average number of rollouts used for single step gradient descent should be computed by total batch size divided by gradient steps per rollout batch.
Hyperparameter
wd/++
MDPO
di
Training step of the best checkpoint
20
150
7500
Training Steps per Rollout batch
1
12
Per-Device Batch Size
4
6
Gradient Accumulation
2
16
2
GPUs used for training
8
4
Total Batch Size
64
128
48
Avg. Rollouts per Step Total Rollouts
64
128
4
1280
19200
30000
We additionally provide reward dynamics in comparison wd1-SFT in training. In Sudoku and Countdown, directly training with wdl without SFT shows significantly more efficient and stable learning process. GSM8k and MATH500, the difference negligible.
C.2 ADDITIONAL ABLATION STUDY
We provide additional ablation study on the combined weight to confirm analysis that the positive and negative samples terms the loss function should be assigned equal proportion, due to the side case of batch of all-negative generated responses (see the paragraph below Equation (9). Assigning equal proportions to positive and negative weights is not arbitrary but rather the most robust design. This can be understood through critical failure modes that arise from imbalanced proportions:
26

---

Published as a conference paper at ICLR 2026
Sudoku
Countdown
GSM8K
MATH500
Reward
0.2
wd w
0.1
d1
wd SFT
08
05
Steps
Steps
Steps
Steps
Figure 3: Reward Dynamics. wdl without SFT demonstrates better rewards in Sudoku and Count- down.
Table Ablation on weight to combine positive w and negative weights w in wdl on Sudoku. Specifically, the final weight assigned to logelikelihood is computed as -A + (1 A)w
Combined Weight
Accuracy
Effective Tokens
0.5
25.63%
326.97
0,4
11.77%
240.04
0.6
14.11%
220.13
When positive weight has larger proportion: In scenarios where all sampled completions have uniformly low rewards, larger proportion of positive weights would paradoxically increase the log-likelihood of negative samples during wdll optimization, which is clearly undesirable and contradicts the learning objective.
When negative weight has larger proportion: Conversely, when all generated completions achieve uniformly high rewards, insufficient proportion of positive weights would result in unlearning high-quality samples.
To empirically validate this analysis, we conducted experiments on the Sudoku dataset with varying mixing proportions. The results, presented in the table below, confirm our theoretical predictions.
Training Reward
Training Reward
2.0
1.4
1.5
#=
== 1.0
1.3
w
1.0
R = 10.0
train/reward
1000
2000
3000
4000
5000
500
1000
1500
2000
Training Steps
Training Steps
Figure 4: Left: Ablation study 1 in the weights (Equation 99) GSM8K. Right: wd/ training on MATH with a random seed different from the seed used in our main experiments. The abrupt decrease of the rewards in the early training (see Figure] 3 MATH500) disappears.
In all benchmark evaluations, we fix the hyperparameter & 1, which controls the scale of the exponential weighting in wdl. To validate this choice, we provide an ablation study on the coefficient d = in the exponential weight of wd/ (Equation 9 below. Larger values & leads to more extreme weight assigned to the samples. According Figure the training of applying different & converges to similar rewards 1 is small. Overly large value (c.g. 10) can cause performance drop, implying that extreme weight assignment is detrimental.
27

---

Published as a conference paper at ICLR 2026
Sudoku
Countdown
GSM8K
MATH500
Completion Langth
wdl
d1
210
110
220
Steps
Staps
Steps
Steps
Figure 5: Completion lengths dynamics of wdl and dl. In math problem-solvingg tasks (GSM8K and MATH500), our method demonstrates smaller completion lengths and better token efficiency.
C.3 CODING BENCHMARKS
We conduct 200 steps of wdl training on AceCode-87K (Zeng al. 2025) following the implemen- tation of Open-R1 (Hugging Face 2025). Our method achieves consistent improvements over the base model.
Table 10: Comparative performance of wd/ improvements compared to base model LLaDA-8B
Instruct. We present the results of wdl/ fine-tuned with AceCode dataset (Zeng et 2025)
Task
Gen Length
Steps
Block Size
wil
LLaDA
Hvaa
256
128
32
34.76 (+3.66)
31.10
EvaanEval
256
256
32
39.02 (+1.82)
37.20
HumanEval
512
512
32
36.59 (+0.61)
35.98
MBPP
128
128
32
39.2 (+2.40)
36.8
MBPP
256
256
32
36.6 (+1.00)
35.6
MBPP
512
512
32
36.8 (+0.40)
36.4
C.4 TRAINING DYNAMICS
Figure Opresents the reward dynamics over gradient steps during training. wdl exhibits a notably faster reward increase compared to dl, highlighting its superior sample efficiency-effectivele lever- aging the reward signal to accelerate policy optimization. In addition, Figure shows the average length of generated completions during training. On math reasoning benchmarks such GSM8K and MATH500, wdl converges to shorter output sequences than dl, suggesting improved token efficiency while maintaining or improving performance.
D LIMITATIONS
Similar to other RL-based approaches, wdl may lose effectiveness when all generations within sampled group receive identical rewards. This situation can occur under several conditions for example, when the training dataset is either too simple or too challenging for the base model. Nonetheless, such cases can be mitigated through careful reward design and the incorporation of curriculum learning strategies.
additional limitation of this work is that the current wd/ framework restricted to text-based reasoning. Extending it to multimodal reasoning unified diffusion-based models (c.g.. Yang et al. 2025 represents a valuable direction for future research.
A final limitation concerns the likelihood approximation used in wdl. Our approach relies on the based approximation, which is computationally efficient but introduces bias. Although some prior works employ L.BB--aaddd estimators (c.g., DCE), they require additional computational overhead (Zhao et al. 2025) often exhibit high variance, as demonstrated in Figure 1. This trade-off highlights important area for further exploration.
28

---

Published as a conference paper at ICLR 2026
D.1 ADDITIONAL ANALYSIS ON UNLEARNING
We provide extended demonstrations for Remark Remark 2 focusing specifically on the theoretical insights underlying the interpretation of the negative-sample reinforcement term in wdl/ form of data unlearning. Under the DCE likelihood approximation, the negative-sample reinforcement term wdl becomes
())
w (q. o). .lgg mo(ola) Eampitt) (xp(-AA(n)) log mo(olg)
(67)
=B =Bxo~p6()
exp ( (-AA --xa))))) bo- moo)
log po (o(xlilz)))
(68)
w/= []mmaakk]
LDCE
=B,, ~paee 1o (1.1))
Σ
log Po( (12
(69)
[mask]
CDCE + ELBO
where Pdata (20) p6(
Eng Pol eexp ((m Equation 69
expl A(r )}
holds by simply applying impor-
tance sampling-
Since DCE is equivalent to the evidence lower bound (ELBO) of masked discrete diffusion models, we draw an analogy between the final objective in Equation (69) and data unlearning in diffusion models (Alberti et al. 2025). Equation (69 can be viewed as a direct masked discrete-diffusion extension of NegGrad (Golatkar et 2020), which aims minimize the evidence lower bound of the log-liiel on samples with lower advantage.
29


Reviews:
[{"review_id": "gOv5fNcd8f", "reviewer": "Reviewer_kDWu", "summary": "This paper introduces wd1 and wd1++, a novel RL algorithm for dLLMs reasoning finetuning. Instead of requiring likelihood approximation in previous work, wd1 is ratio-free and from weighted log-likelihood maximization. Such a formulation reduces the estimation variance and the compute cost. The paper provides a theoretical foundation and demonstrates the connection with energy-based model training. Empirically, wd1 and wd1++ demonstrates superior performance on four benchmarks, while requiring fewer rollouts and no SFT.", "questions": "- Given the wd1 objective, is it easy to collapse to some certain solutions? how diversity the finetuned output is?,\n- Regarding the exponential weights, how easy to tune the weight parameters? is the performance sensitive to the scale?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The idea is well-motivated and paper is well-organized and easy to follow.,\n- The paper has strong theoretical foundation and analysis.,\n- The experiments are thorough and verify the effectiveness of the methods and design choices.,", "weaknesses": "- From my understanding, one key design choice is to choose reverse KL instead of forward KL, which makes the optimization analytic and ratio-free. However, the paper lacks the discussion on this motivation and its impact, e.g., how it will affect the constraint.,\n- The limitations are not discussed in detail. I didn’t find any place to explicitly discuss the limitations of the method.,", "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": "GYeAgpSd7Y", "reviewer": "Reviewer_cbic", "summary": "This paper introduces wd1, a novel reinforcement learning algorithm tailored for diffusion Large Language Models. Unlike traditional on-policy RL, wd1 leverages the closed-form optimal solution of the KL-regularized RL problem. It projects the current model toward this optimum using weighted regression on samples generated by behavior policies.\n\nThe authors also identify a significant issue regarding the under-penalization of negative samples and propose a complementary penalty term to address it. Experiments on popular math and reasoning benchmarks demonstrate that wd1 achieves significant improvements over the d1 baseline. Furthermore, ablation studies successfully validate the effectiveness of the algorithm's key design choices.", "questions": "In Figure 2 (MATH500), all tested RL algorithms seem to cause a significant drop in reward compared to the baseline. Could the authors elaborate on why this might be the case?\n\nSome copy-edits: Lemma 1, A_t(x_0) should be A_t(x_t).", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Wd1 employs a clever approach to bypass the computation of the log-probability ratio, which is often intractable or computationally expensive for diffusion LLMs.\n\nThe introduced penalty term for negative samples is a notable contribution. It appears to effectively address a long-standing issue in weighted behavior-cloning methods—namely, the insufficient penalization of negative samples. The proposed solution is both simple and effective.", "weaknesses": "wd1 appears highly sensitive to the coefficient that mixes positive weights w+ and negative weights w-. In Table 9, varying this coefficient lambda by 0.1 results in an accuracy drop of over 10%, which suggests the model may require careful tuning.", "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": "XUVkcM1wOx", "reviewer": "Reviewer_irEQ", "summary": "This paper presents a weighted log-likelihood objective function for fine-tuning dLLMs via reinforcement learning. Compared to other RL-based approaches, the proposed method, *wd1*, improves the reasoning ability of dLLMs and also reduces the computational overhead due to the ratio-free policy optimization. In addition, the authors provide a theoretical insight to connect their method with energy-guided diffusion sampling, by viewing the energy function as the negative advantage. In experiments, the proposed method is evaluated on four reasoning benchmarks.", "questions": "Please see \"Weaknesses\" and, what's the main limitation of this work? It would be better to provide a discussion on it.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The motivation and the main idea is clear, and the overall paper is well written and organized. \n- The proposed weighted log-likelihood objective is interesting.\n- In experiments, the proposed method outperforms other RL-based fine-tuning approaches.", "weaknesses": "1. The connection between the proposed method and the energy-guided diffusion model looks good but overstated. And the motivation is confusing: if fine-tuning the dLLM with wd1 is equivalent to training a energy-guided diffusion model, why don't you just fine-tuning dLLMs with the energy-based objective? The authors probably need to provide an ablation study or analysis on it.\n\n2. In experiments, the proposed method is built on LLaDA and only evaluated on four benchmarks. But as a large language model, it should be also evaluated on more datasets such as those in the LLaDA paper (Table 2) to illustrate its scalability.\n\n3. It would be better to give more discussion and related works for the “ratio-free” policy optimization. In Table 2, the difference of the training cost between two approaches is actually small. What's other advantages of using this “ratio-free” policy optimization?\n\n4. This RL approach designs different reward functions for different tasks and datasets, which would be unpractical for a foundation LLM. And the proposed model is trained individually on each dataset, which might be unfair to LLaDA model because it is not specifically fine-tuned on these datasets.", "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 (Poster)