wd1: Weighted Policy Optimization for Reasoning in Diffusion Language Models

In plain words

Some language models write text by starting from a sentence full of blanks and filling them in over several rounds, rather than writing one word after the next. Teaching such a model to reason with reinforcement learning is hard, because the training update needs the probability the model gives to a whole answer, and for these models that probability can only be estimated roughly. Existing methods need several rough estimates per update, and the errors add up.

The paper's method removes most of that work. For each question the model writes several answers. Answers that score better than their group average get more weight, and answers that score worse get a second weight that pushes their probability down. The update is then a weighted version of ordinary "make this text more likely" training. Only the model being trained needs its probability estimated.

The authors test this on LLaDA-8B-Instruct with puzzles (Sudoku, Countdown) and math word problems (GSM8K, MATH500). Their method beats a reproduced earlier method on the puzzles, and it is level or slightly ahead on the math sets. An extended version that also uses the intermediate drafts of each answer scores higher on the math sets.

TL;DR

• Reverse-KL-regularised policy optimisation has a closed-form optimum, and fitting the policy to it gives a weighted log-likelihood loss §3.1 eq. 5 eq. 6.
• wd1 adds a negative term so that low-advantage completions are pushed down, and the loss has no importance ratio and no reference-model term §3.2 eq. 8 eq. 9.
• Only one likelihood approximation is needed (the d1 estimator, one masked forward pass per completion) §2.3 §3.2.
• Author-reported: on LLaDA-8B-Instruct wd1 beats the authors' reproduction of d1 on Sudoku, Countdown and GSM8K, and matches it on MATH500 at 256 tokens tab. 1. wd1++ reaches 84.5 on GSM8K and 44.2 on MATH500 tab. 3.
• Many results are reported at a fixed seed with no spread across seeds, and one ablation table is restricted to early checkpoints tab. 4.

Problem

Masked diffusion LLMs (dLLMs) define a likelihood only through an iterative denoising process, so log π(o | q) is intractable §2.1 §2.3. Diffusion-based GRPO needs that likelihood for the current, old and reference policies, and it turns approximation errors into exponentially amplified policy ratios §1 §2.3 app. A. The authors frame two goals: fewer likelihood approximations, and no importance ratio §1 §3.

Method

Notation. Prompt q, completion o, reward R(q, o). For a group of G completions of one prompt, the group-relative advantage is Â_i = R(q, o_i) minus the group mean §2.2. The base model is LLaDA-8B-Instruct, trained with LoRA §5.

Step 1: where the loss comes from. The regularised objective (reverse KL to π_old plus reverse KL to π_ref) has the closed-form optimum in eq. 5. Minimising the KL from that optimum to π_θ gives the WLL loss in eq. 6. Its samples come from a geometric mixture π_old^ref, and its weights are proportional to exp(ψ Â) §3.1. The approximated form normalises the weights over the group eq. 7. With λ = 1 and β = 0 the mixture equals π_old, since the exponents in eq. 5 reduce to that §5 Implementation.

Step 2: wd1 loss. The loss in eq. 8 has a positive term with weight w⁺ and a negative term with weight w⁻ §3.2. Both weights are softmaxes over the advantages, with w⁺ = softmax(ψ Â) and w⁻ = softmax(−ψ Â) eq. 9. The paper sets ψ = 1 in all reported runs tab. 6. The positive term raises the likelihood of high-advantage completions, and the negative term lowers the likelihood of low-advantage ones §3.2. If all advantages are equal, the two weights are equal and the loss has zero gradient §3.2 app. D.

Step 3: likelihood approximation. The estimator follows d1. Each prompt token is masked with probability 0.15, every completion token is masked, one forward pass predicts the completion, and the log-probabilities of the completion tokens are summed §2.3 app. B.4 tab. 6. The paper calls this estimator biased but cheap §2.3 app. D.

Step 4: wd1++. Each completion is expanded to the clean predictions x̂₀|l at every denoising step l. The weights of eq. 9 are computed over this expanded set, and a random step l is sampled per update, with scale L/(G·l) eq. 10 §3.3. The paper's wd1++ results use β = 0 and λ = 1 and a training set derived from OpenR1 §5 app. B.1.

