Title: MRO: Enhancing Reasoning in Diffusion Language Models via Multi-Reward Optimization
Venue: NeurIPS 2025
Accepted: yes
Keywords: Generative Model

Language Model

Reasoning

Diffusion Language Model

Reward Optimization
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces MRO, a novel framework to improve reasoning capabilities in DLMs by directly enhancing token correlation during the denoising process. Empirical results on multiple reasoning benchmarks (e.g., GSM8K, MATH500, GPQA, Countdown, Sudoku) show that MRO significantly improves the reasoning performance of DLMs and even narrows the gap with strong AR LLMs.\n\nSTRENGTHS AND WEAKNESSES:\n**Strength:**\n\nThe method is tested across a diverse set of reasoning benchmarks and consistently demonstrates improvements over strong DLM baselines. The performance gains are especially pronounced in domains such as logical reasoning (e.g., Sudoku, Countdown), where token correlation is critical.\n\n**Weakness:**\n\n- While the paper introduces a multi-reward framework to enhance token correlation, it lacks a rigorous theoretical foundation linking these rewards directly to improved reasoning performance. The definition of intra- and inter-sequence correlation rewards is intuitive, but the paper does not sufficiently justify why these particular reward signals are optimal.\n- The comparisons between DLMs with MRO and autoregressive LLMs like Qwen2.5 and LLaMA-3-8B appear to be superficially favorable but are not fair. What is the setup of AR models? Are they evaluated with the same prompt formats, decoding strategies, or number of shots? Without full control of experimental settings, the claim that MRO narrows the gap with AR models is not strongly substantiated\n- MRO introduces considerable computational complexity due to repeated model evaluations, especially for the token verification reward that requires re-inference of multiple masked variants. This paper does not quantify the actual computational overhead compared to baseline models. The reliance on test-time scaling and beam search with reward evaluation also raises concerns about scalability to real-world applications.\n- The paper only compare the performance with AR models. They didn't compare MRO with other RL methods.\n- The study only looks at STEM tasks; it doesn't assess how well the suggested MRO apply to tasks involving common sense, following instructions, and so on.\n- The experiments are limited to models up to 8B parameters, leaving the effectiveness of MRO on larger models (e.g., 13B, 30B, or more) untested.\n- The paper only consider the experiments with few-shots. It is unclear about the performance of zero-shot.\n\nQUESTIONS:\nsee weakness\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors point out that diffusion language models face the problem of missing intra-sequence and inter-sequence correlations in reasoning tasks. They propose the Multi-Reward Optimization (MRO) framework, introducing three types of rewards—Token Verification, external language-model perplexity, and answer-level reasoning correctness—and present the Step-wise Group Reward Optimization method to reduce variance. Experiments conducted in the three modes of test-time scaling, rejection sampling, and reinforcement learning confirm the effectiveness of the proposed approach.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n1. The paper is the first to investigate how intra-sequence correlation affects reasoning performance in diffusion language models—a crucial issue, especially when the number of sampling steps is small.\n\n## Weakness\n1. The mathematical definition of the x-axis in Figure 1 is never specified.\n\n2. Equation (3) appears incorrect. The left-hand side of the first equality is the marginal distribution of $r_0$, whereas the right-hand side is the joint distribution of $r_0, r_1, \\dots, r_{t-1}$. The second equality is also wrong; please refer to Eq. (7) in [1] for the standard backward-transition formula for masked diffusion models.\n\n3. Token Verification Reward lacks theoretical grounding. The paper introduces this reward but does not explain why it should faithfully measure—or improve—intra-sequence correlation.\n\n4. The authors treat the final-answer correctness (Eq. 6) as an inter-sequence reward, yet final-answer accuracy also depends on intra-sequence correlation.\n\n5. Computational overhead is unreported. The comparison between Figure 2 and Figure 5 is not rigorous because MRO introduces substantial extra computation. The authors should quantify the additional wall-clock time or FLOPs required.\n\n[1] Sahoo et al. Simple and Effective Masked Diffusion Language Models. NeurIPS 2024.\n\nQUESTIONS:\n1. Section 5.3 (“Test-time Scaling”) is puzzling. Lines 304–305 say that policy-gradient updates are carried out during “test time.” If back-propagation is involved, can this still be called test-time scaling? Moreover, a detailed explanation of how the policy gradients are computed would greatly aid reader understanding.\n\n2. When evaluating the LLaDA-Instruct baseline, did you enable semi-autoregressive sampling? If so, what block length did you use?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-29"}

comments: SUMMARY:
This paper proposes Multi-Reward Optimization (MRO), a novel framework designed to enhance the reasoning capabilities of Diffusion Language Models (DLMs). The key motivation is that the parallel token generation process in DLMs lacks token correlation, which hampers performance on reasoning tasks. To address this, the authors introduce intra- and inter-sequence correlation rewards and optimize them using test-time scaling, rejection sampling, and reinforcement learning.

STRENGTHS AND WEAKNESSES:
Pros:
1. The paper studied an important problem. The authors address a key limitation of DLMs, namely the lack of token correlation, by designing multiple reward signals tailored to different aspects of reasoning. 
2. Extensive experiments across five reasoning benchmarks (e.g., GSM8K, MATH500, GPQA) demonstrate that MRO substantially improves performance and sample efficiency.
3. The definition of the problem is clear. The authors pinpoint token-correlation deficiencies as a root cause of poor DLM reasoning.

