wd1: Weighted Policy Optimization for dLLMs

• The paper introduces wd1, a reinforcement learning method that replaces ratio-based surrogates with a weighted-likelihood objective to address intractable likelihoods in dLLMs.
• wd1 reduces computational overhead by requiring only one likelihood approximation, leading to faster convergence and improved efficiency during training.
• Empirical results demonstrate that wd1 achieves up to 16% higher accuracy on reasoning benchmarks like Sudoku and Countdown compared to previous methods.

$\mathtt{wd1}$ is a policy optimization method for reasoning in diffusion-based LLMs (dLLMs). It was introduced to address a specific obstacle in reinforcement learning for dLLMs: the likelihood function is intractable, so existing methods typically approximate the current, old, and reference policy likelihoods at every update, then construct importance-sampling ratios from those approximations. In the formulation reported in “wd1: Weighted Policy Optimization for Reasoning in Diffusion LLMs,” $\mathtt{wd1}$ replaces that ratio-based surrogate with a weighted-likelihood objective that requires only a single approximation for the current parametrized policy likelihood, with the stated goals of reducing computational overhead and mitigating bias induced by denominator errors in policy ratios (Tang et al., 7 Jul 2025).

1. Problem setting and motivation

The method is posed in the standard reinforcement-learning view of sequence generation, where a prompt $q$ and a completion $o$ form a “one-step” MDP. The base objective is

$$\eta(\pi_\theta) = \mathbb{E}_{q\sim\mathcal{D},\,o\sim\pi_\theta(\cdot\mid q)} \bigl[R(q,o)\bigr].$$

In practice, the optimization is regularized with KL terms to prevent uncontrolled policy drift:

$$\max_\theta\; \mathbb{E}_{q,o\sim\pi_\theta}\bigl[R(q,o)\bigr] -\lambda\,\mathbb{E}_q\bigl[D_{\rm KL}(\pi_\theta(\cdot\mid q)\,\|\,\pi_{\rm old}(\cdot\mid q))\bigr] -\beta\,\mathbb{E}_q\bigl[D_{\rm KL}(\pi_\theta(\cdot\mid q)\,\|\,\pi_{\rm ref}(\cdot\mid q))\bigr].$$

For autoregressive models, $\log \pi_\theta(o\mid q)$ decomposes tokenwise. For dLLMs, by contrast, generation proceeds through an iterative denoising chain and only exposes an ELBO or a per-token proxy for $\log \pi_\theta(o\mid q)$. Existing methods such as Diffu-GRPO and d1 therefore approximate $\log\pi_\theta$, $\log\pi_{\rm old}$, and $\mathtt{wd1}$0 and form policy ratios of the form

$\mathtt{wd1}$1

The stated difficulty is twofold: approximation error in the denominator can induce potentially large bias, and the need to evaluate three policies multiplies computational cost (Tang et al., 7 Jul 2025). Within that context, $\mathtt{wd1}$2 is designed as a direct replacement for importance-sampling-based surrogates in dLLM RL.

2. Reverse-KL formulation and closed-form target policy

The derivation begins from a reverse-KL constrained policy optimization problem, written in terms of the old-policy advantage:

$\mathtt{wd1}$3

The paper gives a closed-form optimizer for this objective. The resulting target policy is

$\mathtt{wd1}$4

This expression is central because it converts the policy-update problem into a fitting problem: instead of optimizing a ratio-clipped or ratio-penalized surrogate directly, one can fit $\mathtt{wd1}$5 to $\mathtt{wd1}$6 by maximum likelihood. The negative log-likelihood target is written as

$\mathtt{wd1}$7

A proposal distribution is then introduced,

$\mathtt{wd1}$8

so that one may sample completions from a distribution that depends only on the old and reference policies rather than from $\mathtt{wd1}$9 itself. This suggests a policy-improvement scheme in which the advantage term is absorbed into normalized sample weights rather than into explicit policy ratios (Tang et al., 7 Jul 2025).

3. Weighted-likelihood objective and training procedure

Given a group of $q$0 completions $q$1, the method forms normalized positive weights

$q$2

Using only these weights yields a one-sided weighted-NLL objective,

$q$3

Because $q$4 can drive negative-advantage samples toward zero weight, the method adds a complementary penalization term with normalized negative weights

$q$5