Theory (§4). Sampling from the regularised optimum is energy-guided sampling with energy −A §4; the guidance lemma is eq. 11. The advantage-weighted denoising objective is eq. 13, and with DCE as the likelihood estimate it equals the WLL loss eq. 6 and eq. 14 (Remark 1). The negative term acts as unlearning of low-advantage data (Remark 2), an analogue of NegGrad app. D. The proofs are in app. A. They were not checked line by line for this card.

Rewards app. B.2:

• GSM8K: five additive parts: +0.125 per correct XML tag (with small penalties for extra text after tags); +0.5 soft format; +0.5 strict format; +0.5 integer answer; +2.0 correct answer.
• Countdown: +1.0 if the expression reaches the target with the exact numbers; +0.1 if the numbers are right but the target is missed; 0 otherwise.
• Sudoku: fraction of empty cells filled correctly app. B.2.
• MATH500: format reward 1.00 (answer tags with \boxed inside), 0.75 (tags without \boxed), 0.50 (\boxed only), 0.25 (neither); +2.0 for a correct boxed answer app. B.2.
• wd1++: a Math-Verify verifier scores GSM8K and MATH500 app. B.2.

Training loop (wd1, one prompt). The pseudocode follows alg. 1 and eq. 8-eq. 9. The code normalises weights over a batch, not per prompt (see Paper vs code).

# wd1 update for one prompt  [alg. 1] [eq. 8] [eq. 9]
# Input:  prompt, group size G, reward function, psi (= 1), mu updates
# Output: updated model parameters

def wd1_iteration(model, old_model, prompt, group_size, reward_fn, psi, mu):
    # Sample G completions from the old policy; with lambda = 1 and
    # beta = 0 this is the geometric mixture  [§3.1] [§5]
    completions = sample(old_model, prompt, count=group_size)
    # Group-relative advantage, no division by the standard deviation
    # [§2.2] [code: wd1/trainers/rev_grpo_trainer.py#L628-L638]
    rewards = [reward_fn(prompt, o) for o in completions]
    advantages = rewards - mean(rewards)  # (G,)
    # Positive and negative weights, each a softmax over the group  [eq. 9]
    pos_weights = softmax(psi * advantages)  # (G,)
    neg_weights = softmax(-psi * advantages)  # (G,)
    # Reuse the same rollouts for mu gradient updates  [alg. 1]
    for update in range(mu):
        # Re-estimate each completion's log-likelihood with the current model
        logps = [approx_log_likelihood(model, prompt, o) for o in completions]
        # Lower the weighted loss; the 1/G factor is from [eq. 8]
        loss = 0
        for i in range(group_size):
            loss += (-pos_weights[i] + neg_weights[i]) * logps[i]
        loss = loss / group_size
        gradient_step(model, loss)

# One-pass likelihood estimate (d1-style)  [§2.3] [app. B.4]
# Input:  prompt, completion of length L
# Output: scalar approximation of log pi(completion | prompt)

def approx_log_likelihood(model, prompt, completion):
    # Mask each prompt token with probability 0.15  [tab. 6]
    # [code: wd1/trainers/rev_grpo_trainer.py#L319-L331]
    masked_prompt = mask_tokens(prompt, probability=0.15)
    # Mask all completion tokens, then predict them in one pass  [§2.3]
    masked_completion = mask_all(completion)  # (L,)
    logits = model(concat(masked_prompt, masked_completion))[-L:]  # (L, V)
    # Sum of per-token log-probabilities of the true tokens  [§2.3]
    # [code: wd1/trainers/rev_grpo_trainer.py#L436-L439]
    token_logps = log_softmax(logits)[arange(L), completion]  # (L,)
    return sum(token_logps)

Hyperparameters

Values for wd1 from tab. 6 unless stated. The reproduced d1 column is in tab. 6 and the SFT settings in tab. 5.

Results

Author-reported. Setup: LLaDA-8B-Instruct, LoRA RL, zero-shot pass@1, sequence lengths 256 and 512, four benchmarks §5 tab. 1. The d1 baseline is the authors' reproduction; the d1 numbers the original paper reports are in tab. 7.

d1 as reported by its authors tab. 7: Sudoku 16.7 / 9.5, Countdown 32.0 / 42.2, GSM8K 81.1 / 82.1, MATH500 38.6 / 40.2 (256 / 512 tokens).

Extended method and concurrent methods on GSM8K and MATH500 tab. 3, left:

