Research map / Large language models / Reasoning and chain of thought
MRO: Enhancing Reasoning in Diffusion Language Models via Multi-Reward Optimization
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes MRO, a multi-reward optimization framework (token verification, perplexity, and answer-quality rewards plus step-wise group reward shaping) that improves reasoning of masked diffusion LLMs via test-time scaling, rejection sampling, and REINFORCE.
Problem
Diffusion language models lag autoregressive LLMs on reasoning, especially with fewer denoising steps, because masked tokens are generated independently across steps, weakening intra- and inter-sequence token correlation.
Approach
Define intra-sequence and inter-sequence token correlation and directly optimize them with multiple rewards (token verification, perplexity from an AR LM, and format/accuracy quality), using a potential-based reward shaping over groups of denoising steps (SGRO) to reduce reward variance.
Key results
- LLaDA-MRO-4 via rejection sampling reaches 82.6% on GSM8K at length 512 vs 79.4% for LLaDA, and 36.2% vs 34.4% on MATH500.
- On Countdown at length 128, LLaDA-MRO-4 gains +7.9 points over LLaDA (22.0 vs 14.1).
- LLaDA-MRO-4 with 128 denoising steps matches the LLaDA baseline with 256 steps, indicating faster denoising.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add MRO through rejection sampling with beam size 4 | LLaDA-8B-Instruct · GSM8K, response length 512 | accuracy (%): 79.4 -> 82.6 | Helped |
| Add MRO through rejection sampling with beam size 4 | LLaDA-8B-Instruct · MATH500, response length 512 | accuracy (%): 34.4 -> 36.2 | Helped |
| Add MRO through rejection sampling with beam size 4 | LLaDA-8B-Instruct · GPQA, response length 512 | accuracy (%): 30.3 -> 34.3 | Helped |
| Add MRO through rejection sampling with beam size 4 | LLaDA-8B-Instruct · Countdown, response length 128 | accuracy (%): +7.9 points | Helped |
| Add MRO through rejection sampling with beam size 4 | LLaDA-8B-Instruct · Sudoku, response length 64 | accuracy (%): 11.2 -> 17.2 | Helped |
| Add MRO through REINFORCE | LLaDA-8B-Instruct · MATH500, response length 512 | accuracy (%): 34.4 -> 37.4 | Helped |
| Use MRO test-time scaling instead of confidence-based reward test-time scaling | LLaDA-8B-Instruct · GPQA | score: 30.6 -> 34.6 | Helped |
| Add MRO | LLaDA-8B-Instruct · MATH500 | intra-sequence correlation reward: 3.44±0.18 -> 3.79±0.16 | Helped |
| Add MRO | LLaDA-8B-Instruct · MATH500 | inter-sequence correlation reward: 1.02±0.14 -> 1.58±0.12 | Helped |
| Add SGRO (ablation) | LLaDA-MRO · MATH500, response length 256 | accuracy (%): 32.8 -> 34.2 | Helped |
| Add SGRO (ablation) | LLaDA-MRO · MATH500, response length 512 | accuracy (%): 35.4 -> 36.2 | Helped |
| Add SGRO (ablation) | LLaDA-MRO · GPQA, response length 256 | accuracy (%): 31.3 -> 32.1 | Helped |
Surprising findings
- MRO-v4 (without inter-sequence reward and without perplexity) outperforms full MRO under reinforcement learning, which the authors attribute to variance from the perplexity reward.
- Reinforcement learning with MRO does not clearly beat rejection sampling in general.
Limitations
- SGRO introduces new hyperparameters (group size w, group subset size) that need tuning.
- Larger beam size improves results but significantly increases sampling cost.
- Gains from RL are not substantially larger than rejection sampling; the RL method is basic (REINFORCE).
Experimental setup
Datasets: s1 DeepScaleR Countdown-Tasks-3to4 Sudoku-Dataset GSM8K MATH500
Benchmarks: GSM8K MATH500 GPQA Countdown Sudoku
Metrics: accuracy decoding time intra-sequence correlation score inter-sequence correlation score
Baselines: LLaDA-8B-Instruct LLaDA-s1 LLaMA-3-8B-Instruct Mistral-7B-Instruct DeepSeek-LLM-7B-Chat Qwen2.5-7B-Instruct DiffuLLaMA Dream-7B-Instruct EDLM d1-LLaDA
Builds on: LLaDA d1 DoT EDLM potential-based reward shaping
Closest papers by meaning
- d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learning (NeurIPS 2025 · similarity 0.90)
- Reinforcement Learning for Diffusion LLMs with Entropy-Guided Step Selection and Stepwise Advantages (NeurIPS 2026 · similarity 0.90)
- Denoising Time Matters: Diverse Generation in Diffusion Language Models (NeurIPS 2026 · similarity 0.90)
- d2: Improving Reasoning in Diffusion Language Models via Trajectory Likelihood Estimation (ICML 2026 · similarity 0.90)
- From Next-Token to Next-Block: A Principled Adaptation Path for Diffusion LLMs (NeurIPS 2026 · similarity 0.89)
- Theoretical Benefit and Limitation of Diffusion Language Model (NeurIPS 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.