PaperlayerView exact revision on arXiv

Source-linked AI summary

d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learning

Siyan Zhao, Devaansh Gupta, Qinqing Zheng, Aditya Grover

arXiv:2504.12216v2cs.CLcs.LG

TL;DR

The paper addresses whether masked diffusion language models can benefit from reasoning-oriented online RL, despite lacking autoregressive likelihood factorization. It introduces d1, combining masked SFT with diffu-GRPO and efficient randomized log-probability estimation. Across math and planning benchmarks, d1 achieves the strongest reported improvements, while fixed-length generation constrains the reasoning paths it can discover.

  • Problem

    It remains unclear whether masked diffusion language models can leverage recent reasoning improvements from online RL, because AR-based policy methods do not directly handle their non-sequential generation.

  • Method

    d1 sequentially applies masked SFT on reasoning traces and diffu-GRPO, using one-step log-probability estimation with random prompt masking for masked dLLMs.

  • Results

    Across four math and planning tasks, d1 consistently outperforms the base model, with gains of 26.2% on Countdown and 10.0% on Sudoku.

  • Takeaways & Limitations

    Combining SFT with diffu-GRPO yields the strongest reported post-training recipe and enables policy-gradient RL for masked dLLMs.

  • Takeaways & Limitations

    Because LLaDA requires fixed-length generation, diffu-GRPO may constrain discovery of either concise solutions or extended reasoning traces.

Abstract

from arXiv · show

Recent large language models (LLMs) have demonstrated strong reasoning capabilities that benefits from online reinforcement learning (RL). These capabilities have primarily been demonstrated within the left-to-right autoregressive (AR) generation paradigm. In contrast, non-autoregressive paradigms based on diffusion generate text in a coarse-to-fine manner. Although recent diffusion-based large language models (dLLMs) have achieved competitive language modeling performance compared to their AR counterparts, it remains unclear if dLLMs can also leverage recent advances in LLM reasoning. To this end, we propose d1, a framework to adapt pre-trained masked dLLMs into reasoning models via a combination of supervised finetuning (SFT) and RL. Specifically, we develop and extend techniques to improve reasoning in pretrained dLLMs: (a) we utilize a masked SFT technique to distill knowledge and instill self-improvement behavior directly from existing datasets, and (b) we introduce a novel critic-free, policy-gradient based RL algorithm called diffu-GRPO, the first integration of policy gradient methods to masked dLLMs. Through empirical studies, we investigate the performance of different post-training recipes on multiple mathematical and planning benchmarks. We find that d1 yields the best performance and significantly improves performance of a state-of-the-art dLLM. Our code is released at https://dllm-reasoning.github.io/.

1 Introduction

The paper asks whether masked diffusion language models can gain the reasoning improvements associated with online reinforcement learning. It proposes d1, combining masked SFT with diffu-GRPO, and finds consistent gains across four math and planning tasks.

  • Masked dLLMs generate text through iterative denoising rather than autoregressive next-token prediction, while using bidirectional context.
  • Existing AR-oriented RL methods cannot directly apply to masked dLLMs because their iterative generation lacks a natural sequential log-probability decomposition.
  • d1 combines supervised finetuning on reasoning traces with diffu-GRPO, a policy-gradient method designed for masked dLLMs.The method uses efficient one-step log-probability estimation and random prompt masking during policy updates.
  • Across four math and planning tasks, d1-LLaDA consistently outperforms the base LLaDA-8B-Instruct model.The reported comparisons use the best-performing generation sequence length for each task and model.
  • d1-LLaDA nearly doubles performance on planning tasks and surpasses both SFT-only and diffu-GRPO-only variants.The paper also reports consistent improvements in coding-task extensions and additional ablation and qualitative analyses.

2 Preliminaries

Masked dLLMs corrupt tokens progressively and learn to reconstruct original tokens, while GRPO uses grouped rewards to form advantages and clipped, KL-regularized policy updates.

  • 2.1 Masked Diffusion Large Language Models: Masked dLLMs progressively replace tokens with masks over time, with α_t decreasing until the sequence is fully masked at t = 1.At each timestep, tokens remain unmasked with probability α_t.
  • 2.1 Masked Diffusion Large Language Models: Training predicts the original tokens from corrupted inputs using a negative evidence lower bound that simplifies to a weighted negative log-likelihood.For LLaDA, the forward process uses α_t = 1 − t.
  • 2.1 Masked Diffusion Large Language Models: Masked dLLMs differ from BERT through time-varying masking ratios and multi-step generative decoding from pure noise.
  • 2.2 Group Relative Policy Optimization for Large Language Models: GRPO samples groups of responses for each prompt and derives token advantages from normalized group reward statistics.
  • 2.2 Group Relative Policy Optimization for Large Language Models: Using an unnormalized advantage avoids bias introduced by normalizing Q(q, o_i) − V(q) with a nonzero state function.
  • 2.2 Group Relative Policy Optimization for Large Language Models: The GRPO objective clips policy updates and applies a reverse KL penalty to limit update magnitude and deviation from the reference policy.ε controls clipping, while β controls KL regularization strength.