Cons:
1. The reliance on repeated forward passes (especially in test-time scaling and rejection sampling) increases inference and training overhead. Can the authors comment on this?
2. Can the authors also report the results on long-reasoning benchmarks such as AIME and AMC? This will help the readers to better judge the effectiveness of the proposed method and build future works on top of it.

QUESTIONS:
See weakness above.

LIMITATIONS:
Yes.
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-07-01

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a novel approach to improve the reasoning capabilities of diffusion language models (DLMs) by optimizing token correlations through a Multi-Reward Optimization (MRO) framework. \n\nThe key contributions of this paper are as follows:\n\n1.\tIntra-sequence and Inter-sequence correlation: The authors define two types of token correlations: intra-sequence correlation and inter-sequence correlation, to capture token correlations across denoising steps.\n\n2.\tMulti-Reward Optimization (MRO) framework: The authors propose an MRO approach that enables DLMs to generate reasoning path with an emphasis on token correlation.\n\n3.\tStep-wise Group Reward Optimization (SGRO): The authors introduce SGRO in the optimization phase to reduce reward variance, leading to more stable optimization.\n\n4.\tComprehensive evaluation: Through extensive experiments on various reasoning tasks, the authors demonstrate that MRO significantly improves the reasoning performance of DLMs.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1.\tThorough analysis: The authors provide a thorough analysis of the limitations of current DLMs in reasoning tasks, specifically focusing on the lack of token correlation across denoising steps. \n\n2.\tDetailed framework explanations: The authors provide detailed explanations of the proposed MRO framework, including the design of intra-sequence, inter-sequence correlation rewards and SGRO. \n\n3.\tComprehensive experiments: The paper includes extensive experiments on multiple reasoning tasks across different models and settings. The results consistently demonstrate the effectiveness of the proposed MRO framework.\n\nWeaknesses\n\n1.\tLack of clarity in the Introduction: The main contributions of the paper are not clearly articulated in the introduction.\n\n2.\tInadequate figures and pseudocodes: The paper lacks a comprehensive figure illustrating the MRO framework workflow, which would help readers in understanding all components of MRO. Additionally, the inclusion of brief pseudocode for algorithms such as SGRO is recommended.\n\n3.\tTypos and notation abuse: There are some typos and inconsistent notations throughout the paper, which hinder clarity.\n\n4.\tAmbiguous reward settings: Some reward settings appear unreasonable; further details can be found in the Question section.\n\nQUESTIONS:\n1.\tIn Equation 5, the rationale for setting the upper bound of the perplexity reward to 100 is unclear. Have your test alternative parameters, and if so, how do they impact the RL training process?\n2.\tIn Equation 6, considering only r_0 in the inter-sequence correlation rewards is questionable. In a Markov setting, this reward should be able to measure the correlation between consecutive state pairs (e.g., s_t and s_{t-1}), and all consecutive state pairs should be considered in a complete reasoning trajectory.\n3.\tIn Equation 8, each separate reward has a different range, resulting in varying weights in the total reward. Why is reward normalization not utilized to address this problem?\n1.\tTypos and notation abuse:\n(1)\tLine 159: the response x_0 should be r_0\n(2)\tLine 168: p_{theta}(x_0|x_t) is for the pretraining stage, not the SFT stage\n(3)\tLine 253: the “<” in the variance inequality should be “>”\n(4)\tFor the figure 7 in the appendix, the title and the x-axis label should be “group size”, not “temperature”\n\nLIMITATIONS:\nYES", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper starts with an empirical analysis that identifies the critical issue of why Diffusion Language Models struggle with reasoning tasks due to their inability to capture dependencies between tokens generated in different denoising steps. This paper observes two critical types of token correlations–intra-sequence (within a step) and inter-sequence (across steps)-as essential for coherent and accurate reasoning. To address this, the authors propose a Multi-Reward Optimization (MRO) framework that: (1) introduces multiple reward signals (e.g., token verification, perplexity, reasoning correctness and format adherence) to guide token generation; (2) models the denoising process as a Markov Decision Process (MDP), enabling reinforcement learning-based optimization; and (3) incorporates Step-wise Group Reward Optimization (SGRO) to reduce reward variance and stabilize training. The paper then conducts experiments across various benchmarks demonstrating that MRO significantly enhances reasoning accuracy, while also reducing decoding time by enabling fewer denoising steps.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. The analysis and the identification of the effect of token correlation are insightful and could lead to potential research in the future.\n2. The proposed reward design follows the empirical analysis and resolves reward variance.\n3. The proposed method improves the performance consistently across 4 benchmarks and especially improves accuracy on lower denoising steps.\n4. Overall, the writing is clear and easy to follow.\n\nWeakness:\n1. The paper does not provide any efficiency analysis on the decoding time. The proposed method requires the calculation of multiple rewards and policy gradients, which poses a challenge due to latency issues.\n2. Although the paper claims that the method could especially benefit reasoning problems with CoTs, it generally mitigates the decoding problem of DLMs and could benefit from further verification on other tasks.\n3. The paper could benefit from adding more baselines that require training on the dataset in order to better verify the proposed method.\n\nQUESTIONS:\n1. Can the authors provide a latency analysis of the proposed method compared with vanilla decoding?\n2. Could the authors provide analysis on the intra-/inter- sequence correlations with and without the proposed method?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-03"}