Our new X account is live! Follow @wizwand_team for updates
WorkDL logo mark

d-TreeRPO: Towards More Reliable Policy Optimization for Diffusion Language Models

About

Reinforcement learning (RL) is pivotal for enhancing the reasoning capabilities of diffusion large language models (dLLMs). However, existing dLLM policy optimization methods suffer from two critical reliability bottlenecks: (1) reward sparsity, arising from coarse or unverifiable signals that impede accurate advantage calculation; and (2) their probability estimates do not account for the gap to the unbiased expectation over all decoding orders, which are intractable to compute. To mitigate these issues, we propose d-TreeRPO, a reliable RL framework for dLLMs that leverages tree-structured rollouts and bottom-up advantage computation based on verifiable outcome rewards to provide fine-grained and verifiable step-wise reward signals. Furthermore, we provide a theoretical proof demonstrating that increasing prediction confidence effectively minimizes the gap between unbiased expected prediction probabilities and its single-step forward pass estimate. Guided by this analysis, we introduce a time-scheduled self-distillation loss during training that enhances prediction confidence in later training stages, thereby enabling more accurate probability estimation and better performance. Experiments demonstrate that d-TreeRPO outperforms existing baselines and achieves significant improvements across multiple reasoning benchmarks. Specifically, it achieves +86.2% on Sudoku, +51.6% on Countdown, +4.5% on GSM8K, and +5.3% on Math500 compared to the base model.

Leyi Pan, Shuchang Tao, Yunpeng Zhai, Zheyu Fu, Liancheng Fang, Minghua He, Lingzhe Zhang, Zhaoyang Liu, Bolin Ding, Aiwei Liu, Lijie Wen• 2025

Related benchmarks

TaskDatasetResultRank
Mathematical ReasoningGSM8K (train)
Accuracy82.6
22
Mathematical ReasoningMATH-500 (train)
Accuracy38.9
21
Advanced Mathematical ReasoningMath500 256 tokens
Pass@1 Accuracy41.2
15
Advanced Mathematical ReasoningMath500 512 tokens
Pass@1 Accuracy46.3
15
Arithmetic ReasoningCountdown 256 tokens
Pass@171.1
15
Arithmetic ReasoningCountdown 512 tokens
Pass@162.1
15
Grade School Math Word ProblemsGSM8k 256 tokens
Pass@181.2
15
Grade School Math Word ProblemsGSM8k 512 tokens
Pass@182.6
15
Sudoku SolvingSudoku 256 tokens
Pass@192.9
15
Sudoku SolvingSudoku 512 tokens
Pass@180.3
15
Showing 10 of 10 rows

Other info

Follow for update