MRO: Enhancing Reasoning in Diffusion Language Models via Multi-Reward Optimization

AI summary

Overview

• Research area: Natural Language Processing, specifically diffusion language models (DLMs) and their reasoning capabilities.
• Technical level: Advanced. The paper assumes familiarity with autoregressive LLMs, masked diffusion training objectives, Markov Decision Processes, policy gradient methods, and potential-based reward shaping.
• One-sentence scope: The paper proposes Multi-Reward Optimization (MRO), a set of rewards combined with test-time scaling, rejection sampling, and reinforcement learning, to strengthen token correlations in DLMs and thereby improve reasoning accuracy and denoising speed.

What This Paper Is About

Diffusion language models generate text by predicting many masked tokens in parallel at each denoising step, but they still trail autoregressive LLMs on reasoning tasks, especially when few denoising steps are used. The authors argue this gap stems from the fact that tokens predicted within a step, and across steps, are generated independently, so the model fails to capture the dependencies among them. Their goal is to optimize this "token correlation" directly, without the cost of scaling data or model size, and to show that doing so improves reasoning accuracy and allows fewer denoising steps.

Key Contributions

1. A diagnosis and formalization of the problem. The paper defines two kinds of token correlation in DLMs: intra-sequence correlation (dependencies among tokens generated within a single denoising step) and inter-sequence correlation (alignment between sequences produced at different denoising steps). An analysis of 50 decoding runs on a subset of GSM8K shows that higher correlation tends to yield higher reasoning accuracy.