3 d1: Adapting Pre-trained Masked dLLMs to Reasoning Models

d1 adapts masked dLLMs to reasoning through sequential SFT and online policy-gradient RL. Its core mechanism estimates diffusion-model log-probabilities efficiently, enabling diffu-GRPO updates with randomized prompt masking.

  • 3 d1: Adapting Pre-trained Masked dLLMs to Reasoning Models: d1 sequentially combines SFT and online RL to enhance reasoning in pretrained masked dLLMs.
  • 3 d1: Adapting Pre-trained Masked dLLMs to Reasoning Models: Masked dLLMs require specialized likelihood estimation because their sequence probabilities lack autoregressive factorization and their denoising predictor runs repeatedly.
  • 3.1 Efficient Log Probability Estimation for Masked dLLMs: The method approximates sequence log-probability with independent per-token terms and estimates each token probability using one forward pass.
  • 3.1 Efficient Log Probability Estimation for Masked dLLMs: For a randomly masked prompt q′, one-step unmasking of the fully masked completion supplies estimated per-token log-probabilities for diffu-GRPO.
  • 3.1 Efficient Log Probability Estimation for Masked dLLMs: The estimator avoids LLaDA’s Monte Carlo approach, which uses 128 samples and hundreds of forward passes, making online policy optimization inefficient.
  • 3.2 diffu-GRPO: Policy Gradient Optimization for Masked dLLMs: Using these estimators, diffu-GRPO extends GRPO to masked dLLMs and can also be applied to PPO or REINFORCE.
  • 3.2 diffu-GRPO: Policy Gradient Optimization for Masked dLLMs: Randomly remasking prompts creates perturbed views of fixed prompt-completion pairs, regularizing optimization and enabling more inner updates per batch.
  • 3.3 SFT Recipe: SFT trains LLaDA on 1,000 high-quality reasoning questions with randomly masked tokens and traces containing verification and backtracking behaviors.

4 Experiments

Experiments evaluate SFT, diffu-GRPO, and their sequential combination across mathematical, planning, coding, and unified-task settings. The combined d1 recipe generally achieves the strongest reasoning performance, while results also examine scaling behavior and design choices.

  • Experimental Setup: The experiments compare SFT, diffu-GRPO, and sequential d1 training on six reasoning tasks spanning mathematics and planning.LLaDA-8B-Instruct is used as the base model, with task-specific post-training recipes and zero-shot evaluation across multiple generation lengths.
  • Main Results: diffu-GRPO outperforms both LLaDA-8B-Instruct and SFT in all 12 evaluated setups, while SFT beats the baseline in 7 of 12.LLaDA+diffu-GRPO also improves over the base model across all setups.
  • Main Results: A unified model trained across GSM8K, MATH500, Countdown, and Sudoku retains strong performance without sacrificing accuracy relative to per-task diffu-GRPO.The combined-task training data are subsampled so each task contributes the same number of examples.
  • Main Results: diffu-GRPO consistently improves coding performance over the initialization checkpoint, while s1k is unsuitable for coding because it lacks code examples.The coding evaluation uses KodCode-Light-RL-10K and includes HumanEval and MBPP.
  • Discussion: At generation length 512, SFT and d1-LLaDA produce qualitative self-correction and backtracking behaviors described as “aha moments.”These behaviors include verification of intermediate results and backtracking to the correct answer.
  • Discussion: Longer generation lengths improve GSM8K and MATH500 performance but produce mixed scaling on Countdown and declining performance on Sudoku.The reported average jump from 128 to 256 tokens is ∼7.1%, compared with ∼2.5% from 256 to 512; effective token usage increases without truncation at 128 tokens.
  • Design Choices and Ablations: Random masking consistently outperforms fixed masking and permits higher policy optimization update values while maintaining or improving performance.The method supports µ values of 12 or 24, whereas conventional approaches typically limit µ to 2 because of diminishing returns and overfitting risks.

5 Related Works

Diffusion language models evolved from continuous latent approaches toward scalable masked discrete diffusion, while recent work has begun exploring reasoning and related generation strategies.

  • Masked diffusion became an effective discrete approach for applying diffusion models to language.
  • Recent diffusion-language-model research has explored chain-of-thought reasoning and block-based generation.

6 Conclusion

The paper studies reasoning improvements for diffusion LLMs through SFT and diffu-GRPO, then combines them in the d1 pipeline. The combined recipe produces the largest improvements over baseline, while future work targets efficient longer decoding.

  • SFT on reasoning datasets improves diffusion-LLM performance and reveals “Aha moments”.
  • diffu-GRPO consistently outperforms SFT across the evaluated benchmarks.
  • Combining SFT and diffu-GRPO in d1 delivers the most significant improvements over the baseline.
  • Future work should develop efficient decoding strategies to scale generation length for more effective RL training.

A Limitations