Training cost on 4× A100 tab. 2 (μ = 8 gradient steps per rollout):

Rollouts to the best model tab. 8: wd1++ reaches its best GSM8K/MATH500 model at training step 20 with 1280 total rollouts; MDPO at step 150 with 19200; d1 at step 7500 with 30000. The paper summarises this as "10× fewer rollouts".

Ablation on the negative term and on SFT tab. 4 (restricted to checkpoints before 5000 steps, see Pitfalls):

Combined weight cw, Sudoku tab. 9: cw 0.4 gives 11.77 % accuracy, cw 0.5 gives 25.63 %, cw 0.6 gives 14.11 %. The text also reports cw = 0 (negative only) and cw = 0.8 as worse than the default, without numbers §5.2 fig. 2.

Coding, 200 RL steps on AceCode-87K, wd1 vs LLaDA-8B-Instruct tab. 10: HumanEval 34.76 vs 31.10 (length 256, 128 steps), 39.02 vs 37.20 (256, 256 steps), 36.59 vs 35.98 (512, 512 steps); MBPP 39.2 vs 36.8 (128), 36.6 vs 35.6 (256), 36.8 vs 36.4 (512).

Paper vs code

Code: github.com/xiaohangt/wd1 at commit 1a2e254, Apache-2.0. The first author links it in an author footnote on the title page, and the owner is an author, so it qualifies as official. The repository holds the training code for the WLL variants and the d1 baseline, and evaluation scripts. It has no code for wd1++ (see below).