The final $q$6 objective, explicitly named Weighted Diffusion Policy Optimization 1, is

$q$7

The operational consequence is that only $q$8 must be approximated, using the usual ELBO or masked-prompt surrogate; there is no ratio of the form $q$9 or $o$0. In the practical loop described in the paper, the model is initialized with $o$1, then at each outer iteration one sets $o$2, samples prompts, draws $o$3 completions from $o$4, computes group-relative advantages $o$5, constructs $o$6 and $o$7, and performs $o$8 inner gradient steps on $o$9 (Tang et al., 7 Jul 2025).

In most experiments, the reported setting is $$\eta(\pi_\theta) = \mathbb{E}_{q\sim\mathcal{D},\,o\sim\pi_\theta(\cdot\mid q)} \bigl[R(q,o)\bigr].$$0 and $$\eta(\pi_\theta) = \mathbb{E}_{q\sim\mathcal{D},\,o\sim\pi_\theta(\cdot\mid q)} \bigl[R(q,o)\bigr].$$1, so $$\eta(\pi_\theta) = \mathbb{E}_{q\sim\mathcal{D},\,o\sim\pi_\theta(\cdot\mid q)} \bigl[R(q,o)\bigr].$$2 and $$\eta(\pi_\theta) = \mathbb{E}_{q\sim\mathcal{D},\,o\sim\pi_\theta(\cdot\mid q)} \bigl[R(q,o)\bigr].$$3. The same $$\eta(\pi_\theta) = \mathbb{E}_{q\sim\mathcal{D},\,o\sim\pi_\theta(\cdot\mid q)} \bigl[R(q,o)\bigr].$$4 samples are reused over $$\eta(\pi_\theta) = \mathbb{E}_{q\sim\mathcal{D},\,o\sim\pi_\theta(\cdot\mid q)} \bigl[R(q,o)\bigr].$$5 gradient updates to amortize sampling cost, and log-likelihood is approximated as in d1 by randomly masking a subset of the prompt at each gradient step.

4. Relationship to prior dLLM RL methods

The immediate comparison class consists of Diffu-GRPO and d1. These methods are described as PPO-style objectives that clip or penalize ratios built from likelihood proxies and therefore require three calls to the proxy $$\eta(\pi_\theta) = \mathbb{E}_{q\sim\mathcal{D},\,o\sim\pi_\theta(\cdot\mid q)} \bigl[R(q,o)\bigr].$$6: one for $$\eta(\pi_\theta) = \mathbb{E}_{q\sim\mathcal{D},\,o\sim\pi_\theta(\cdot\mid q)} \bigl[R(q,o)\bigr].$$7, one for $$\eta(\pi_\theta) = \mathbb{E}_{q\sim\mathcal{D},\,o\sim\pi_\theta(\cdot\mid q)} \bigl[R(q,o)\bigr].$$8, and one for $$\eta(\pi_\theta) = \mathbb{E}_{q\sim\mathcal{D},\,o\sim\pi_\theta(\cdot\mid q)} \bigl[R(q,o)\bigr].$$9. The paper’s stated criticism is that small errors in approximating $$\max_\theta\; \mathbb{E}_{q,o\sim\pi_\theta}\bigl[R(q,o)\bigr] -\lambda\,\mathbb{E}_q\bigl[D_{\rm KL}(\pi_\theta(\cdot\mid q)\,\|\,\pi_{\rm old}(\cdot\mid q))\bigr] -\beta\,\mathbb{E}_q\bigl[D_{\rm KL}(\pi_\theta(\cdot\mid q)\,\|\,\pi_{\rm ref}(\cdot\mid q))\bigr].$$0 in the denominator can destabilize the ratio and that this design imposes avoidable overhead (Tang et al., 7 Jul 2025).

By contrast, $$\max_\theta\; \mathbb{E}_{q,o\sim\pi_\theta}\bigl[R(q,o)\bigr] -\lambda\,\mathbb{E}_q\bigl[D_{\rm KL}(\pi_\theta(\cdot\mid q)\,\|\,\pi_{\rm old}(\cdot\mid q))\bigr] -\beta\,\mathbb{E}_q\bigl[D_{\rm KL}(\pi_\theta(\cdot\mid q)\,\|\,\pi_{\rm ref}(\cdot\mid q))\bigr].$$1 removes all policy ratios. The update is expressed as weighted likelihood regression onto the reverse-KL solution, with positive weights rewarding advantageous completions and negative weights suppressing low-reward completions. In that sense, the method is not presented as a minor variant of ratio-clipped PPO for dLLMs, but as a reformulation of the surrogate itself.