Masked dLLMs use time-varying token masking and iterative denoising, with LLaDA adopting a linear noise schedule. Their formulation differs from BERT through multi-step generative decoding.

  • Masked dLLMs gradually corrupt tokens with a mask token, with masking probability governed by a decreasing noise schedule α_t.
  • At t = 1, all tokens are masked, and the reverse process gradually unmasks the sequence to recover the data distribution.
  • The bidirectional predictor fθ simultaneously predicts masked tokens from corrupted inputs using a weighted NLL derived from the NELBO objective.
  • LLaDA uses the linear noise schedule α_t = 1 − t.

D Experiment Details

The experiments use semi-autoregressive masked diffusion decoding and specified hardware, optimization, sequence-length, and parameter-efficient training settings for diffu-GRPO.

  • Sequences are decoded in N^2 denoising steps, unmasking 2 tokens per step.
  • The default decoding strategy processes 32-token blocks from left to right and unmasks the highest-confidence tokens within each block.
  • diffu-GRPO uses LoRA with rank r = 128 and scaling factor α = 64.
  • For math and planning tasks, training uses 8 NVIDIA A100-80G GPUs, sequence length 256, per-GPU batch size 6, and two accumulation steps.
  • Coding experiments use 4 NVIDIA RTX A5000 GPUs, with 7500 steps from the base model and 9000 steps from the SFT model.

D.1.1 Reward Functions, RL Training, and Evaluation Datasets

The evaluation combines task-specific rewards with RL training across mathematical, planning, and coding tasks. Results indicate that dataset difficulty matters for scaling reasoning, while d1-LLaDA maintains stronger or comparable reward trajectories.

  • RL Training: Across four reasoning tasks, d1-LLaDA achieves consistently higher or comparable reward trajectories than LLaDA+diffu-GRPO.The comparison is shown in Figure 7 during RL training.
  • Reward Functions: Reward functions combine formatting, reasoning, and solution correctness, with task-specific granularity.Countdown rewards exact target achievement and number usage; Sudoku rewards the proportion of correctly filled empty cells.
  • Reward Functions: GSM8K uses composite rewards for correctness, safe code, XML structure, soft and strict formatting, and valid integer answers.Correctness includes exact ground-truth matching, while code safety assigns zero reward when blocked modules are used.
  • Reward Functions: MATH500 and coding use structured rewards that separately assess answer formatting, correctness, and code organization.MATH500 awards formatting points based on answer tags and boxed answers; coding adds rewards for XML structure and Python code fencing.
  • RL Training: RL training samples partially masked responses before gradient updates, with SFT using long reasoning traces and LoRA-based adaptation.The SFT setup uses sequence length 4096, while the training data are truncated to that maximum length.
  • Evaluation Datasets: Overly hard datasets can degrade performance through overfitting, whereas weaker models benefit from a mixture of easier and harder examples.The reported relationship links dataset difficulty and model strength to reasoning-scaling behavior.

E.1 LLaDA-Instruct

LLaDA-Instruct produces more detailed reasoning when given longer generation sequences. At sequence length 512, its examples show explicit calculations and, for some tasks, correct self-contained solutions.

  • Generation Length: At sequence length 512, dLLM generations contain more detailed step-by-step solutions than at length 128.The longer outputs include explicit mathematical formulas and enumerated calculations.
  • Examples: At length 128, LLaDA-Instruct correctly answers the flag-row problem with 8 rows of 5 stars.The reasoning subtracts 24 and 12 stars from 76, then divides the remaining 40 by 5.
  • Examples: At length 512, LLaDA-Instruct correctly solves the flag problem by calculating 40 remaining stars and dividing by 5.The resulting answer is 8 rows of 5 stars.
  • Examples: At length 512, LLaDA-Instruct correctly solves the bakery problem and reports 5 afternoon loaves.The example computes 40 sold in the morning, 20 remaining, and 10 sold across afternoon and evening before dividing equally.

E.2 SFT on Reasoning Traces from S1K

SFT on S1K reasoning traces produces examples with mixed correctness and visible self-correction behavior. The generated solutions sometimes revise an initially incorrect interpretation before answering.

  • SFT Examples: LLaDA+SFT correctly solves the flag problem by subtracting 36 specified-row stars from 76 and dividing 40 by 5.The example verifies that 8 rows of 5 stars account for the remaining 40 stars.
  • Self-Correction: The SFT bakery trace explicitly revisits its calculation after identifying that the initial interpretation seems incorrect.Several passages describe a correction attempt before the final answer.

E.4 d1-LLaDA

d1-LLaDA examples show variable reasoning quality across tasks. The traces include a correct flag solution, an incorrect bakery answer, and an incorrect percentage answer accompanied by an attempted recalculation.

  • Flag Example: d1-LLaDA correctly solves the flag problem with 8 rows of 5 stars.The trace computes 24 stars in three 8-star rows, 12 in two 6-star rows, and 40 remaining stars.
  • Percentage Example: The puppy-percentage example aggregates 20 puppies and 7 with spots before reporting an answer of 35.The trace includes an explicit mistake check before its final response.
  • Bakery Example: d1-LLaDA gives an incorrect bakery answer of 5 for the afternoon-loaves problem.The trace states that 40 loaves are sold in the morning and ultimately reports 5.
Loading 2504.12216v2…