2. The MRO framework. A Multi-Reward Optimization approach that directly optimizes token correlation using several rewards: a token verification reward (re-masking a predicted token and re-checking the model's probability for it), a perplexity reward computed with a pre-trained AR language model via the lmppl tool, and a quality reward that scores the final response 2, 1, or 0 depending on whether it follows the required <think> / <answer> format and whether the answer is correct. These are combined in an MDP formulation where each denoising step is an action, and are applied through test-time scaling, rejection sampling, and reinforcement learning.

3. Step-wise Group Reward Optimization (SGRO). A method that groups denoising steps so each group yields one reshaped reward instead of one per step, introduced to counteract the increased reward variance that potential-based reward shaping can cause in long decision trajectories. The authors provide theoretical analysis (Properties 1, 2, and 3 in Appendix B), including a proof that the token verification reward improves intra-sequence correlation from a mutual-information perspective.

4. Empirical validation across task types and optimization algorithms. Experiments on LLaDA-8B-Instruct and LLaDA-8B-Instruct-s1 across mathematical, scientific, and logical reasoning benchmarks show accuracy gains averaging 3 points under test-time scaling, plus a denoising speedup.

Main Findings

• Higher token correlation tracks with higher accuracy: An analysis of 50 decoding runs on a subset of GSM8K (Figure 1) shows that higher intra- or inter-sequence correlation tends to produce higher reasoning performance.

• Test-time scaling gains: With beam search (k set to 4, temperature 0.25), LLaDA with MRO reaches 82.6% on GSM8K at a response length of 512 versus 79.4% for LLaDA without MRO. On MATH500 at length 512, LLaDA + MRO improves from 34.4% to 38.0%.

• Rejection sampling gains across task categories: Table 1 reports that rejection sampling with MRO improves reasoning across mathematical, scientific, and logical tasks. The improvement is especially pronounced on logical reasoning: LLaDA-MRO-4 improves the Countdown task by +7.9 points over LLaDA at a response length of 128.

• Closing the gap with a strong AR LLM on math: On GSM8K, the model with MRO reaches 82.6%, close to Qwen2.5-7B-Instruct's 85.4%. On MATH500, MRO reaches 36.2% versus Qwen2.5-7B-Instruct's 41.1%.

• Larger beam size generally helps: MRO-4 outperforms MRO-2 in most test cases, which the authors attribute to exploring a broader range of denoising sequences, at a significant computational cost.

• Measurable correlation increases: Table 3, based on 200 sampled examples from MATH500 and GPQA across five decoding runs with different random seeds, shows LLaDA-MRO raising intra-sequence correlation from 3.44 (± 0.18) to 3.79 (± 0.16) and inter-sequence correlation from 1.02 (± 0.14) to 1.58 (± 0.12) on MATH500, and from 2.76 (± 0.21) to 3.34 (± 0.19) and 1.02 (± 0.15) to 1.27 (± 0.13) respectively on GPQA.

• Modest decoding overhead: Table 2 compares decoding time and score on MATH500 and GPQA. Vanilla LLaDA uses 0.35h to 0.39h and scores 33.2 on MATH500; the confidence-based reward baseline (LLaDA-TTS + CBR) uses 0.73h to 0.81h and scores 35.2; LLaDA-TTS + MRO uses 0.84h to 0.85h and scores 36.0. On GPQA, the corresponding figures are 0.16h to 0.21h / 29.2, 0.27h to 0.34h / 30.6, and 0.31h to 0.44h / 34.6.

• Denoising speedup: At a fixed response length of 512 and denoising step sizes of {64, 128, 256, 512}, LLaDA-MRO-4 outperforms the baseline at all step sizes, and achieves performance comparable to the baseline's 256 steps using only 128 steps. The authors attribute this to the inter-sequence correlation reward encouraging collaboration across steps.

• SGRO configuration: In the test-time scaling setup, SGRO was applied with the grouping parameter w set to 32, meaning beam search with MRO was applied every 32 steps. The authors also report that SGRO reduces reward evaluation frequency and thus the computational overhead of reward estimation (Appendix D).

Methodology in Plain English

The researchers start from the observation that a DLM predicts several masked tokens at once, and those predictions are treated as independent of each other, which creates inconsistencies in a chain-of-thought reasoning path. They reframe the denoising process as a sequence of decisions in a Markov Decision Process: the state is the prompt plus the currently partially denoised response plus the step number, the action is the newly predicted response, and the policy is the model's own denoising distribution.

With that framing, they can attach rewards to decisions. Three rewards are designed. The token verification reward re-masks one predicted token at a time and asks the model how likely that token is given the others, averaging over a randomly sampled subset of tokens to keep the cost down. The perplexity reward uses a separate pre-trained autoregressive model to score how readable the current intermediate response is, with a fixed upper-bound constant and a scaling factor to keep it on the same scale as the other rewards. The quality reward is a terminal signal given only at the end: 2 if the response follows the required format and the answer is correct, 1 if the format is right but the answer is wrong, and 0 if the format is violated.

These rewards are combined using potential-based reward shaping, where the intra-sequence rewards shape the intermediate steps and the quality reward arrives as a delayed terminal reward. Because this shaping can raise reward variance over long trajectories, the authors group denoising steps together (SGRO) so that fewer reshaped reward computations are needed.

The resulting reward signal is then used in three ways: at inference time via beam search (test-time scaling), by collecting high-reward decoded sequences and using them for supervised fine-tuning (rejection sampling), and via policy-gradient style reinforcement learning.

For evaluation, they use LLaDA-8B-Instruct and LLaDA-8B-Instruct-s1, test on GSM8K, MATH500, GPQA, 4×4 Sudoku and Countdown with 3 numbers, compare against LLaDA, LLaDA-s1, and several ~8B open-source AR LLMs, and compute final accuracy with the lm-evaluation-harness toolkit. The test-time scaling configuration used response lengths of 64, 128, 256, and 512, with denoising steps set to half the response length, temperature 0.25, and block lengths of 8, 64, and 64 for GSM8K, MATH500, and GPQA. During rejection sampling, they randomly selected two consecutive denoising steps to optimize and used a block length of 64 for Countdown and Sudoku; gradient accumulation was used to manage memory.

Why This Matters

This work shifts attention from scaling DLMs up to fixing a structural property of how they decode. If independent parallel token prediction is indeed a core reason DLMs reason poorly, then optimizing correlation offers a cheaper path to better reasoning than adding data and parameters, and it also makes faster (fewer-step) decoding viable.

Real-world applications suggested by the tasks in the paper:

• Mathematical problem solving: tutoring or homework assistants that need multi-step arithmetic and algebra reasoning (GSM8K, MATH500).
• Scientific question answering: graduate-level biology, physics, and chemistry reasoning (GPQA).
• Structured puzzle and planning tasks: logic-grid and constraint problems such as 4×4 Sudoku and Countdown with 3 numbers.
• Latency-sensitive generation: settings where a DLM must produce coherent output in fewer denoising steps, since MRO allowed comparable accuracy at roughly half the denoising steps in the tested setup.

Industry relevance: the models tested are open-source and around 8B parameters (LLaDA-8B-Instruct, LLaDA-8B-Instruct-s1), and the reward machinery uses off-the-shelf tools (lmppl, lm-evaluation-harness). Practitioners deploying diffusion LMs for reasoning workloads can apply MRO-style rewards at inference time without retraining, or use rejection sampling to bake the gains into weights. The reported decoding-time overhead of 0.84h to 0.85h versus 0.35h to 0.39h for vanilla decoding on MATH500 is the main practical caveat for inference-time use.

Future Directions

• Reducing the cost of larger beam sizes. MRO-4 usually beats MRO-2 but requires more sampling; the authors explicitly call for methods that lower the computational burden of DLM training while still searching for better denoising sequences.
• Broadening the correlation analysis. The correlation measurements came from 200 examples each from MATH500 and GPQA over five decoding runs; extending this to more datasets, task types, and model scales would test how general the correlation–accuracy link is.
• Generalizing the reward design. The token verification reward, perplexity reward, and format-based quality reward were tuned for chain-of-thought reasoning. Whether analogous rewards work for non-reasoning generation tasks is not reported.
• Comparing against other optimization methods at scale. The paper compares MRO with vanilla decoding, a confidence-based reward baseline, and unoptimized AR LLMs; comparisons against other RL fine-tuning approaches for DLMs, and against larger AR models, are not reported.

Target Audience

This paper is most useful to researchers and engineers working on diffusion language models, efficient or non-autoregressive text generation, and reinforcement-learning-based post-training for LLMs. It is also relevant to practitioners who need reasoning-capable models at low decoding latency, and to readers interested in why token-level dependency structure affects chain-of-thought accuracy. A background in masked diffusion training objectives and policy-gradient style optimization is needed to follow the method sections and the theoretical results in the appendix.

Authors’ abstract

Recent advances in diffusion language models (DLMs) have presented a promising alternative to traditional autoregressive large language models (LLMs). However, DLMs still lag behind LLMs in reasoning performance, especially as the number of denoising steps decreases. Our analysis reveals that this shortcoming arises primarily from the independent generation of masked tokens across denoising steps, which fails to capture the token correlation. In this paper, we define two types of token correlation: intra-sequence correlation and inter-sequence correlation, and demonstrate that enhancing these correlations improves reasoning performance. To this end, we propose a Multi-Reward Optimization (MRO) approach, which encourages DLMs to consider the token correlation during the denoising process. More specifically, our MRO approach leverages test-time scaling, reject sampling, and reinforcement learning to directly optimize the token correlation with multiple elaborate rewards. Additionally, we introduce group step and importance sampling strategies to mitigate reward variance and enhance sampling efficiency. Through extensive experiments, we demonstrate that MRO not only improves reasoning performance but also achieves significant sampling speedups while maintaining high performance on reasoning benchmarks.