A second point of contrast concerns supervision. The experiments are explicitly described as R1-Zero-like training, with no SFT stage and no supervised data. This matters because d1 is reported with an SFT stage in the benchmark comparison. A plausible implication is that the paper is attempting to separate gains from RL objective design from gains inherited from an initial supervised fine-tuning phase; the reported results are framed accordingly (Tang et al., 7 Jul 2025).

5. Empirical results on reasoning benchmarks

The empirical study is reported on widely used reasoning benchmarks, with all models built on LLaDA-8B-Instruct, evaluated zero-shot, using the best of maximum lengths 256 and 512. The base model, Diffu-GRPO, d1, and $$\max_\theta\; \mathbb{E}_{q,o\sim\pi_\theta}\bigl[R(q,o)\bigr] -\lambda\,\mathbb{E}_q\bigl[D_{\rm KL}(\pi_\theta(\cdot\mid q)\,\|\,\pi_{\rm old}(\cdot\mid q))\bigr] -\beta\,\mathbb{E}_q\bigl[D_{\rm KL}(\pi_\theta(\cdot\mid q)\,\|\,\pi_{\rm ref}(\cdot\mid q))\bigr].$$2 are all compared.

On Sudoku (256), the reported accuracies are 6.7% for LLaDA base, 16.1% for Diffu-GRPO, 17.6% for d1, and 25.2% for $$\max_\theta\; \mathbb{E}_{q,o\sim\pi_\theta}\bigl[R(q,o)\bigr] -\lambda\,\mathbb{E}_q\bigl[D_{\rm KL}(\pi_\theta(\cdot\mid q)\,\|\,\pi_{\rm old}(\cdot\mid q))\bigr] -\beta\,\mathbb{E}_q\bigl[D_{\rm KL}(\pi_\theta(\cdot\mid q)\,\|\,\pi_{\rm ref}(\cdot\mid q))\bigr].$$3. On Countdown, they are 19.5%, 27.0%, 25.8%, and 51.2%, respectively. On GSM8k, they are 76.7%, 80.7%, 78.2%, and 80.8%. On MATH500, they are 32.4%, 34.4%, 34.4%, and 34.4%. The summary in the paper states that $$\max_\theta\; \mathbb{E}_{q,o\sim\pi_\theta}\bigl[R(q,o)\bigr] -\lambda\,\mathbb{E}_q\bigl[D_{\rm KL}(\pi_\theta(\cdot\mid q)\,\|\,\pi_{\rm old}(\cdot\mid q))\bigr] -\beta\,\mathbb{E}_q\bigl[D_{\rm KL}(\pi_\theta(\cdot\mid q)\,\|\,\pi_{\rm ref}(\cdot\mid q))\bigr].$$4 achieves up to 16% higher accuracy and that, on Sudoku, it is 8 percentage points above d1, while on Countdown it is over 25 percentage points higher (Tang et al., 7 Jul 2025).

The learning-curve description is also specific. Average reward is reported to climb significantly faster than d1 across all four tasks. On the math benchmarks, $$\max_\theta\; \mathbb{E}_{q,o\sim\pi_\theta}\bigl[R(q,o)\bigr] -\lambda\,\mathbb{E}_q\bigl[D_{\rm KL}(\pi_\theta(\cdot\mid q)\,\|\,\pi_{\rm old}(\cdot\mid q))\bigr] -\beta\,\mathbb{E}_q\bigl[D_{\rm KL}(\pi_\theta(\cdot\mid q)\,\|\,\pi_{\rm ref}(\cdot\mid q))\bigr].$$5 is said to converge to shorter completions, which the paper interprets as better token efficiency. These observations are consistent with the algorithmic claim that removing likelihood-ratio estimation simplifies and stabilizes the optimization target.

6. Computational profile, limitations, and interpretation

