Pith. sign in

REVIEW 3 major objections 5 minor 48 references

This paper argues that reinforcement learning for masked diffusion language models has been missing half the story: at each denoising step the model chooses both which tokens to write and which positions to reveal, and the policy gradient m

Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

T0 review · deepseek-v4-flash

2026-08-01 23:51 UTC pith:PNEOTERX

load-bearing objection Real, reproducible-looking gains from the masking term, but the paper's formal claim and implementation use different unmasking scores, and the SOTA comparison rests on a re-run baseline that disagrees with the original SPG number. the 3 major comments →

arxiv 2607.15200 v1 pith:PNEOTERX submitted 2026-07-16 cs.CL cs.AIcs.LG

Mask-Aware Policy Gradients for Diffusion Language Models

Haran Raajesh , Kulin Shah , Adam Klivans , Philipp Kr\"ahenb\"uhl This is my paper
classification cs.CL cs.AIcs.LG
keywords masked diffusion language modelspolicy gradientsunmasking ordertwo-stage action MDPPlackett-Luce samplingreinforcement learningmathematical reasoningcode generation
verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

Reinforcement learning for masked diffusion language models normally estimates an intractable log-likelihood by looking only at token predictions, ignoring the order in which positions get unmasked. This paper observes that an MDLM trajectory contains a second decision at every step: after predicting tokens at all masked positions, the model must pick a subset to keep revealed and remask the rest. The authors formalize this as a two-stage action MDP and show the policy gradient decomposes into a token term and a masking term, with the masking term computable from the model's own log-likelihoods at no extra forward-pass cost. If correct, this means trajectory-based RL for MDLMs that ignores remasking is using a biased gradient, and adding the masking term yields consistent 2-4 point gains across mathematics and code benchmarks. The paper supports this with a minimal counterexample where a token-only gradient is zero along a direction that strictly improves expected reward.

Core claim

Existing RL methods for MDLMs approximate the trajectory log-likelihood with token predictions only, either through ELBO variants or by treating generation as a sequence of token actions. The paper shows that an MDLM trajectory has a richer structure: at each step the model first samples tokens at all masked positions and then samples which of those positions to unmask, remasking the rest. By replacing the non-differentiable greedy top-K selection with a probabilistic Plackett–Luce distribution over position subsets, the remasking decision becomes differentiable and its log-probability can be included in the trajectory likelihood. The resulting policy gradient (Eq. 11) decomposes into a toke

What carries the argument

The central object is a two-stage action MDP for MDLM decoding: at step t the policy first samples tokens at masked positions and then samples an unmasking set Ut from a Plackett–Luce distribution whose weights are exp(v_k / τ), where v_k is the model's log-likelihood for the predicted token at position k. This makes the otherwise non-differentiable greedy top-K remasking into a differentiable distribution over position subsets. The trajectory likelihood (Eq. 10) multiplies token and unmasking probabilities, so the policy gradient (Eq. 11) splits into a token term and an unmasking term. The unmasking term is computed from the same forward passes as the token term, adding no extra parameters

Load-bearing premise

The method assumes the model's own token likelihood is a trustworthy guide to which positions should be unmasked, and that behavior learned with randomized remasking still transfers to the greedy confidence-based unmasking used at evaluation; if that confidence signal is miscalibrated or the train/inference mismatch is large, the masking term may reinforce behavior that is never executed at test time.

What would settle it