• Loss sign and terms match. The code computes −Σ w⁺·log π + Σ w⁻·log π over completion tokens, the same as eq. 8 [code: wd1/trainers/rev_grpo_trainer.py#L145-L160].
• Normalisation differs. eq. 9 normalises over the G completions of one group, and §3.2 says the constant uses "completions from all groups". The code takes softmax over the advantages of the local process slice, after the batch has been gathered and sliced rev_grpo_trainer.py rev_grpo_trainer.py. The code has no 1/G factor rev_grpo_trainer.py. Either the code or the equation is off, and the card cannot tell which one the reported numbers used.
• Beta is not used. tab. 6 lists β = 0 for wd1, but the shipped train.yaml sets β = 0.04 and ε = 0.5 train.yaml. The WLL trainer's loss has no KL or ratio term at all [code: wd1/trainers/rev_grpo_trainer.py#L117-L160], so these settings do nothing for wd1.
• Starting model. train.yaml sets model_path: LLaDA-sft-s1k-merged train.yaml. The WLL run script used for the no-SFT GSM8K experiment sets GSAI-ML/LLaDA-8B-Instruct wll_NP_gsm8k.sh. Use the script, which matches the paper's "no SFT" claim.
• Mask seeds. The paper says the prompt mask is redrawn at every gradient step §3.2. The code draws one seed per inner iteration when a rollout batch is generated, and reuses those seeds for the batch [code: wd1/trainers/rev_grpo_trainer.py#L547-L550].
• Sampling. The paper samples from the geometric mixture §3.1. The WLL trainer calls the same generate on the current model, and no mixture path was found rev_grpo_trainer.py. With β = 0 the two coincide, but the code does not implement the general mixture.
• Advantage. No division by the standard deviation in the update, matching §2.2 rev_grpo_trainer.py. (The std is computed for logging only rev_grpo_trainer.py.)
• Countdown reward is more generous than the paper. The paper gives 0.1 only when the numbers are right and the target is missed app. B.2. The code returns 0.1 when the numbers are wrong reward_func.py and when the target is missed reward_func.py. It returns 1.0 on an exact hit reward_func.py.
• MATH format reward never reaches the top two levels. The code tests for "\boxed" in a normal (non-raw) Python string [code: wd1/math500_utils.py#L24]. In Python, \b is a backspace, so the test looks for a backspace character and fails. The result is that the format reward is 0.75 for answer tags and 0.25 otherwise (the paper's 1.00 and 0.50 levels are unreachable). Not run here; this follows from Python's string escape rules.
• Sudoku reward is padded. The code pads or truncates the answer to 16 characters before scoring reward_func.py. The paper's description (fraction of empty cells correct) matches, but the padding is not described.
• wd1++ is not in the repository. The released trainers are the WLL variant with the negative term (wll_d1_neg), the positive-only variant (wll_d1_pos_only) and a reference-policy variant [code: wd1/run_train.py#L153-L230]. No intermediate-step (stepwise) loss was found in the repository, so the wd1++ numbers in tab. 3 cannot be reproduced from it.
• Timing settings differ. tab. 6 uses 12 inner updates, while the cost measurement in app. B.5 uses μ = 8 with batch size 4. tab. 2 also reports NFEs "including μ = 8".
• Training temperature is not stated in the paper. The config documents a default of 0.9 diffu_grpo_config.py, and the eval callback uses 0.0 run_train.py.

No official code is missing; the discrepancies above are the findings.

Pitfalls

• Figures in the paper disagree with tables. The text says wd1 beats d1 on Sudoku by "43%" §5.1 and the introduction gives "+58.8%" §1. The table gives 76.4 vs 17.6, a gain of 58.8 percentage points, a ratio of about 4.3×. Several other headline deltas mix points, relative changes and different lengths. Take the numbers from the tables.
• Sudoku numbers differ between tables. tab. 1 reports 76.4 (256 tokens). The ablation tab. 4 reports 25.2 for the same name, from checkpoints before 5000 steps. The footnote of tab. 1 says the 25.2 and 24.2 values come from an earlier technical report at 5000 steps, and that training was then extended to 12500 steps. Do not compare tab. 4 with tab. 1.
• The reported "10× fewer rollouts" is a ratio of 15× to 23× on the table's numbers (1280 against 19200 for MDPO and 30000 for d1) tab. 8.
• Timing claim is smaller than the wording suggests. The text says the RL step is "substantially lower" in cost §5.1. tab. 2 shows 81.16 s against 103.5 s per step, about 22 % less. The SFT stage (2.01 h) is a separate saving.
• Single seed. The main results fix seed 42 tab. 6. Variation across seeds appears only as a reward curve fig. 4 and no error bars are reported.
• Uniform-reward groups. With equal advantages in one group the two weights match and the loss cancels §3.2. In the code the softmax runs over the whole local batch, so the cancellation holds only when every row has the same advantage. This is a reading of the code, not something the paper states.
• Loss is a sum over tokens. rev_grpo_trainer.py sums log-probabilities over completion positions without length normalisation, so long completions have larger magnitude. The paper does not discuss this.
• Negative-only and positive-heavy weighting are unstable. The paper reports that cw = 0 (negative only) collapses learning and cw = 0.8 degrades more quickly §5.2 fig. 2. The text values are not in any table.
• λ is overloaded. λ is the reverse-KL weight in eq. 4. §5.2 calls the positive/negative mixing weight "cw" and says it "corresponds to λ" §5.2. These are different quantities; do not mix them up.
• ψ is sensitive. ψ = 1 everywhere; ψ = 10 degrades results app. C.2 fig. 4.
• Large ψ concentrates weight. Larger ψ gives more extreme weights on individual samples app. C.2. The paper's ψ sweep is the only check fig. 4.
• Biased likelihood. The approximation is biased and the paper says so §2.3 app. D. The theory in §4 assumes the DCE likelihood, not the d1 estimator used in experiments.
• No monotonic-improvement result for the final loss. The main text shows a monotonic-improvement result only for the regularised objective §3.1. The negative term is added to the loss without an analogous guarantee in the main text §3.2 §4.
• Coding results are not compared with d1. They are a single 200-step run against the base model app. C.3.
• Reward and checkpoint choices. Evaluation checkpoints differ from d1's app. B.1, so comparisons with the d1 reproduction use different checkpoint grids.

Lineage

• Builds on: diffusion-based GRPO and d1 arXiv:2504.12216: masked diffusion RL with a single-step likelihood estimate and policy ratios. wd1 keeps the estimator and drops the ratios and reference term.
• Builds on: GRPO arXiv:2402.03300 for the group-relative advantage §2.2.
• Builds on: masked diffusion with DCE / ELBO for the base model arXiv:2406.03736 and LLaDA arXiv:2502.09992 §2.1.
• Compared against: d1 (reproduced and reported), diffu-GRPO, SDPO, TCR and MDPO tab. 3 tab. 7; SPG and d2 are discussed in Related Work §6.
• Related weighted-regression work: advantage-weighted regression (AWR) is cited as the closest off-policy relative §6.
• Superseded by: not stated.