The computational comparison is measured on 4×A100 with $$\max_\theta\; \mathbb{E}_{q,o\sim\pi_\theta}\bigl[R(q,o)\bigr] -\lambda\,\mathbb{E}_q\bigl[D_{\rm KL}(\pi_\theta(\cdot\mid q)\,\|\,\pi_{\rm old}(\cdot\mid q))\bigr] -\beta\,\mathbb{E}_q\bigl[D_{\rm KL}(\pi_\theta(\cdot\mid q)\,\|\,\pi_{\rm ref}(\cdot\mid q))\bigr].$$6 inner steps per global step. In that setting, d1 is reported to require approximately 2 hours of supervised fine-tuning, whereas $$\max_\theta\; \mathbb{E}_{q,o\sim\pi_\theta}\bigl[R(q,o)\bigr] -\lambda\,\mathbb{E}_q\bigl[D_{\rm KL}(\pi_\theta(\cdot\mid q)\,\|\,\pi_{\rm old}(\cdot\mid q))\bigr] -\beta\,\mathbb{E}_q\bigl[D_{\rm KL}(\pi_\theta(\cdot\mid q)\,\|\,\pi_{\rm ref}(\cdot\mid q))\bigr].$$7 requires 0 hours because there is no SFT stage. The RL per-step time is given as 103.5 s for d1 and 82.2 s for $$\max_\theta\; \mathbb{E}_{q,o\sim\pi_\theta}\bigl[R(q,o)\bigr] -\lambda\,\mathbb{E}_q\bigl[D_{\rm KL}(\pi_\theta(\cdot\mid q)\,\|\,\pi_{\rm old}(\cdot\mid q))\bigr] -\beta\,\mathbb{E}_q\bigl[D_{\rm KL}(\pi_\theta(\cdot\mid q)\,\|\,\pi_{\rm ref}(\cdot\mid q))\bigr].$$8. FLOPs per global step are reported as $$\max_\theta\; \mathbb{E}_{q,o\sim\pi_\theta}\bigl[R(q,o)\bigr] -\lambda\,\mathbb{E}_q\bigl[D_{\rm KL}(\pi_\theta(\cdot\mid q)\,\|\,\pi_{\rm old}(\cdot\mid q))\bigr] -\beta\,\mathbb{E}_q\bigl[D_{\rm KL}(\pi_\theta(\cdot\mid q)\,\|\,\pi_{\rm ref}(\cdot\mid q))\bigr].$$9 for d1 and $\log \pi_\theta(o\mid q)$0 for $\log \pi_\theta(o\mid q)$1. Forward passes, measured as NFEs per step, are summarized as $\log \pi_\theta(o\mid q)$2 for d1 and $\log \pi_\theta(o\mid q)$3 for $\log \pi_\theta(o\mid q)$4 (Tang et al., 7 Jul 2025).

The practical hyperparameters reported as robust across tasks are $\log \pi_\theta(o\mid q)$5, $\log \pi_\theta(o\mid q)$6, group size $\log \pi_\theta(o\mid q)$7, $\log \pi_\theta(o\mid q)$8, learning rate $\log \pi_\theta(o\mid q)$9, diffusion steps 128, and block length 32. The paper also emphasizes that the only approximation retained in $\log \pi_\theta(o\mid q)$0 is the approximation of $\log \pi_\theta(o\mid q)$1 through the usual ELBO or masked-prompt surrogate; there is no need to approximate $\log \pi_\theta(o\mid q)$2 or $\log \pi_\theta(o\mid q)$3.

The main limitation discussed is theoretical rather than empirical. The monotonic-improvement proof for reverse-KL updates does not directly apply once the negative-weight term $\log \pi_\theta(o\mid q)$4 is added, so convergence guarantees for the full two-sided weighting are left for future work. The paper also identifies several extensions: reintroducing $\log \pi_\theta(o\mid q)$5 to include an explicit reference policy, tuning $\log \pi_\theta(o\mid q)$6 for reward-variance control, adapting the method to preference-based rewards rather than scalar rewards, and proving convergence under the full two-sided weighting.

Taken together, these properties define $\log \pi_\theta(o\mid q)$7 as a weighted-likelihood reformulation of reinforcement learning for dLLMs. Its distinctive feature is not merely empirical performance, but the structural claim that one can optimize reasoning behavior in diffusion LLMs without the policy-ratio machinery that has dominated earlier RL formulations for this setting (Tang et al., 7 Jul 2025).