Run the exact same pipeline with the masking term removed (token-only trajectory gradient, matching StepMerge's estimator) at the same N=32, K=12 configuration and seeds; if the 2-4 point gains on GSM8K, MATH500, HumanEval, and MBPP vanish or shrink to noise, the central claim fails. A second check: evaluate with probabilistic remasking exactly matching training instead of greedy confidence-based unmasking; if the gains vanish or invert, the transfer assumption is the real driver.

Watch this falsifier. Get emailed when new claim-graph text bears on it.

If this is right

  • Trajectory-based RL for MDLMs that ignores the remasking step yields a biased policy gradient; including the masking term is necessary for a correct gradient estimate.
  • Jointly optimizing token and masking terms gives state-of-the-art results on GSM8K (87.1%), MBPP (53.4%), MATH500, and HumanEval, with gains of 2-4 points over prior ELBO- and trajectory-based methods.
  • The benefit grows with inference block size: at full-sequence decoding the margin over SPG widens to +4.3 on GSM8K, suggesting that position selection matters more when more positions are masked per step.
  • The method generalizes to a second base model (Dream-7B) and to planning tasks (Sudoku and Countdown), indicating the decomposition transfers beyond the primary setup.
  • Because the unmasking term adds no forward passes and no architectural changes, it can be dropped into existing trajectory-based estimators such as StepMerge with only the masking term added.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • A natural extension the authors do not pursue is applying the same two-term decomposition to other absorbing-state discrete diffusion settings, such as graph or protein generation, where an analogous 'which node to reveal' decision exists; the theory suggests the masking term should matter there too.
  • The train/inference mismatch is a testable risk: training uses probabilistic remasking while evaluation uses greedy confidence-based unmasking. If a base model's confidence is miscalibrated, the masking term could reinforce behaviors that are never executed at evaluation, so the gains may depend on calibration quality.
  • The theoretical counterexample is deliberately minimal; quantifying how much of the empirical improvement comes from variance reduction versus genuinely new gradient directions would clarify the mechanism behind the 2-4 point gains.
  • If the decomposition is as general as claimed, the masking term should also help preference-alignment RL for MDLMs, not just verifiable-reward reasoning; this is a cheap, direct test of the framework's reach.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit.

Referee Report

3 major / 5 minor

Summary. The paper proposes a policy-gradient method for masked diffusion language models (MDLMs) that treats generation as a two-stage action MDP: at each denoising step the model both predicts tokens and chooses which positions to unmask. The authors replace deterministic confidence-based top-K remasking with a probabilistic Plackett–Luce selection, derive a trajectory policy gradient that decomposes into a token term and an unmasking term (Eq. 11), and show that omitting the unmasking term can miss improving directions (Appendix G). Empirically, the method is evaluated on GSM8K, MATH500, HumanEval, and MBPP with LLaDA-8B-Instruct, reporting gains of 2–4 points over baselines and state-of-the-art accuracy (87.1% on GSM8K, 53.4% on MBPP at length 512), with additional results on Dream-7B and planning tasks.

Significance. If the central claims hold, this is a meaningful contribution: the decomposition gives a principled way to include unmasking-order decisions in RL for MDLMs, and the consistent gains over StepMerge—which shares the same trajectory-based framework but omits the position term—support the importance of the added signal. The paper provides a formal derivation, a toy counterexample, a large set of ablations, and publicly available code, which are strengths. However, the formal unbiasedness claim is currently tied to one definition of the unmasking score while the experimental section describes another, and the headline SOTA comparison is weakened by the SPG baseline footnote. These issues are load-bearing for the paper's central claims and need to be resolved before the results can be fully assessed.

major comments (3)
  1. [§3.1 Eq. (9); §4; Appendix G] The unmasking score is defined in three incompatible ways. Eq. (9) sets v_t^k = log π(hat z_t^k | z_{t-1}), the log-likelihood of the sampled token; Section 4 says it is the maximum logit; Appendix G uses the maximum token probability σ(|s|). Eq. (11) is an unbiased estimator of ∇J only if p_unmask in Eq. (10) is the actual rollout distribution. If the implementation samples with max-logit scores while Eq. (11) uses log-likelihood weights, the position term is the score function of the wrong distribution and the estimator is biased. If the implementation uses log-likelihood, the Section 4 sentence is wrong and the claimed link to greedy max-logit inference is unsupported. Either way, the formal and experimental descriptions cannot both be correct; this must be fixed before the unbiased-gradient claim and the empirical results can be assessed.
  2. [Table 1, SPG row; §4.1] The footnote to Table 1 reports that SPG's GSM8K@256 result is 83.9 after running the public code, replacing the originally reported 86.1. Ours is 85.9 at this configuration, so the original published baseline beats the proposed method. The 'best in every configuration' and state-of-the-art claims therefore depend on rejecting the original SPG number without explanation. Please report both numbers, explain the 2.2-point discrepancy (evaluation protocol, sampling, checkpoint, etc.), and restate the SOTA claim accordingly.
  3. [Appendix G, Proposition 2] The counterexample does not instantiate the paper's formal definitions. It sets v_t^k = max_x π_θ(x|z_t,k) (maximum token probability) and then defines p_unmask with exp(v/τ); this is neither Eq. (9)'s log-likelihood of the sampled token nor Section 4's maximum logit. Since this proposition is the main evidence for the incompleteness claim under the proposed method, it must be restated using Eq. (9)'s v. The zero-token-gradient argument is valid for the toy model as stated, but it does not currently support the paper's central claim because it uses a different unmasking distribution.
minor comments (5)
  1. [§3.1, Eq. (9)] The notation \w\t^k in Eq. (9) is undefined; it should be \z_t^k to match the rest of the section.
  2. [Figure 3 caption] The caption says the bars are 'proportional to predicted token log-likelihoods,' but Eq. (9) uses exp(v/τ), i.e., a temperature-scaled power of the likelihood. Please align the wording with the equation.
  3. [§4.1, Figure 5] Training uses stochastic remasking with τ_pos=0.5, while inference uses deterministic confidence-based unmasking. The paper does not discuss this train/inference mismatch; a sentence on why the learned position log-probabilities transfer to greedy selection would be helpful.
  4. [§4, generation parameters] The relationship between block size 32, 128 diffusion steps, and |U_t|=2 is not explained. Please state explicitly how many positions are unmasked per denoising step within a block.
  5. [Appendix F, Prop. 1] In Eq. (13) the stability assumption bounds a per-step log ratio by |U_t| ϵ_unmask, which makes the summation in Eq. (25) appear to include an extra factor. Clarify whether ϵ_unmask is per-position or per-set.

Circularity Check

0 steps flagged

No load-bearing circularity; the policy-gradient derivation is self-contained, with one minor non-load-bearing self-citation.

full rationale

The central derivation is self-contained rather than circular. Section 3 defines the extended trajectory distribution in Eq. 10 as the product of token likelihoods, a Plackett–Luce position-selection distribution, and a deterministic remasking indicator, and Eq. 11 is the exact REINFORCE score-function identity for that distribution. No fitted parameter enters: the unmasking scores v_t^k are model logits, and the position term is the derivative of the log-probability of the same stochastic remasking policy used during rollouts. The reported benchmark gains are external evaluations, not quantities defined by the estimator. The only self-citation is Kim et al. (2025) in Section 2 ('In practice however, a deterministic strategy that selects positions to unmask based on confidence scores derived from the model’s logits (Chang et al., 2022; Kim et al., 2025) often performs much better.'), which shares author Kulin Shah with this paper. This citation is not load-bearing: the paper immediately replaces deterministic top-K remasking with a probabilistic Plackett–Luce policy, and the formal claims (Eqs. 9–11, Appendix G) are established in the paper itself without relying on Kim et al.'s results. There is a separate correctness concern—Eq. 9 defines v_t^k = log π(ŵ_t^k|z_{t-1}), while Section 4 states 'The unmasking score at each masked position is the maximum logit from the model’s output,' and Appendix G uses maximum token probability—but this is an internal inconsistency between the formal estimator and the implementation, not a reduction of the claimed result to its own inputs. Therefore no significant circularity is present; the minor self-citation keeps the score at 2 rather than 0.

Axiom & Free-Parameter Ledger

5 free parameters · 5 axioms · 0 invented entities

The method relies on standard MDLM assumptions (factorizable token prediction, Markovian generation, Plackett-Luce sampling) and hand-selected hyperparameters. No new physical or architectural entities are introduced: the position score is the model's own token log-likelihood, not a separately trained head.

free parameters (5)
  • position temperature τ_pos = 0.5
    Controls randomness of Plackett-Luce remasking (Eq. 9); ablated in Table 8; directly shapes the masking gradient.
  • token temperature τ_tok = 0.9
    Token sampling temperature during RL rollouts; ablated in Table 8.
  • StepMerge segments N = 32
    Number of macro-steps in the trajectory log-probability approximation; controls approximation error (Prop. 1) and compute.
  • subsampled boundaries K = 12
    Number of segment boundaries evaluated for the gradient; ablated in Table 7b; affects variance and bias of the estimator.
  • unmask set size |U_t| = 2
    Number of positions unmasked per denoising step; chosen to match d1/wd1/SPG; ablated in Table 9.
axioms (5)
  • domain assumption Token prediction factorizes over positions: log π_θ(z_t | z_{t-1}, c) = sum over masked positions of per-token log-probabilities
    Used implicitly in Eq. 3 and Eq. 10; standard for LLaDA-style masked diffusion. If positions are coupled, the trajectory likelihood is mis-specified.
  • domain assumption Denoising is Markovian: π_θ(z_t | c, z_{≤t-1}) = π_θ(z_t | c, z_{t-1})
    Eq. 8; needed for the product factorization of the trajectory likelihood.
  • standard math Plackett-Luce sampling without replacement is a valid distribution over unmasking subsets
    Eq. 9; standard statistical model for sampling K items without replacement proportional to weights.
  • domain assumption Within-segment stability of token and unmasking log-likelihoods (Eq. 12, Eq. 13)
    Needed for Proposition 1's StepMerge error bound; not verified empirically for LLaDA. If violated, the position term is computed on an inaccurate trajectory probability.
  • domain assumption Rule-based verifiers produce correct rewards for GSM8K, MATH500, HumanEval, MBPP
    All RL training relies on these verifiers; errors in reward correctness bias the learned policy.

pith-pipeline@v1.3.0-alltime-deepseek · 22902 in / 20004 out tokens · 160898 ms · 2026-08-01T23:51:49.997946+00:00 · methodology

0 comments
read the original abstract

Reinforcement learning has proven effective for improving reasoning in large language models, but extending it to Masked Diffusion Language Models (MDLMs) remains challenging due to the intractability of the log-likelihood estimation. Existing approaches approximate this log-likelihood by modeling only the token predictions, ignoring the order in which positions are unmasked during generation. We observe that MDLM generation involves two decisions at each step: what tokens to place at each masked position and which positions to remask. We formalize this as a two-stage action MDP, showing that the policy gradient naturally decomposes into a token term and a masking term. Combining optimization of both terms leads to state-of-the-art outcomes on mathematical reasoning and coding benchmarks, with scores of 87.1% on GSM8K and 53.4% on MBPP.

Figures

Figures reproduced from arXiv: 2607.15200 by Adam Klivans, Haran Raajesh, Kulin Shah, Philipp Kr\"ahenb\"uhl.

Figure 1
Figure 1. Figure 1: Test accuracy of our method and baseline methods on two mathematical reasoning [PITH_FULL_IMAGE:figures/full_fig_p001_1.png] view at source ↗
Figure 2
Figure 2. Figure 2: Overview of our approach. Given a prompt [PITH_FULL_IMAGE:figures/full_fig_p002_2.png] view at source ↗
Figure 3
Figure 3. Figure 3: Illustration of a single denoising step with probabilistic remasking. (1) The model [PITH_FULL_IMAGE:figures/full_fig_p005_3.png] view at source ↗
Figure 4
Figure 4. Figure 4: Wall-clock training efficiency on GSM8K (8 H100 GPUs, generation length 128). [PITH_FULL_IMAGE:figures/full_fig_p008_4.png] view at source ↗
Figure 5
Figure 5. Figure 5: Ablation over inference block size. All methods use confidence-based unmasking [PITH_FULL_IMAGE:figures/full_fig_p009_5.png] view at source ↗
Figure 6
Figure 6. Figure 6: Training reward dynamics for our method on GSM8K, MATH500, and KodCode. [PITH_FULL_IMAGE:figures/full_fig_p016_6.png] view at source ↗

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

48 extracted references · 25 linked inside Pith

  1. [1]

    Hence,∂ θ2v1 2 ̸=0 for allθ 2 ̸=0

    for all finite θ2, the factor σ(|θ2|)(1 −σ(|θ 2|)) is strictly positive, and sign(θ2)̸=0 wheneverθ 2 ̸=0. Hence,∂ θ2v1 2 ̸=0 for allθ 2 ̸=0. Combining the above observations, we conclude that ∂θ2 punmask(2|z 1)̸=0 for allθ 2 ̸=0, and therefore ∂θ2 J(θ)̸=0 wheneverσ(θ 1)>σ(θ 3)andθ 2 ̸=0. 24 Published as a conference paper at COLM 2026 Finally, sincev= (0,...

    2026

  2. [2]

    Where-to-unmask: Ground-truth-guided unmasking order learning for masked diffusion language models

    Hikaru Asano, Tadashi Kozuno, Kuniaki Saito, and Yukino Baba. Where-to-unmask: Ground-truth-guided unmasking order learning for masked diffusion language models. arXiv preprint arXiv:2602.09501,

    arXiv

  3. [7]

    Soft policy optimization: Online off-policy rl for sequence models.arXiv preprint arXiv:2503.05453,

    Taco Cohen, David W Zhang, Kunhao Zheng, Yunhao Tang, Remi Munos, and Gabriel Synnaeve. Soft policy optimization: Online off-policy rl for sequence models.arXiv preprint arXiv:2503.05453,

    Pith/arXiv arXiv

  4. [8]

    Flashattention-2: Faster attention with better parallelism and work partitioning

    Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.08691,

    Pith/arXiv arXiv

  5. [10]

    Scaling diffusion language models via adapta- tion from autoregressive models.arXiv preprint arXiv:2410.17891,

    Shansan Gong, Shivam Agarwal, Yizhe Zhang, Jiacheng Ye, Lin Zheng, Mukai Li, Chenxin An, Peilin Zhao, Wei Bi, Jiawei Han, et al. Scaling diffusion language models via adapta- tion from autoregressive models.arXiv preprint arXiv:2410.17891,

    Pith/arXiv arXiv

  6. [12]

    Zemin Huang, Zhiyang Chen, Zijun Wang, Tiancheng Li, and Guo-Jun Qi

    URLhttps://openreview.net/forum?id=Lm8T39vLDTE. Zemin Huang, Zhiyang Chen, Zijun Wang, Tiancheng Li, and Guo-Jun Qi. Reinforcing the diffusion chain of lateral thought with diffusion language models.arXiv preprint arXiv:2505.10446,

    arXiv

  7. [13]

    Openai o1 system card

    Aaron Jaech, Adam Kalai, Adam Lerer, Adam Richardson, Ahmed El-Kishky, Aiden Low, Alec Helyar, Aleksander Madry, Alex Beutel, Alex Carney, et al. Openai o1 system card. arXiv preprint arXiv:2412.16720,

    Pith/arXiv arXiv

  8. [14]

    Learning unmasking policies for diffusion language models.arXiv preprint arXiv:2512.09106,

    Metod Jazbec, Theo X Olausson, Louis Béthune, Pierre Ablin, Michael Kirchhof, João Monteiro, Victor Turrisi, Jason Ramapuram, and Marco Cuturi. Learning unmasking policies for diffusion language models.arXiv preprint arXiv:2512.09106,

    Pith/arXiv arXiv

  9. [15]

    net/forum?id=DjJmre5IkP

    URL https://openreview. net/forum?id=DjJmre5IkP. 11 Published as a conference paper at COLM 2026 Wouter Kool, Herke van Hoof, and Max Welling. Buy 4 reinforce samples, get a baseline for free! InDeepRLStructPred@ICLR,

    2026

  10. [17]

    dllm-cache: Accelerating diffusion large language models with adaptive caching.github, 2025a

    Zhiyuan Liu, Yicun Yang, Yaojie Zhang, Junjie Chen, Chang Zou, Qingyan Wei, Shaobo Wang, and Linfeng Zhang. dllm-cache: Accelerating diffusion large language models with adaptive caching.github, 2025a. 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 perspectiv...

    Pith/arXiv arXiv

  11. [18]

    Discrete diffusion modeling by estimating the ratios of the data distribution.arXiv preprint arXiv:2310.16834,

    Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution.arXiv preprint arXiv:2310.16834,

    Pith/arXiv arXiv

  12. [20]

    Large language diffusion models.arXiv preprint arXiv:2502.09992,

    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.09992,

    Pith/arXiv arXiv

  13. [21]

    Your absorbing discrete diffusion secretly models the conditional distributions of clean data.arXiv preprint arXiv:2406.03736,

    Jingyang Ou, Shen Nie, Kaiwen Xue, Fengqi Zhu, Jiacheng Sun, Zhenguo Li, and Chongxuan Li. Your absorbing discrete diffusion secretly models the conditional distributions of clean data.arXiv preprint arXiv:2406.03736,

    Pith/arXiv arXiv

  14. [22]

    Principled rl for diffusion llms emerges from a sequence-level perspective

    Jingyang Ou, Jiaqi Han, Minkai Xu, Shaoxuan Xu, Jianwen Xie, Stefano Ermon, Yi Wu, and Chongxuan Li. Principled rl for diffusion llms emerges from a sequence-level perspective. arXiv preprint arXiv:2512.03759,

    arXiv

  15. [23]

    Accessed: 2025-01-24. 12 Published as a conference paper at COLM 2026 Fred Zhangzhi Peng, Zachary Bezemek, Sawan Patel, Jarrid Rector-Brooks, Sherwood Yao, Avishek Joey Bose, Alexander Tong, and Pranam Chatterjee. Path planning for masked diffusion model sampling.arXiv preprint arXiv:2502.03540,

    arXiv 2025

  16. [25]

    Deepseekmath: Pushing the limits of mathe- matical reasoning in open language models.arXiv preprint arXiv:2402.03300,

    Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Yang Wu, et al. Deepseekmath: Pushing the limits of mathe- matical reasoning in open language models.arXiv preprint arXiv:2402.03300,

    Pith/arXiv arXiv

  17. [27]

    wd1: Weighted policy optimization for reasoning in diffusion language models.arXiv preprint arXiv:2507.08838,

    Xiaohang Tang, Rares Dolga, Sangwoong Yoon, and Ilija Bogunovic. wd1: Weighted policy optimization for reasoning in diffusion language models.arXiv preprint arXiv:2507.08838,

    arXiv

  18. [28]

    Kimi Team, Angang Du, Bofei Gao, Bowei Xing, Changjiu Jiang, Cheng Chen, Cheng Li, Chenjun Xiao, Chenzhuang Du, Chonghua Liao, et al

    URL https://openreview.net/forum?id= L2rfd2Czbj. Kimi Team, Angang Du, Bofei Gao, Bowei Xing, Changjiu Jiang, Cheng Chen, Cheng Li, Chenjun Xiao, Chenzhuang Du, Chonghua Liao, et al. Kimi k1. 5: Scaling reinforcement learning with llms.arXiv preprint arXiv:2501.12599,

    Pith/arXiv arXiv

  19. [29]

    Fine-tuning discrete diffusion models via reward optimization with applications to dna and protein design

    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,

    Pith/arXiv arXiv

  20. [30]

    13 Published as a conference paper at COLM 2026 Guanghan Wang, Gilad Turok, Yair Schiff, Marianne Arriola, and Volodymyr Kuleshov

    URL https://openreview.net/ forum?id=18j5Q49GwN. 13 Published as a conference paper at COLM 2026 Guanghan Wang, Gilad Turok, Yair Schiff, Marianne Arriola, and Volodymyr Kuleshov. d2: Improved techniques for training reasoning diffusion language models.arXiv preprint arXiv:2509.21474, 2025a. Jin Wang, Yao Lai, Aoxue Li, Shifeng Zhang, Jiacheng Sun, Ning K...

    Pith/arXiv arXiv 2026

  21. [32]

    Kodcode: A diverse, challenging, and verifiable synthetic dataset for coding

    Zhangchen Xu, Yang Liu, Yueqin Yin, Mingyuan Zhou, and Radha Poovendran. Kodcode: A diverse, challenging, and verifiable synthetic dataset for coding. InFindings of the Association for Computational Linguistics: ACL 2025, pp. 6980–7008,

    2025

  22. [33]

    Taming masked diffusion language models via consistency trajectory reinforcement learning with fewer decoding step.arXiv preprint arXiv:2509.23924, 2025a

    Jingyi Yang, Guanxu Chen, Xuhao Hu, and Jing Shao. Taming masked diffusion language models via consistency trajectory reinforcement learning with fewer decoding step.arXiv preprint arXiv:2509.23924, 2025a. Ling Yang, Ye Tian, Bowen Li, Xinchen Zhang, Ke Shen, Yunhai Tong, and Mengdi Wang. Mmada: Multimodal large diffusion language models.arXiv preprint ar...

    arXiv

  23. [34]

    Llada 1.5: Variance-reduced preference optimization for large language diffusion models.arXiv preprint arXiv:2505.19223,

    Fengqi Zhu, Rongzhen Wang, Shen Nie, Xiaolu Zhang, Chunwei Wu, Jun Hu, Jun Zhou, Jianfei Chen, Yankai Lin, Ji-Rong Wen, et al. Llada 1.5: Variance-reduced preference optimization for large language diffusion models.arXiv preprint arXiv:2505.19223,

    Pith/arXiv arXiv

  24. [35]

    14 Published as a conference paper at COLM 2026 Optimizer.We use AdamW (Loshchilov & Hutter,

    for memory-efficient attention computation. 14 Published as a conference paper at COLM 2026 Optimizer.We use AdamW (Loshchilov & Hutter,

    2026

  25. [36]

    We follow the same train-test splitting as d1 (Zhao et al., 2025b) and SPG (Wang et al., 2026)

    test set. We follow the same train-test splitting as d1 (Zhao et al., 2025b) and SPG (Wang et al., 2026). For code generation, we train on KodCode-Light-RL-10K (Xu et al., 2025), a dataset of coding problems at varying difficulty levels with synthetic unit tests for reward computation. We evaluate on HumanEval (Chen et al.,

    2026

  26. [37]

    We follow the same dataset setup as GDPO (Rojas et al., 2026)

    (257 crowd-sourced Python tasks). We follow the same dataset setup as GDPO (Rojas et al., 2026). Reward functions.We follow the reward functions of d1 (Zhao et al., 2025b) and SPG (Wang et al.,

    2026

  27. [38]

    GSM8K.We use a composite reward with five additive components: • XML Structure Reward:+0.125 per correctly placed opening or closing tag

    for code generation. GSM8K.We use a composite reward with five additive components: • XML Structure Reward:+0.125 per correctly placed opening or closing tag. • Soft Format Reward: +0.5 if the output matches the pattern <reasoning>...</reasoning><answer>...</answer>. • Strict Format Reward:+0.5 for exact formatting with correct line breaks. • Integer Answ...

    2026

  28. [39]

    and Count- down (Pan et al., 2025), following the evaluation protocol of SPG (Wang et al.,

    2025

  29. [40]

    Following SPG, Sudoku uses 3 few-shot examples during rollouts and evaluation

    and d1 (Zhao et al., 2025b). Following SPG, Sudoku uses 3 few-shot examples during rollouts and evaluation. Table 3 reports accuracy at generation lengths 128, 256, and 512; results for prior methods are taken from SPG (Wang et al., 2026). Our method achieves the best results on both tasks at every generation length, outperforming the strongest baseline, ...

    2026

  30. [44]

    Best results inbold

    Underlined entries denote the configuration used in our main results. Best results inbold. E.5 Unmasking-set size At each denoising step, our method selects a set Ut of positions to unmask; for all main results we use |Ut|= 2 per step, matching d1 (Zhao et al., 2025b), wd1 (Tang et al., 2026), and SPG (Wang et al., 2026). The selection distribution punmas...

    2026

  31. [45]

    |U t|Gen Len 128 Gen Len 256 Gen Len 512 2 81.0 85.9 87.1 4 80.4 85.1 86.8 8 80.1 84.9 86.5 Table 9: Ablation over the unmasking-set size|Ut| on GSM8K

    This modest drop is expected, since larger sets commit more positions per step and yield lower-quality rollouts during training. |U t|Gen Len 128 Gen Len 256 Gen Len 512 2 81.0 85.9 87.1 4 80.4 85.1 86.8 8 80.1 84.9 86.5 Table 9: Ablation over the unmasking-set size|Ut| on GSM8K. We unmask|Ut|= 2 positions per step in all main results, matching the config...

    2026

  32. [46]

    Recall that at each denoising stept, the model produces probability πθ(·|z t) over the full vocabulary at every position

    G Incompleteness of the token-only policy gradient In this section, we show that optimizing only the token gradient can miss a direction that improves the reward, in the following setting. Recall that at each denoising stept, the model produces probability πθ(·|z t) over the full vocabulary at every position. Let Mt denote the set of positions that are st...

    2026

  33. [47]

    To show that, we write the expectation as sum and then change the order of expectation and derivative

    Therefore, we can write the above as ∂θ2 J|tok =E ξ∼π θ h R(ξ)1{u 1 =2}∂ θ2 logπ θ(x2 |z 1,u 1 =2) i =p unmask(2|z 1)E ξ∼π θ h R(ξ)∂ θ2 logπ θ(x2 |z 1,u 1 =2)|u 1 =2 i =p unmask(2|z 1)E x2∼πθ (·|z1,2) h ∂θ2 logπ θ(x2 |z 1,u 1 =2)Eu2,x1 h R(ξ)|u 1 =2,x 2 ii =p unmask(2|z 1)σ(θ 3)E x2∼πθ (·|z1,2) h ∂θ2 logπ θ(x2 |z 1,u 1 =2) i 23 Published as a conference p...

    2026

  34. [128]

    The defaults strike a good balance, with results remaining reasonably stable across the tested range

    Performance degrades as τtok is lowered, since less diverse rollouts weaken the group-relative advantage signal, and as τpos is raised, since more random selection unmasks low-confidence positions. The defaults strike a good balance, with results remaining reasonably stable across the tested range. 18 Published as a conference paper at COLM 2026 Algorithm...

    2026

  35. [256]

    E.2 RL algorithm Our decomposition is compatible with any policy gradient algorithm

    Results for baselines are from SPG (Wang et al., 2026). E.2 RL algorithm Our decomposition is compatible with any policy gradient algorithm. We compare three widely used algorithms, RLOO (Kool et al., 2019), GRPO, and GSPO, in Table 7a withN= 32 and K= 12, evaluating on GSM8K and MATH500 at generation length

    2026

  36. [512]

    Best results inbold, second best underlined

    SPG results are obtained by running their publicly available code. Best results inbold, second best underlined . D Comparison to LLaDOU LLaDOU (Huang et al., 2025), trained with the Diffusion Chain of Lateral Thought (DCoLT) algorithm, shares our high-level insight that the unmasking order should be part of the policy, and arrives at a similar probabilist...

    2025

  37. [1959]

    dkv-cache: The cache for diffusion language models.arXiv preprint arXiv:2505.15781,

    Xinyin Ma, Runpeng Yu, Gongfan Fang, and Xinchao Wang. dkv-cache: The cache for diffusion language models.arXiv preprint arXiv:2505.15781,

    Pith/arXiv arXiv

  38. [1992]

    Fast-dllm: Training-free acceleration of diffusion llm by enabling kv cache and parallel decoding.arXiv preprint arXiv:2505.22618,

    Chengyue Wu, Hao Zhang, Shuchen Xue, Zhijian Liu, Shizhe Diao, Ligeng Zhu, Ping Luo, Song Han, and Enze Xie. Fast-dllm: Training-free acceleration of diffusion llm by enabling kv cache and parallel decoding.arXiv preprint arXiv:2505.22618,

    Pith/arXiv arXiv

  39. [1999]

    Unified multimodal discrete diffusion.arXiv preprint arXiv:2503.20853,

    Alexander Swerdlow, Mihir Prabhudesai, Siddharth Gandhi, Deepak Pathak, and Katerina Fragkiadaki. Unified multimodal discrete diffusion.arXiv preprint arXiv:2503.20853,

    Pith/arXiv arXiv

  40. [2015]

    Proximal policy optimization algorithms.arXiv preprint arXiv:1707.06347,

    John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms.arXiv preprint arXiv:1707.06347,

    Pith/arXiv arXiv

  41. [2017]

    Training verifiers to solve math word problems.arXiv preprint arXiv:2110.14168,

    Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems.arXiv preprint arXiv:2110.14168,

    Pith/arXiv arXiv

  42. [2019]

    Shufan Li, Konstantinos Kallidromitis, Hritik Bansal, Akash Gokul, Yusuke Kato, Kazuki Kozuka, Jason Kuen, Zhe Lin, Kai-Wei Chang, and Aditya Grover

    URL https://api.semanticscholar.org/ CorpusID:198489118. Shufan Li, Konstantinos Kallidromitis, Hritik Bansal, Akash Gokul, Yusuke Kato, Kazuki Kozuka, Jason Kuen, Zhe Lin, Kai-Wei Chang, and Aditya Grover. Lavida: A large diffusion language model for multimodal understanding.arXiv preprint arXiv:2505.16839,

    Pith/arXiv arXiv

  43. [2021]

    dultra: Ultra-fast diffusion language models via reinforcement learning.arXiv preprint arXiv:2512.21446,

    Shirui Chen, Jiantao Jiao, Lillian J Ratliff, and Banghua Zhu. dultra: Ultra-fast diffusion language models via reinforcement learning.arXiv preprint arXiv:2512.21446,

    arXiv

  44. [2022]

    Evalu- ating large language models trained on code.arXiv preprint arXiv:2107.03374,

    Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evalu- ating large language models trained on code.arXiv preprint arXiv:2107.03374,

    Pith/arXiv arXiv

  45. [2023]

    Diffuseq: Se- quence to sequence text generation with diffusion models.arXiv preprint arXiv:2210.08933,

    Shansan Gong, Mukai Li, Jiangtao Feng, Zhiyong Wu, and LingPeng Kong. Diffuseq: Se- quence to sequence text generation with diffusion models.arXiv preprint arXiv:2210.08933,

    Pith/arXiv arXiv

  46. [2024]

    Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning.arXiv preprint arXiv:2501.12948,

    Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning.arXiv preprint arXiv:2501.12948,

    Pith/arXiv arXiv

  47. [2025]

    Block diffusion: Interpolating between autoregressive and diffusion language models.arXiv preprint arXiv:2503.09573,

    Marianne Arriola, Aaron Gokaslan, Justin T Chiu, Zhihan Yang, Zhixuan Qi, Jiaqi Han, Subham Sekhar Sahoo, and Volodymyr Kuleshov. Block diffusion: Interpolating between autoregressive and diffusion language models.arXiv preprint arXiv:2503.09573,

    Pith/arXiv arXiv

  48. [2026]

    Program synthesis with large language models.arXiv preprint arXiv:2108.07732,

    10 Published as a conference paper at COLM 2026 Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, et al. Program synthesis with large language models.arXiv preprint arXiv:2108.07732,

    Pith/arXiv arXiv 2026