MRO: Enhancing Reasoning in Diffusion Language Models via Multi-Reward Optimization
Authors: Chenglong Wang
Yang Gan
Hang Zhou
Chi Hu
Yongyu Mu
Kai Song
MuRun Yang
Bei Li
Chunliang Zhang
Tongran Liu
JingBo Zhu
Zhengtao Yu
Tong Xiao
Venue: NeurIPS 2025 poster ()
Keywords: Generative Model
Language Model
Reasoning
Diffusion Language Model
Reward Optimization
Primary area: deep_learning
TL;DR: We propose the Multi-Reward Optimization (MRO) approach, which enhances token correlation during the denoising process in diffusion language models, improving reasoning performance and sampling efficiency.

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.

MRO: Enhancing Reasoning in Diffusion Language Models via Multi-Reward Optimization
Chenglong Wang'
Hang Zhou' Bei Li Jingbo Zhu 1.3 Zhengtao Yu
Yang Gan
Chi Hu
Yongyu Mu²
Kai Song
Murun Yang²
Chunliang Zhang
Tongran Liu
Tong Xiao 1.3t
School of Computer Science and Engineering, Northeastern University. Shenyang, China ByteeDance Niu'Trans Research, Shenyang, China
*CAS Key Laboratory of Behavioral Science, Institute of Psychology, CAS, Beijing, China Kunmming University of Science and Technology (xiaotong, zhujingbo)@mail neu, edu.
(clwang1119, 2zhu8250]88gail_ com
Abstract
Recent advances in diffusion language models (DLMs) have presented promising alternative to traditional autoregressive large language models (LLMs). However, DLMs still lag behind LLMs 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-Rewan Optimization (MRO) approach, which encourages DLMs to consider the token correlation during the denoising pro- cess. More specifically, our MRO approach leverages test-time scaling, rejection sampling, and reinforcement learning to directly optimize token correlation with multiple elaborate rewards. Furthermore, we introduce a step-wise group reward optimization approach mitigate reward variance during the reward optimization. Through extensive experiments, we demonstrate that MRO not only improves reasoning performance but also achieves significant denoising speedups while maintaining high performance across reasoning tasks.
Introduction
Large language models (LLMs) have recently made remarkable strides, profoundly impacting the
entire field artificial intelligence. Almost all of these models share a typical recipe: learn a model
that maximizes data likelihoods using an autoregressive (AR) paradigm 2 3 This AR paradigm
enables L.LMs to perform exceptionally well in a wide range of downstream tasks, particularly
demonstrating impressive abilities solve complex reasoning problems using chain of thought (CoT)
5 7 8) Models such as OpenAl's ol [9] Qwen2.5 [10), and DeepSeek-RI [I] the forefront of this progress. Despite its empirical success, the AR paradigm has inherent limitations. During decoding, AR models generate tokens in left-to-right, token-by-token manner, which constrains both the efficiency and flexibility of generation [3]. Furthermore, the generation process heavily depends on previously generated tokens, which often leads to error accumulation (also known as exposure bias)
methods 4
'Wok was done when Yang Gan was interning at ByteDance. Corresponding author.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

To overcome these limitations, recent research has explored alternative approaches to developing language models [66 18). Among these, diffusion language models (DLMs) have emerged as a promising and competitive direction 13 19 20 21. Unlike traditional AR L.Mss, DLMs generate sequences by iteratively predicting multiple masked tokens in parallel at each intermediate denoising step, enabling bidirectional and controllable generation while improving sampling efficiency. As part of ongoing efforts in this area, substantial progress has been made in pretraining and scaling DLMs, which have demonstrated strong capabilities in text generation tasks [22 [33] More recently. Ye et al. 24 have proved the potential of DLMs in solving planning tasks, including Countdown and Sudoku.
Despite their potential, DLMs still underperform AR models by a large margin in reasoning tasks, as also evidenced in Table IT As a result, developing DLMs with strong reasoning capabilities remains an open research problem. To address this, DoT [25] introduces the CoT technique into the denoise process of the DLM. Instead of denoising the entire sequence at each step, DoT distributes the reasoning process across denoising steps, allowing each step to correspond to distinct reasoning stage. However, this approach still partially relies on the AR paradigm, making susceptible to error accumulation during reasoning. Additionally, dl (26) improves the reasoning performance of DLMs by fine-tuning them with a policy gradient algorithm.
All of the existing efforts, however, overlook a fundamental factor con- tributing to the poor reasoning perfor- mance: the independent generation of multiple masked tokens across denois- ing steps fails to capture the dependen- cies among them in CoT--ttyle reason- ing. We argue that this independence introduces inconsistencies in the rea- soning path, ultimately leading to in- correct outcomes, as the generated to- kens may lack correlation with one another. Specifically,
Line of Best Fit
95% Confidence Interval
Stundardized Aucuracc
Stunsholioed Avcuracy
Standardized lntre sequence Correlation
Standardized Inter sequenes Correlation
Figure I: Visualization of the relationship between token cor- relation and reasoning accuracy. The results show that higher intra- or inter-sequence correlation tends to yield higher rea- soning performance.
define two types of token correlation in DLMs: intra-sequence correlation and inter- sequence correlation. Intra-sequence correlation measures the degree of dependency among tokens generated within single denoising step. In contrast, inter-sequence correlation captures the alignment between sequences generated at different denoising steps. Furthermore, we examine the relationship between the correlation and reasoning accuracy by performing 50 decoding runs on subset of GSMSK, as shown in Figure From the results, we observe that higher token correlation leads to more accurate reasoning outcomes. This finding aligns with the observation in Lightman al., 27)ss work, that suggests that inconsistencies within the reasoning path often lead to incorrect outcomes.
How can equip DLMs with enhanced token correlation in reasoning? One straightforward approach is to scale the training data and model size [22 23 enabling DLMs to better capture the correlation between tokens. While this approach does not directly optimize the correlation for DLMs, scaling the data and model size could potentially enhance by improving the accuracy of masked token predictions. Unfortunately, such scaling often comes with significantly increased requirements for labeled data and computational resources. In this paper, we propose an approach that enhances the token correlation without these burdens. Our approach, enables DLMS to generate reasoning paths with an emphasis on token correlation.
Optimization (MRO),
To develop MRO, we begin by designing multiple rewards and incorporating test-time scaling, rejection sampling, and reinforcement learning to optimize these rewards simultancously. Intuitively, the optimization objective during DLM training is predict masked tokens, but the decoding joint distribution is parameterized as a product of token-wise independent distributions. To this end, MRO enhances the correlation through multi-reward optimization mechanism, capturing dependencies among tokens and bridging this gap. Furthermore, we find that MRO encounters the issue of reward variance when handling multiple rewards simultaneously. To address this, we introduce step-wise group reward optimization (namely SGRO) approach to optimizing DLMs with rewards. We provide theoretical analysis of how our SGRO reduces the reward variance from potential-baser reward shaping perspective, as shown in Appendix B

---

Through extensive experiments on MRO with cest-time scaling, rejection sampling, and reinforcement learning, we show that multi-reward optimization is effective. For example, MRO yields an average improvement of 3 points across various reasoning tasks with test-time scaling. As another bonus, MRO enables DLMSs to generate high-quality reasoning paths with fewer denoising steps, thereby accelerating the overall decoding process.
2
Related Work
2.1 Diffusion Language Models
Building previous advancements in diffusion models for image generation 28
29 30 recent
efforts have attempted to explore the application of diffusion models for text generation tasks. For example, text-based continuous diffusion models 31 32 have introduced an embedding function to map discrete text into a continuous space. Furthermore, to accommodate the inherently discrete nature of text, researchers have explored discrete diffusion models for language modeling [33]. Despite their significant success, challenges remain, particularly with scalability. More recently, the masked diffusion model 21 22 has emerged as a breakthrough instance of the discrete diffusion model. For example, DiffuLLaMA 3 extended the masked diffusion model by initializing it with parameters from LLaMA models. In another line of research, based on these established DLMs, researchers attempted to optimize the denoising process to further enhance them. This includes integrating the CoT approach 25 and AR paradigm 18] during the denoising process. Excitingly, in the context of scaling DLMs, LLaDA (23] and Dream (24) showed that it is possible to scale DLMs to the 8B and 7B sizes, respectively, with pre-training and SFT training, achieving text generation capabilities comparable to AR LLMs. Motivated by the recent success of DeepSeek-RI [ΑI, which leverages reinforcement learning to scale the reasoning capability, Zhao et al. (26) proved the effectiveness of applying reinforcement learning to DLMs for improving reasoning performance. Although previous work improves the performance of DLMs on various text generation tasks, they often overlook a fundamental limitation: the inherent parallel and independent generation of tokens in DLMs, which leads to weak token correlations across denoising steps. This limitation can negatively impact tasks that require high consistency, such as CoT-style reasoning, where coherence and alignment of content are critical. Researchers have been aware of this (35), but it is still rare to see studies on this issue.
2.2 Reward Optimization for Language Models
Reward optimization is a fundamental concept in reinforcement learning that helps agents make better decisions. In the context of LLMs, reward optimization can guide the generated content to better align with human preferences. A widely used approach for reward optimization is reinforcement learning from human feedback [5] 36 37 (38), where a reward model is trained based on human feedback, and the LLM is optimized relative to this reward using algorithms like proximal policy optimization [39 and rejection sampling 44). This area has seen significant progress, with recent examples including models developed by OpenAl [9] and DeepSeek [1I, which employ reward optimization approaches to teach models to human-like thinking and complex reasoning. Building on this foundation, there is a growing body of work exploring how design reward functions that
enable LLMs to perform more effectively across a wide range of tasks [1 42
3 Preliminaries
In this section, we outline some basic concepts and notation of
diffusion language models.
3.1 Training Diffusion Language Models
Discrete diffusion models 33 have emerged as promising approach for language modeling, providing an alternative to traditional AR models. One of the most popular methods in this context the "masked" approach, also known as masked diffusion models [1] 43 22). More excitingly, (3 scales masked diffusion models to a 7B parameter model and integrates unsupervised pre-training with instruction fine-tuning during the training of DLMs, achieving impressive text generation performance. Specifically, at the pre-training stage, the training objective involves minimizing a cross-entropy loss computed only for masked tokens. Given an unlabeled training sequence zo.
3

---

time step with a masking ratio sampled uniformly from [0, 1], and the sequence Iy is obtained by masking each token independently with probability t. The loss function
L 1 - log Pro(zc) mx)
p = - -Bo
Σ
(I)
where L is the length of the sequence FD. z. is the i-th token, M is the masked token, and 1[:] is the indicator function. Here, (ro(-) denotes the mask predictor, and 0 is the set of its parameters. is typically parameterized by stacked Transformer model without a causal mask. This loss function encourages the model to accurately predict the original tokens from the masked sequence. The basic idea for pre-training DLMs is similar to that of pre-training AR LLMs. Both aim capture linguistic knowledge from large amounts of unsupervised data, whereas AR LLMs rely on next-token prediction, and DLMs use masked token prediction, a strategy proven effective in encoder-only language models, such as BERT 44] and RoBERTa [45).
After pre-training, we can employ labeled data to enhance DLMS' ability to follow instructions with supervised fine-tuning (SFT). This phase involves training the model on labeled paired data (po. ro). where po is the prompt, such as "If + 5 = 12, what is the value of x?" and ro is the corresponding response, such as "The answer is 7". During SFT, the prompt Po is kept unchanged, and tokens the response to are masked independently. The masked response re then fed into the pre-trained mask predictor to compute the loss:
L Σ
Lun (0)
-Bor.ooroorr
1|r
log Pro(ró po.
(2)
where I,, is the length of ro- The training stage is similar to the pre-training, with the key difference being that the model now learns to predict the response tokens conditioned on the given prompt.
3.2 Applying Diffusion Language Models
fundamental application of DLMs is cloze (fill-n-thhebblane tasks, which align with the training objective. However, since the ultimate goal is to replace AR LLMs, key expectation is that DLMs should also be capable of generating textual sequences. To achieve this, we typically use masked token prediction to simulate the denoising step. More specifically, given an unseen prompt po. the model starts with a fully masked response rr and iteratively predicts the masked tokens to fully reconstruct the response To over T' time steps. The process also involves re-masking a fraction of the predicted tokens at each step to ensure that the reverse process aligns with the forward process.
Multi-Reward Optimization for Diffusion Language Models
4.1 Problem Definition
During the DLM decoding process, we use the learned denoising distribution Præ(-) to construct the decoding distribution Qe(rrr1,.... .To] Po.rr), where [/T-1.... To] denotes the sequence of intermediate responses progressively refined through the itcrative denoising steps. This process predicts the masked tokens across T denoising iterations and can be expressed
Lr
Qo(Ir-I..".. rol po. PT) = II Uo(ri- Po. r()
M M/Pro(ri- Po.re) (3)
inter sequence correlation
intra sequeecce correaatinn
where Ua(-) denotes the denoising distribution that jointly predicts all masked tokens at step - 1 conditioned on the current partially denoised sequence т₁ and the initial prompt Pu- While this approach effectively constructs the decoding distribution Qo(-) through iterative denoising, it fails to capture the dependencies between tokens across denoising steps. Consequently, Prø(-) cannot perfectly approximate the true decoding distribution, leading to weakened token-levee correlations and inconsistencies within the generated sequence. These inconsistencies can be particularly harmful in reasoning-intensie tasks (e.g., CoT-style reasoning) where precision and coherence are essential. As shown Eq. 3 we define these correlations as intra-sequence correlation and inter-sequence correlation, with their definitions outlined below:

---

Intra-sequence Correlation. This refers to the dependencies among tokens generated within single denoising step. In reasoning tasks, it measures how well these tokens collaborate to form coherent and logical segment of the reasoning path, ensuring fluency and consistency.
Inter-sequence Correlation, It captures the alignment between tokens generated at different denoising steps. This correlation ensures that the sequence generated at time step is consistent with the sequence generated at time step t, leading to coherent reasoning path.
4.2
Multi-Reward Optimization
We aim to design multiple rewards that optimize token correlations during the DLM decoding process. We begin by describing separate reward designs targeting intra-sequence and inter-sequence correlations, respectively, and then explain how combine these rewards in the optimization process.
4.2.1 Intra-sequence Correlation Rewards
For token correlation within single sequence at a denoising step, we evaluate it from two aspects. First, during decoding, we assess whether the tokens predicted in parallel at one denoising step exhibit strong correlations. Second, based on these generated tokens, we check whether the current sequence forms coherent and readable response.
Token Verification Reward. To address the first aspect, we design token verification reward (denoted as R" This basic idea is that, at each denoising step, after predicting the masked tokens, we re-enter the model to verify the correlations between the generated tokens. Specifically, consider an instance with N masked tokens: at the 1-th denoising step, given the prompt Po and the intermediate response rı, the DLM predicts the masked tokens (r(-). 1-11--........
""")) ). which together form rz-1, where {m,, m22..... my) }o represents the indices of the masked token positions. We then use the
token verification method compute the reward Rd for this time step. denoted as R
N Σ Pro(r?" m Po, 1t-1/r(1) N
R =
where 1r-1/r(" denotes the response where only the token r( is masked from 71-1- Furthermore, provide a theoretical proof that the token verification reward enhances intra-sequence token correlations from the perspective of mutual information (see Property 3 in Appendix BB. In practice, however, calculating this reward introduces additional computational overhead. To mitigate this issue, we adopt three optimization strategies as follows. First, for each denoising step. we randomly sample a subset of tokens for verification instead of checking all tokens. Second, we leverage GPU parallelism to compute the probabilities of different masked tokens simultaneously, thereby reducing time cost. Third, we introduce a step-wise group optimization strategy further lower the reward computation overhead, as described in Section 4.3
Perplexity Reward. For the second aspect, considering perplexity widely used as a measure of textual sequence consistency and readability [66.. we design perplexity reward (denoted as Reel).. Specifically, at the t-th denoising step, we use a pre-trained AR language model to compute the perplexity of the response T(-1 generated at that step, and then compute R." as follows
max ( (Cppl PPL(
Reol
(5)
Fppl
where PPL(-) denotes the perplexity computation function, Oppi is fixed upper bound constant controlling the maximum reward value, and Fppi is a scaling factor used to ensure that the range of R is comparable to that of other rewards. Here, we use 1mpp1 to implement it.
4.2.2 Inter-sequence Correlation Rewards
For token correlation across different sequences during the denoising steps, we primarily assess whether the combination of these steps results in a high-quality response. Specifically, for reasoning
httt::://gttaab..co/aaaa////////
5

---

tasks, we build on recent work [TT to evaluate the quality of CoT-style reasoning generated by DLMs, focusing on both CoT format and accuracy. Given a prompt Po. once we obtain ro, the corresponding reward Rg is computed as follows
2, if the ro follows the required format and its answer is correct
R =
1, if the ro follows the required format but its answer is incorrect
(6)
0, if the ro does not follow the required format
we define a format reward, which checks whether the CoT strictly appears within the "<think> </think>" tag and whether the answer strictly appears within the "<answer> </answer>" tag. Note that our inter- and intra-sequence correlation rewards are not entirely in- dependent. A stronger intra-sequence correlation reward, which enforces local token-level coherence, can indirectly enhance the inter-sequence reward by improving the overall fluency and consistency of each intermediate response. Nevertheless, the two rewards focus on different aspects of correlation: the intra-sequence reward emphasizes token-level dependencies within single response, while the inter-sequence reward evaluates the global consistency and task-level quality across denoising steps.
Following M.
4.2.3 Reward Optimization via Markov Decision Process
To achieve the reward optimization, we model the iterative denoising process of DLMs as a Markov Decision Process (MDP) 47 In this formulation, each denoising step corresponds to an action, and the goal is to optimize the rewards at each step by considering both intra-sequence and inter-sequence token correlations. Formally, we define the key MDP components as follows:
St A (Pos Tt, (), ag rr-1: To(uu (8))= _ Up(ri-1 Po.rr)
where s, is the state at the t-th step, which includes the current time step the given prompt Po, and the generated tokens at step rt. ae represents the action taken at step t, which corresponds to the predicted response Tı-1- (a | st) denotes the policy of the DLM, which predicts the tokens at the previous time step /. - 1 based on the current tokens r, and the prompt po- Based on this MDP formulation, can easily apply optimization algorithms like policy gradient [48 maximize an arbitrary reward (denoted as R((e, at)) during the denoising process.
Here, we construct R(s;, az) by incorporating intra-sequence and inter-sequence correlation rewards. Specifically, we define the intra-sequence correlation rewards in a potential-based reshaping reward [49]. At each denoising step f, after generating the tokens Tı-1, we shape the reward by considering both the token verification reward R and the perplexity reward Reel, Additionally, we treat the inter- seter-sequence correlation reward R as delayed reward, as the alignment between sequences generated across different denoising steps becomes evident only after the entire sequence is generated. Based on this reward allocation scheme, we can define the R(st, at) by
R R + R.P!
RePl ift=
R(s,, az)
(8)
otherwise
The benefit of this formulation is that if we use a standard sampling strategy with (r | po. rt) as described in Eq. 3 this policy can optimize this reward and adjust the parameter 0 when using algorithms like policy gradient. This enables DLMs to directly optimize token correlations, which are typically absent during pre-training and SFT training processes.
4.3 Step-wise Group Reward Optimization
As mentioned earlier, we combine the rewards using a potential- based reshaping approach. Let ((-) denote the potential function. We can express the difference as XPP(s ((:)) = R + R,,, where is the discount factor. Thus, we have R(s,, ax) = A(s,, az) + (0( - ((s)). where R(s,, az) denotes the original sparse reward. While potential- based reshaping
effectively optimize token correlations, (50) shows that, in the case of longer decision trajectories in MDPs, this approach can lead to higher reward variance during the optimization of parameters. We establish similar result for our nelti-reward optimization in Property 1.
Property 1. Under potential-basee reward shaping, the expected reward remains the same, i.e. E[R(s, a)] = [[ a)), but Var(Ri, s, a)) is higher than Var Var[R(s, a)).

---

The proof can be found in Appendix B This conclusion highlights a limitation in our reward combi- nation method. To address this limitation, we propose the Step-wise Group Reward Optimization (SGRO) approach. This approach groups w denoising steps together, with each group providing reshaped reward, thereby replacing the computation of the intra-sequence correlation reward each individual denoising step. By reducing the number of reshaped reward computations, SGRO helps lower reward variance. Specifically, we further prove that Var(R a)) Var R((*) (s, a)), where
R())()) denotes the grouped reward after applying SGRO (see Property 2 in Appendix B). In addition to variance reduction, find that SGRO also decreases the frequency of reward evaluations, thus
reducing the computational overhead of reward estimation, described in Appendix D
5
Experiments
We evaluate the effectiveness of our multi-reward optimization (MRO) approach with various opti- mization algorithms, including test-time scaling, rejection sampling, and reinforcement learning.
5.1
Setups
We conducted our experiments using the LL model [23] a state-of-the-art open- source DLM that employs fully masked-based training approach for both pre-training and SFT training. Additionally, we performed experiments on the MaADA-8B-Instructe DLM derived by fine-tuning the pre-trained LLaDA-8B model with the sl training dataset (51). For evaluation, we considered five reasoning benchmarks across three categories: (1) Mathematical reasoning: GSM8K and MATH500; (2) Scientific reasoning: GPQA, which focuses on biology, physics, and chemistry reasoning; (3) Logical reasoning: 4x4 Sudoku and the Countdown task with
model, reasoning
numbers. More experimental details can be found in Appendix C
5.2 Baselines
We compared our MRO with several strong baselines: the LLaDA-8B LaaDA-8B-Instruct and LLaDA-8B -Isttrrcct- s/ models, which are two DLMs that have not undergone reward optimization (denoted as LLaDA and LLaDA-sI) and AR LLMs, including serval open-source AR LLMs with around 8B parameters, such as LLLMA-3-8B-Instrucee Ins-urttutt [52), Mistral-7B-Insttuus
5 Deepseek-LLMM77b- [1] and
Qwenn.557B-IInstruct 10
5.3 Test-time Scaling
In this subsection, we examine whether our MRO enables DLMs to generate accurate reasoning outcomes. To achieve this, we perform test-time scaling to evaluate the effectiveness of MRO.
Task Setup. As illustrated in Figure 2 we per- formed beam search to explore better decoding results. At the t-th denoising step, we generated k different responses frr,1, P2,2, ... rt,k) } using temperature-baat bassed sampling, where & was set to 4. We then computed the reward for each response us- ing Eq. and selected the response with the high- reward the final result for that step. Note that in the test-time scaling experiment, we did have access to the ground outcome when com- puting RJ. Therefore, we employed a majority voting approach to obtain pseudo-ground out- come, which was then used to compute Rg. The temperature was set to 0.25, and we present the performance with other temperature values in Fig-
marked token
non-maakked token
300
20c
1= T-1
Figure 2: Illustration of the test-time scaling
procedure in DLMs. The bolded line denotes that the response obtains the highest reward and is selected the final output.
We also tested different response lengths, including 64, 128, 256, and 512. For each setting, we set the number of denoising steps half the generated response length. Additionally, we applied our SGRO, setting w to 32, meaning that beam search with MRO was applied every 32 steps during
ure 5
7

---

LLaDA
LLaDA MRO
LLaDA
LLaDA MRO
80.0
35.0
Accuracy (56
70.0
Accuracy (56)
35.0
Accuracy (96
30.0
60.0
30.0
25.0
50.0
25.0
15.0
30.0
10.0
64
128
256
512
64
128
256
512
64
128
256
512
Response Length
Response Length
Response Length
GSM8K (4-shot)
(b) MATH500 (4-shot)
GPQA (5-shot)
Figure 3: Accuracies (%) three reasoning tasks (GSM8K, MATH500, and GPQA) with test-time scaling, showing results for different response lengths.
the beam search process. For evaluation, we adhered the official evaluation procedure provided by LLaDA, which employs semi-autoregressive sampling by default [[2]. The block lengths were set 8, 64, and 64 for GSM8K, MATH500, and GPQA, respectively. The final accuracy metrics were obtained using the eeeeevvlaaionaaareeeee toolkit.
Results. As shown in Figure MRO significantly improves reasoning performance across different response lengths and models. Specifically, when comparing models with and without MRO, we observe notable improvements in accuracy. Firstly, the token correlation rewards we introduced proven to be effective. For example, in the GSM8K task, the LLaDA model with MRO achieves an accuracy of 82.6% at response length of 512, whereas the LLaDA model without MRO reaches only 79.4%. Similarly, for the MATH500 task. LLaDA + MRO improves from 34.4% to 38.00 38.0% response length of 512. Secondly, the results also confirm the claim made Section 4. that DLMs like LLaDA often miss important token correlations during decoding, which can limit their performance. These observations confirm that the proposed reward design effectively enhances the model's 's reasoning capability, providing a useful insight: the designed reward could be leveraged to further optimize the DLM via parameter adjustment using techniques like reinforcement learning or rejection sampling. We validate this insight in the following section.
5.4 Rejection Sampling
While test-time scaling effectively implements MRO, the decoding time is significantly scaled. Ideally, we aim optimize the parameters of this DLM to consider the token correlation during the decoding stage. Here, we employ rejection sampling to achieve this goal.
Task Setup. We used beam search and SGRO, as described in Section 5.3 sequence fr,,,--1...... fn] that yields high rewards. This sequence then used for SFT training of our DL.. The corresponding loss function can be given by
to obtain a sampled
Lp Σ
(T)( = -Eo, Por.po..TT.. 1. fo
i = M] log Pro po, ft)
where Ly denotes the length of the response f. During our training process, experimented with setting the beam size for each expansion step to 2 and 4. Furthermore, we found that computing the gradients for all denoising steps simultancously consumed large amount of GPU memory. To address this issue, we developed two strategies. First, we used gradient accumulation, where we computed the gradients for each individual step and accumulated them until the entire sequence was processed, after which backpropagation was performed. Second, instead of optimizing over the entire sequence, we selected segment of denoising to optimize. The basic idea that token correlations between different denoising steps can be captured not only by using the full denoising sequence, but also through combinations of smaller subsets.

---

GSM8K
MATH1500
GPQA
Countdown
Sudoku
Model/Length
128
256
512
128
256
512
128
256
512
64
128
64
18.4
25.9
3.7
0.0
Histral-7B-Instrual
36.2
13.1
24.7
1.2
0.0
Doopsck-LLLM-7--CMae
17.4
6.0
19.5
8.5
0.0
Qwen2.5-7B-Instee
85.4
41.1
36.4
6.2
21.0
LLaDA
74.3
78.8
79.4
28.0
33.2
34.4
22.2
29.2
30.3
13.8
14.1
11.2
LLaDA-MRO-2
75.1
80.0
82.5
30.0
34.0
35.6
25.3
30.0
32.8
16.7
16.9
16.0
LLLDA-MRO-4
76.9
79.6
82.6
31.0
34.2
36.2
26.3
32.1
34.3
21.4
22.0
17.2
LLaDA-sl
70.7
76.4
78.7
25.0
26.0
28.8
21.7
25.3
28.3
10.2
12.4
8.4
LLaDA-sI-MRO-2
71.6
76.7
79.6
28.2
27.6
30.0
22,7
27.3
30.3
14,5
15.2
12.0
LLaDA-sI-MRO-4
73.3
77.7
80.1
27.8
28.0
29.0
24.2
29.3
32.8
17.3
17.8
15.2
Accuracies (%) on mathematical reasoning, scientific reasoning, and logical reasoning tasks. The best performance in each group is highlighted in bold. The suffixes "-2" and "-4" indicate the beam size & used during rejection sampling, set to 2 and 4, respectively. Note that for GSM8K and MATH500, we use 4-shot, while for GPQA, Countdown, and Sudoku, we use 5-shot. f indicates that the results for GSM8K and GPQA are taken from 23
Table 1:
In the experiments presented in this subsection, for given sampled sequence, randomly selected two consecutive denoising steps for optimization. The block length was set to 64 for Countdown and Sudoku. Further training details can be found in Appendix
Results. The experimental results listed in Table 1 From the results, we observe that rejection sampling with MRO can effectively enhance the reasoning capabilities of models across various types of reasoning tasks. This further validates that our MRO approach is able to adjust model parameters to better capture token correlations during the decoding process, leading to improved reasoning performance. We also observe that our approach can yield particularly significant improvements in logical reasoning tasks, such as Countdown and Sudoku. Notably, based on the LLaDA model, MRO-4 achieves an improvement of +7.9 points on the Countdown task when the response length is set to 128. This is because these tasks seem to rely more heavily on token correlations, which may explain the more pronounced performance gains. Excitingly, the results show that MRO helps DLMs approach the performance of the strong LLM, Qwen2.5-7BB in mathematical reasoning. For example, on the GSM8K task, our model with MRO achieves an accuracy of 82.6%, which is very close to Qwen2.5-7BB 85.4%, indicating that MRO has the potential to narrow the performance gap in mathematical reasoning.
Additionally, we further compare the performance of MRO-2 and MRO-4. In most test cases, MRO-4 achieve better performance, indicating that expanding the search space by increasing the beam size allows the model to explore a broader range of denoising sequences. However, this comes with significant computational cost due the increased sampling requirements. This motivates future research into methods for reducing the computational burden associated with DLM training while still ensuring the ability to search for more optimal denoising sequences.
5.5 Analysis
Decoding Efficiency Analysis. We present experimental results compar- ing decoding efficiency under test- time scaling setting. Our results demonstrate that MRO is computa- tionally efficient, introducing no no- ticeable decoding burden while sig- nificantly improving generation accu- racy. Specifically, we provide a de- tailed comparison of decoding time and performance between vanilla decoding and MRO, as summa- rized in T Since test-time scaling inherently involves additional sampling, which introduces time overhead, we further compare against a confidence-based reward test-time scaling baseline (denoted as aDA-ASSS + CBR) that incurs the same sampling cost as MRO but omits reward computation. From the results, observe that although reward computation introduces modest
MATH500
GPQA
Method
Decoding Time Score
Decoding Time
Score
Vanilla (LLaDA)
0.35b=0.39h
33.2
0.16h
29.2
LLTDA-TTS + CBR
0.73h-0.81h
35.2
0.27h-0.34h
30.6
LLaDA-TTS + MRO
0.84b-0.85h
36.0
0.3lh-0.444
34.6
Table 2: Decoding time and performance comparison on MATH500 and GPQA benchmarks.

---

amount of additional latency, the overhead remains within reasonable range and is well justified by the substantial performance gains it yields. Moreover, compared with the confidence-bic reward baseline, we find that after excluding the sampling time, our reward design adds only minimal computational overhead while consistently delivering superior performance.
Quantitative Analysis of Intra- and Inter- sequence Correlations. We conduct a quantitative analysis of intra- and inter- sequence correlations to better understand our proposed MRO. Specifically, we ran- domly sampled 200 examples from the MATH500 and GPQA datasets and computed the corresponding intra- and inter-sequence reward scores across five decoding runs with different random seeds. As shown in Table MRO can consistently enhance both intra- and inter-sequence correlations. indicating that it strengthens token correlations within each denoising step and promotes more coherent transitions across steps.
MATHS00
GPQA
Model
intra-corr
inter-corr
intra-corr
inter-corr
LLaDA LLaDA-MRO
3.44:0.18
1.02:0.14
2.76=0.21 1.02=0.15 1.27=0013
3.79+0.16 16
1.58+0.12
3.34
Table 3: Intra- and inter-correlation analysis on MATH500 and GPQA benchmarks.
Performance with Different Denoising Step Sizes. With a fixed response length of 512, we further ex-
LLaDA
LLDDA MRO-4
plore the performance of LLaDA and LLDDA-MRO. 4 with different denoising step sizes. Specifically, we test denoising step sizes of (64, 128, 256, 512).
45.0
10.0
Speedup:
35.0
30.0
The results are summarized in Figure 4 From the results, we observe that our model outperforms the baseline across all denoising step sizes. Notably, we find that MRO helps accelerate the denoising process. For example, ALaDA-MRO- achieves performance comparable to the baseline with 256 denoising steps, even when using only 128 steps. We attribute this to the inter-sequence correlation rewards we designed during optimization, which encourage the model to consider the collaboration of different steps and facilitate a faster and more efficient denoising process.
25.0
20.0
15.0
10.0
5.0
64
128
256
384
512
Denoising Step Size
Figure 4: Performance of MRO with different denoising step sizes.
Performance with Different Temperature Settings. When optimizing DLMs with our MRO approach, we employ temperature-base diverse responses at each denoising step. Here, we investigate the performance of MRO with different temperature settings. Specifically, we perform test- time scaling on LLaDA and LLaDA-sI I with temper- ature coefficients of {0, 0.25, 0.5, 0.75, 1.0). We evaluate the results on the GSMSK task with a length of 256, shown in Figure 5 From these results, we can observe that a temperature coefficient of 0.25 achieves the best performance. Similar trends can be observed for other reasoning tasks as well. Therefore, in this work, we select a temperature of 0.25 for optimizing the denoising steps.
LLaDA + MRO
LLaDA MRO
sampling to generate
82.0
Accuracy
80.0
78.0
76.0
74.0
0.25
0.5
Temperature
Figure 5: Performance of MRO with different temperature settings.
6 Conclusion
In this paper, we introduced the Multi-Reward Optimization (MRO) approach to enhance reasoning in diffusson language models. Specifically, address the fundamental limitation of token correlation across denoising steps, we designed intra-sequence and inter-sequence correlation rewards, optimized through test-time scaling, rejection sampling, and reinforcement learning. Our experiments show that MRO can significantly improve reasoning, as evidenced by its performance across multiple reasoning benchmarks. Furthermore, we proposed a step-wise group reward optimization approach to tackle reward variance during the optimization process, ensuring efficient optimization. We have demonstrated the effectiveness of this approach through both theoretical analysis and experiments. Our codebase could be found at htttttt///// //tpt/ggtttttttt// com/ comm/waggclnlp/MR.
10

---

Acknowledgments
This work was supported in part by the National Science Foundation of China (Nos. 62276056 and U24A20334), the Yunnan Fundamental Research Projects NN.2002401B0000222 the Yunnan Science and Technology Major Project (No. 202502AD0880144) the Fundamental Research Funds for the Central Universities (Nos. N25BSS054 and N25BS5094), and the Program of Introducing Talents of Discipline to Universities, Plan 111 (No.B16009)
References
[1] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 5998-5999999 6008,, 2017.
[2] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Merbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33; Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
[3] Tong Xiao and Jingbo Zhu. Introduction to transformers: an nlp perspective. ArXiv preprint. 0bs/2331.1766 2023.
[4] Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel M. Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F. Christiano. Learning to summarize with human feedback. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020.
[5] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul F. Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28 December 2022, 2022.
[6] Jie Huang and Kevin Chen-Chuan Chang. Towards reasoning in large language models: A survey. In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki, editors, Findings of the Association for Computational Linguistics: ACL 2023, pages 1049-1065, Toronto, Canada, 2023. Association for Computational Linguistics.
[7] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed H. Chi, Quoc V. Le, and Denny Zhou. Chain-of-thought prompting elicits reasoning in large language models. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems 35; Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022.
[8] Denny Zhou, Nathanael Schärli, Le Hou, Jason Wei, Nathan Scales, Xuezhi Wang, Dale Schuurmans, Claire Cui, Olivier Bousquet, Quoe V. Le, and H. Chi. Least-to-most prompting enables complex reasoning in large language models. In The Eleventh International Conference
11

---

on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023, OpenReview.net, 2023.
[9] OpenAl. Learning to reason with 1lms, 2024.
[10] An Yang. Baosong Yang. Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, et al. Qwen2. 5 technical report. ArXiv preprint, abs/2412.15115, 2024.
[11] Daya Guo, Dejian Yang, Haowei Zhaagg Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang. Xiao Bi, et Deepseek-rl: Incentivizing reasoning capability in 2025.
Iims via reinforcement learning. ArXiv preprint,
[12] Emiel Hoogeboom, Alexey A. Gritsenko, Jasmijn Bastings, Ben Poole, Rianne van den Berg, and Tim Salimans. Autoregressive diffusion models. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.nee 2022.
[13] Ishaan Gulrajani and Tatsunori B. Hashimoto. Likelihood based diffusion language models. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Advances in Neural Information Processing Systems 36: Annual Conference Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10- 16, 2023, 2023.
[14] Florian Schmidt. Generalization in generation: A closer look exposure bias. In Alexandra Birch, Andrew Finch, Hiroaki Hayashi, Ioannis Konstas, Thang Luong, Graham Neubig, Yusuke Oda, and Katsuhito Sudoh, editors, Proceedings of the 3rd Workshop on Neural Generation and Translation, pages 157-167, Hong Kong, 2019. Association for Computational Linguistics.
[15] Kushal Arora, Layla EI Asri, Hareesh Bahuleyan, and Jackie Cheung. Why exposure bias matters: An imitation learning perspective error accumulation language generation. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio, editors, Findings of the Association for Computational Linguistics: ACI. 2022, pages 700-710, Dublin, Ireland, 2022. Association for Computational Linguistics.
[16] Jiatao Gu, James Bradbury, Caiming Xiong, Victor O. K. Li, and Richard Socher. Non- autoregressive neural machine translation. In 6th International Conference on Learning Rep- resentations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018
[17] Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz, and Gabriel Synnaeve. Better faster large language models via multi-token prediction. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenRe- view.net, 2024.
[18) Marianne Arriola, Aaron Gokaslan, Justin T Chiu, Zhihan Yang, Zhixuan Qi, Jiaqi Han, Subham Sekhar Sahoo, and Volodymyr Kuleshov. Block diffusion: Interpolating between autoregressive and diffusion language models. ArXiv preprint, abs/2503.00 009533 2025.
[19] Jiasheng Ye, Zaixiang Zheng, Yu Bao, Lihua Qian, and Quanquan Gu. Diffusion language models perform many tasks with scaling and instruction-finett finetuning. ArXiv preprint, abs/2308. 12219, 2023.
[20] Zhenghao Lin, Yeyun Gong, Yelong Shen, Tong Wu, Zhihao Fan, Chen Lin, Nan Duan, and Weizhu Chen. Text generation with diffusion language models: A pre-training approach with continuous paragraph denoise. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honoiulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pages 21051-21064, PMLR, 2023.
[21] Subham S. Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. Amir Globersons, Lester Mackey, Danielle Belgrave, Angela Fan, UI- rich Paquet, Jakub M. Tomczak, and Cheng Zhang, editors, Advances in Neural Information
12

---

Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurIPS 2024, Vancouver, BC. Canada, December 10 15. 2024, 2024.
[22] Shen Nie, Fengqi Zhu, Chao Du, Tianyu Pang. Qian Liu, Guangtao Zeng, Min Lin, and Chongxuan Li. Scaling up masked diffusion models on text. ArXiv preprint, abs/2410.18514, 2024.
[23] Shen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, Jun Zhou, Yankai Lin, Ji-Rong Wen, and Chongxuan Li. Large language diffusion models. ArXiv preprint, abs/2502.09992, 2025.
[24] Jiacheng Ye, Zhihui Xie, Lin Zheng. Jiahui Gao, Zirui Wu, Xin Jiang, Zhenguo Li, and Lingpeng Kong. Dream 7b: Diffusion large language models. ArXiv preprint, abs/2508. 15487, 2025.
[25] Jiacheng Ye, Shansan Gong, Liheng Chen, Lin Zheng, Jiahui Gao, Han Shi, Chuan Wu, Xin Jiang, Zhenguo Li, Wei Bi, et al. Diffusion of thoughts: Chain-of-thought reasoning in diffusion language models. ArXiv preprint, abs/2402.07754, 2024.
[26] Siyan Zhao, Devaansh Gupta, Qinqing Zheng, and Aditya Grover. dl: Scaling reasoning in diffusion large language models via reinforcement learning. ArXiv preprint, abs/2504. 12216, 2025.
[27] Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024.
(28] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan- Tien Lin, editors, Advances in Neural Information Processing Systems 33; Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
[29] Jonathan Ho, Chitwan Saharia, William Chan, David Fleet, Mohammad Norouzi, and Tim Salimans. Cascaded diffusion models for high fidelity image generation. J. Mach. Learn. Res., 23:47:1-47:33, 2022.
[30] Omer Bal-Tal, Lior Yariv, Yaron Lipman, and Tali Dekel. Multidiffusion: Fusing diffusion paths for controlled image generation. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pages 1737-1752. PMLR, 2023.
[31] Shansan Gong, Mukai Li, Jiangtao Feng, Zhiyong Wu, and Lingpeng Kong, Diffuseq: Sequence to sequence text generation with diffusion models. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023.
[32] Xiang Li, John Thickstun, Ishaan Gulrajani, Percy Liang, and Tatsunori Hashimoto. Diffusion- Im improves controllable text generation. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022.
[33] Jacob Austin, Daniel D. Johnson, Jonathan Ho. Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Anmual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pages 17981-17993, 2021.
[34] Shansan Gong, Shivam Agarwal, Yizhe Zhang, Jiacheng Ye, Lin Zheng, Mukai Li, Chenxin An, Peilin Zhao, Wei Bi, Jiawei Han, et al. Scaling diffusion language models via adaptation from autoregressive models. ArXiv preprint, abs/2410.17891. 2024.
13

---

[35] Minkai Xu, Tomas Geffner, Karsten Kreis, Weili Nie, Yilun Xu, Jure Leskovec, Stefano Ermon, and Arash Vahdat. Energy-based diffusion language models for text generation. ArXiv preprint. abs/2410.21337,, 2024.
[36] Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort. Deep Ganguli, Tom Henighan, et al. Training a helpful and harmless assistant with reinforcement learning from human feedback. ArXiv preprint, abs/2204.0586. 2022.
[37] Chenglong Wang, Hang Zhou, Kaiyan Chang, Bei Li, Yongyu Mu, Tong Xiao, Tongran Liu, and Jingho Zhu. Hybrid alignment training for large language models. ArXiv preprint, fbs/2406.15118. 2024.
[38] Tong Zheng, Hongming Zhang, Wenhao Yu, Xiaoyang Wang. Xinyu Yang, Runpeng Dai, Rui Liu, Huiwen Bao, Chengsong Huang. Heng Huang, et al. Parallel-rl: Towards parallel thinking via reinforcement learning. ArXiv preprint, abs/2509.07980, 2025.
[39] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. ArXiv preprint, abs/ 0ttp/1777.dur
1101107777.3777777 .0..6063772 2017.
[40] Leo Gao, John Schulman, and Jacob Hilton, Scaling laws for reward model overoptimization. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barhara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, International Conference Machine Learning. ICML 2023. 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pages 10835-10866, PMLR, 2023.
[41] Lunjun Zhang, Arian Hosseini, Hritik Bansal, Mehran Kazemi, Aviral Kumar, and Rishabh Agarwal. Generative verifiers: Reward modeling as next-token prediction. ArXiv preprint, 0is/24008.15240, 2024.
[42] Chenglong Wang, Yang Gan, Yifu Huo, Yongyu Mu, Qiaozhi He, Murun Yang, Bei Li, Tong Xiao, Chunliang Zhang, Tongran Liu, al. Gram: A generative foundation reward model for reward generalization. ArXiv preprint, abs/2506. 14175, 2025.
[43] Kaiwen Zheng, Yongxin Chen, Hanzi Mao, Ming-Yu Yu Liu, Jun Zhu, and Qinsheng Zhang. Masked diffusion models are secretly time-agnostic masked models and exploit inaccurate categorical sampling. ArXiv preprint, abs/2409.02908,, 2024.
[44] Jacob Devlin, Ming- Wiig-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pages 4171-4186, Minneapolis, Minnesota, 2019. Association for Computational Linguisties.
[45] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis. Luke Zettlemoyer, and Veselin Stoyanov. Roberta: robustly optimized bert pretraining approach. ArXiv preprint, abs/1907.1 1692, 2019.
[46] Alex Wang and Kyunghyun Cho. BERT has mouth, and it must speak: BERT as Markov random field language model. In Antoine Bosselut, Asli Celikyilmaz, Marjan Ghazvininejad, Srinivasan Iyer, Urvashi Khandelwal. Hannah Rashkin, and Thomas Wolf, editors, Proceedings of the Workshop on Methods for Optimizing and Evaluating Neural Language Generation, pages 30-36, Minneapolis, Minnesota, 2019. Association for Computational Linguisties.
[47] Martin Puterman. Markov decision processes. Handbooks in operations research and management science, 1990.
[48] David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin A. Ried- miller. Deterministic policy gradient algorithms. In Proceedings of the 3 31th International Conference on Machine Learning, ICMI. 2014, Beijing, China, 21-26 June 2014, volume 32 of JMLR Workshop and Conference Proceedings, pages 387-395. JMLR org. 2014.
14

---

[49] Richard S Sutton, Andrew G Cambridge, 1998.
Barto, et al. Reinforcement learning: An introduction. MIT press
[50] Dhawal Gupta, Yash Chandak, Scott M. Jordan, Philip S. Thomas, and Bruno C. da Silva. Behavior alignment via reward function optimization. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Advances in Neural Infor- mation Processing Systems 36: Anmual Conference Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
[51] Niklas Muennighoff, Zitong Yang. Weijia Shi, Xiang Lisa Li, Li Fei-Fei, Hannaneh Hajishirzi, Luke Zettlemoyer. Percy Liang, Emmanuel Candès, and Tatsunori Hashimoto. sl: Simple test-time scaling. ArXiv preprini, abs/2501. 01925001.1939,, 2025.
[52] Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama
3 herd of models. ArXiv preprint, abs/2247.22188. 2024.
[53] Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang. Timothée Lacroix, and William EI Sayed. Mistral 7b, 2023.
[54] Andrew Y Ng, Daishi Harada, and Stuart Russell. Policy invariance under reward transforma- tions: Theory and application reward shaping. In leml, volume 99, pages 278 287. Citeseer, 1999.
[55] Michael Luo, Sijun Tan, Justin Wong, Xiaoxiang Shi, William Y. Tang, Manan Roongta, Colin Cai. Jeffrey Luo, Li Erran Li, Raluca Ada Popa, and Ion Stoica. Deepscaler: Surpassing ol -p-preveew with a 1.5b model by scaling rl, 2025. Notion Blog.
(56) Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. ArXiv preprint, abs/21 14168, 2021.
[57] Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards. Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. In The Twelfth International Conference Learning Representations, ICLR 2024, Vienna, Austria, May 1- 2024. OpenReview.net, 2024.
[58] David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang, Julien Dirani, Julian Michael, and Samuel Bowman. Gpqa: A graduate-level google-proot q&a benchmark. In First Conference on Language Modeling, 2024.
[59] Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforce- ment learning. Machine learning, 8:229-256, 1992.
[60] Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Y Wu, al. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. ArXiv preprint, abss/4002.33300, 2024.
15

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Please see Appendix A.I
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
16

---

Justification:
Please see Appendix B
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: The models evaluated in the experiments are all public, so it is quite easy to reproduce the results
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
17

---

Answer:
[Yes]
Justification:
The data and code are uploaded to GitHub.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ /pipps.ccppullc//uuuii wwwiii/gguie///eeeeeeoecc CoodeSubeiiie ssi wsionnolli onnolicy)))) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Please see Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We show the error bars for reinforcement learning experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
18

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: The models in this paper are all public, and only the inference is needed. The computer resources for running these models are well known.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that didn't 't make into the paper).
9. Code ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification: Please Appendix A.2
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work primarily enhances the reasoning capabilities of DLMs. Its societal impact is limited to our understanding of how reward optimization methods can improve
DLMs.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
19

---

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification: There is no high risk for misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We have cited the original paper or attached the link to the existing assets used in this paper.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the Datasets.papersstt paperswi thcode. cww/ datasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
20

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: For the newly proposed data sets, we provide detailed description.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/cooe/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
No crowdsourcing in this study.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
21

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: We simply use GPT to embellish our text. There is need to provide detailed information.
Guidelines:
The answer NA means that the method development in this research does involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy for what should or should be described.
22

---

Supplementary Materials for MRO
A Limitations and Ethics Statement
A.1 Limitations
We propose the step-wise group reward optimization (SGRO) approach, which can effectively reduce the reward variance during the reward optimization. However, this introduces new hyperparameters to our MRO, including the size of denoising steps w per group and the size of groups chosen for optimizing DLMs the rejection sampling and reinforcement learning. The selection of these hyperparameters requires careful consideration. To address this limitation, conduct comprehensive ablation experiments, as shown in Figure 7 These experimental results provide valuable guidance for selecting the optimal hyperparameters, Additionally, in future work, we will consider designing techniques to determine these hyperparameters automatically.
A.2 Ethics Statement
This work does not involve any ethical concern.. All data collected for training our DLMs through rejection sampling and reinforcement learning are sourced exclusively from open-source materials. Additionally, this paper may reference certain case study content. However, these references are presented an elliptical manner, and any potentially harmful content will not be explicitly presented.
B Proofs for Theoretical Results
In this section, provide the proofs for three theoretical results. The first result shows that introducing a potential-based shaping in reward optimization leads to a higher reward variance. The second result demonstrates that using SGRO mitigate this issue. The third result proves that token verification reward encourages DLMs to generate responses with stronger internal sequence correlation.
Property 1. Under potential-bun reward shuping, the expected reward remains the same, i.e., E[R(s, a)] = E[ a)), but Var(B(ss, a)) is higher than Var (R((Rss,))) a)).
Proof:
Let's denote the original reward as M(s,, at) and the potential function as ((8)). The potential-based reward shaping is given by
R(st, R(s), (0((541) ((st)
(10)
where À is the discount factor. The expected reward under the diential-basee shaping is
E[R(se, ar]] E[A(s,, at) + XB B(
E[ft(se, ar) + (0(()+ ((s))
(11)
(12)
Since the expectation of the potential function the next state ([((s++1) equal to the expectation of the potential function at the current state (B[(s;; (because the potential function is a function of the state and the state transitions are Markovian) [49 we have
E ((((,,,)))
E[[((,; at)) + (x]) ((+())))
(13)
= E[R(s, a:)] + (A {s:)l
(14)
For the variance, we have
Var (ar( = Var(fi(se, ar) + (0((5(1) -((ss)
(15)
= Var(f((s az)) +
Var(X 1)) + 2Cov az), ((s))) (A (bss))
2Cov R(s,, at). (+(
(16)
additional variance introduced by the potential function terms
23

---

Algorithm Simplified Grouped Reward Optimization (SGRO)
Input: Pre-trained DLM TH: Group Size w; Discount X; Learning Rate 7/; Response Length L
Output: Fine-tuned DLM Te*
for each prompt Po dataset do
% Roll-out full denoising trajectory PT FULLMASK(L)
states t [], actions <- for t < T downto do
sample r To(- Po.
states. secees..ppnddd (ooppen(((p ((po, Tt, ()) aapponn..pppeddcccc
D 81
>
end for
% SGRO: compute rewards and returns per group
grad +
for g 0 g < <g
do start t- g. end <- min T)
Rintra t-
Intra-sequence rewards
for k f start to end- 1 do
Rintra Rintra + TOKEN TONEENVERIFCAATONNEEWWWW actions (A])
Rintra + Rintra + SERPLEXITYREWANT (wt.///AAD(aaoossiio.ooaoooooola
end for
Rg +0
Inter-sequence reward (if final group)
if end T then
R2 & TISSAACUUACCAATTTOAAARRRAAAACCCC
end if
Qstart t POTENTTALFUUCCTIO((sttt.. (start)
b Potential-based shaping
Dend
POTENTIALFUNCTION (states (end-1))
shaping t X Gond 2start
Ryroup +- Rintra + Ry + shaping
D Group return
for k: 4 slarl to end-1 do
log <- log Th (actionssk) states (k:)
grad <- grad + (o Rgroup
REINFORCE update for steps in this group
end for end for
% Parameter update
I 0 77 grad
end for
Since (s:) and ((8++1) are not independent, the covariance terms are not zero. However, the variance of the potential function terms could add to the variance of the original reward, leading to higher overall variance, the general reinforcement learning, [50) has also proven that this result
holds when (s) >> 2f((;; 4)) for all a. Furthermore, (54) suggests that such a condition is indeed possible, even when using optimal baseline reward technique. Therefore, we have
Var(R(s;,aa) > Var(it(sr,.:))
(17)
This completes the proof of Property 1.
Property Using step-wise group reward optimization can reduce the reward variance, i.e. Var(R(s,,) > group reward optimization.
where R(w)(.) represents the reward after applying step-wise
Proof: SGRO groups w denoising steps together and provides reshaped reward for each group. This means that the number of times the potential function computed and added the reward is reduced. Let's denote the reward for group of w steps as
10
R(w) (se, ar)
M(004 a.+c) + x0 - ((s))
(18)
24

---

The variance of this group reward is
( az)) Var
R(se+i; (0( 25())
(19)
Following the derivation in Eq. 16 obtain
Var(R (st,ar)
Var(R(s), +
/Var(D((LtHI) + Var(O(si)) + 2XCov( oz), ((8++)))
3Cov(R(s;, ((ss)) 2XCov(@(5 (s:)))
(20)
Compared to Eq. 16 we can observe that due to the larger interval w between St+w and the correlation between (s)=w)) and ((x)) is typically weaker than the correlation between ((++1) and (sx)). Therefore, we have
[Cor(((s4ww); ((s))) (Cov((())2)
(21)
Here, the negative term 222 (8x)) in the variance expression of ((w)(s, 0z) has larger absolute value compared to -2ACov (CCCoo(((s))) ($1+1), ((s))) )) in the variance expression of R(s,, az). as the correlation between (($++w) and (s;)) is weaker. Therefore, obtain that this results in smaller overall variance for (st, az):
Var (ar(R(() a.)) < Var(R(s,, a,))
(22)
This completes the proof of Property 2. Here, one problem may arise when only a single group provides potential-basce reward, the reward that could guide the model toward better token correlation may be lost. Therefore, in this approach, we strike a balance between the reward signal and the reward variance for MRO. The implementation of SGRO is simple, illustrated in Algorithm A.
Property 3. When DLM is optimized to achieve higher token verification rewards, it consequently
enhances its intra-sequence correlation.
Proof: We provide a theoretical proof for the token verification reward by analyzing it from the perspective of mutual information. Firstly, we recall the definition of the token verification reward the denoising step t:
1 N Σ Pro(r?"
R
Po, 11-r/""))
(23)
where M = {m...... myy is the set of masked token indices, rM is the set of the predicted tokens at these positions, and rr-1/r,, is the set of unmasked token set. The joint probability over the masked positions can be factorized autoregressively (within the masked set) as:
N
Pro(r (o,Tt))
II
Pro(r Po.Tt-1/1M
(24)
where the rest of the sequence is fixed and only the masked tokens are predicted. The token verification reward approximates this joint modeling by computing ceave-one-out log bogg-robabilities:
R
log Pro (o-r/-1/77
(25)
==1
Now define the empirical average pairwise mutual information (PMI) among the masked tokens:
PM
2
Σ
10 Po. (
(26)
N(N -)
ISiSJEN
25

---

GSMSK
MATH500
GPQA
Countdown
Sudoku
Model/Length
128
256
512
128
256
512
128
256
512
64
128
64
LLaDA
74.3
78.8
79.4
28.0
33.2
34.4
22.2
29.2
30.3
13.8
14.1
11.2
LLaDA-MRO-RS
76.9
79.6
82.6
31.0
34.2
36.2
26.3
32.1
34.3
21.4
22.0
17.2
LLaDA-MRO-RL
77.1
80.9
81.8
33.4
35.2
37.4
28.8
33.8
33.8
24.6
27.2
20.2
LLaDA-s1
70.7
76.4
78.7
25.0
26.0
28.8
21.7
25.3
28.3
10.2
12.4
8.4
MLDDA-1-MMR--SS
73.3
77.7
80.1
27.8
28.0
29.0
24.2
29.3
32.8
17.3
17.8
15.2
LLDDA-SIMMR--II 71.8
75.9
78.0
26.6
27.2
29.4
23.7
28.3
32.3
17.1
16.2
13.8
Table 4: Results of the MRO reinforcement learning. The suffixes "-RS" and "-RL" denote the results obtained using rejection sampling with k = 4 and reinforcement learning, respectively.
Using standard second-order Taylor expansion around the independence assumption (as in energy- based models), this quantity be approximated as:
Pro(r 1-1. r21 1.)
PMIlagg ~ N(N - 1) $- Pro(r Pro my
log
(27)
Importantly, the leave-one-out fog-probabilities computed in TVR serve as sufficient statistics for esti- mating these pairwise interactions. Therefore, maximizing the average leave-one-out log-probable pr-obabbbiltty is first-order equivalent maximizing PMIlvg thereby promoting stronger intra-sequence correlation. This completes the proof of Property 3.
C Experiments
In this section, we provide additional experimental details and present the experimental results of MRO with reinforcement learning.
C.1 Experimental Details
Training Setups. For training LLaDA-s1, we used a pre-trained version of LLaDA. The learning rate set to 2c-5. We trained this model on the sl dataset for 3 epochs. In contrast to (26), found that training for more epochs on the sl dataset did not result in further performance improvements. For rejection sampling and reinforcement learning, we performed model validation every 50 steps and selected the best model based on performance on the validation set as our final model. For computing R,, we sampled one token from the predicted masked tokens at each denoising step. For N2ppl we set Cpp and Fpp to 100 and 100, respectively. In Figure [T] the x-axis corresponds to the intra- and inter-sequence correlation rewards defined in Section 4.2: the left part presents the cumulative intra-sequence rewards R, R.pll while the right part shows the inter-sequence reward Rg. As the reward scales differ, we applied standardization for visualization.
set the learning rate 2e-6. During training,
Training Datasets. For both rejection sampling and reinforcement learning, we utilized DeepScaleR [55] in conjunction with the 10k Countdown and Sudoky datasets. These datasets were randomly shuffled to ensure a
data distribution.
Evaluation. For the evaluation, we focus on five reasoning tasks: GSM8K (56), MATH500 (57] GPQA (58). Countdowr and Sudoku During testing, we set the sampling temperature to 0.25.
C.2
Reinforcement Learning
We explore the use of reinforcement learning to implement our MRO. The experimental setup and results presented below.
26

---

Longth of 256
Length of 512
39.0
38.0
38.0
37.4
37.0
36.8
36.4
36.0
Accurucy
35.0
36.0
34.4
35.2
35.0
34.8
35.6
36.4
34:4
35.0
34.0
33.6
33.0
LLaDA
MRO
MRO-v1
MRO-v2 Performance on Test-time Scaling
MRO-v4
MRO-15
MRO-v6
(a)
39.0
38.0
37.8
Accuracy (55)
37.0
36.2
36.0
35.0
34.4
35.2
34.6
31.8
31.8
34.8
35.4
33..0
"33.2
33.6
33.6
$3.0
32.2
32.0
31.0
LLaDA
MRO
MRO-v1
MRO-v2
MRO-v3
MRO-v4
MRO-v5
MRO-v6
Performance Reinforcement Learning
Figure 6: Performance comparison of different MRO variants the MATH500 benchmark for both lest-time scaling and reinforcement learning.
Task Setup. We employed the REINFORCE [59 to perform this optimization. Specifically, during the optimization process, used temperature sampling obtain denoising sequence (fr, f1. J. Subsequently, we use REINFORCE to train our DLMs through cumulative reward. The loss function can be given by
I
Ca( = -Eo,, (fr
1lft = M] log Pro(ff--llos Race
(28)
where Race is the cumulative reward, computed as: Race Ei_r R(st, a.). Similarly, we apply the SGRO in reinforcement learning, where we group the sampled sequence, with each group containing w steps. For each group, we provide reshaped reward. Additionally, during optimization, similar to rejection sampling, instead of using all groups, sample only subset of the groups. Building upon this, we integrated quality evaluation scores into the reinforcement learning training process through shaping mechanism. Specifically, we selected group fastw- Then, we computed the quality rewards Ri and Ri_.. for the first and last steps of the group, respectively. The final quality reward for the group was determined by the difference between these two quality rewards. Although we implemented MRO the DLM using REINFORCE, any reinforcement learning algorithm (such PPO (39 GRPO 60) could be used to achieve MRO, outlined in the modeling framework presented in Section 4.2.3 The primary focus of this work is to demonstrate the effectiveness of MRO in enhancing token correlation, rather than to explore the performance of different reinforcement learning algorithms in the DLM training process. Therefore, we did not conduct tests of these algorithms by one.
Results. We compare the performance of our DLMs trained via reinforcement learning to those trained with rejection sampling and the LLaDA instruction model. The results listed Table 4 From the results, we observe that with reinforcement learning, our MRO still achieves significant improvements across various benchmarks. However, when compared rejection sampling, our reinforcement learning approach does not show a substantial advantage. We identify two main reasons for this observation. First, our rejection sampling approach has been enhanced compared to the original; it combines offline data construction with online sampling techniques. Second, our reinforcement learning approach is relatively basic and does not incorporate advanced modifications or improvements.
27

---

C.3
Ablation Study
In this subsection, we design several MRO variants to further describe the functionality of intra- sequence and inter-sequence rewards, as well as SGRO.
Reward Design. As shown in Table 5 we design six MRO variants to investigate the impacts of intra-sequence and inter-sequence rewards.
Variant MRO-v1 MRO-v2 MRO-v3 MRO-v4
MRO-v5
MRO-v6
R
Table 5: Description of MRO variants.
We conduct experiments on test-time scaling and reinforcement learning evaluate the performance of these variants. As illustrated in Figure the results reveal that achieving superior performance with single reward combination challenging. For example, MRO-v2 and MRO-v3 demonstrate relatively poor performance compared to other variants in both test-time scaling and reinforcement learning. Furthermore, the token verification reward proves to be highly effective. We can see that MRO-v2, MRO-v3, and MRO-v5 perform worse in comparison. However, when compared to LLaDA, it is clear that all of our reward designs are effective, except for MRO-v2 in reinforcement learning, which exhibits some performance degradation. Other variants, varying degrees, lead to performance improvements. This further supports the correctness our design approach, centered around enhancing token correlation through tailored reward strategies. Interestingly, results in reinforcement learning show that MRO-v4 outperforms MRO itself. This improvement could be attributed to sampling and the potential reward variance caused by the perplexity reward. Nevertheless, since other experiments have confirmed the usefulness of this reward, we chose not to discard it.
Step-wise Group Reward Optimization. We also conduct an ablation study on our SGRO. Specif- ically, we test the case where SGRO is not applied, meaning that each denoising step receives shaping reward in rejection sampling. As shown in Table we find that SGRO is effective and helps the MRO achieve better performance.
D
Analysis
D.1
Comparison of MRO with Other Reasoning-Enhanced DLM Approaches
We compare MRO with other existing reasoning- enhanced models and approaches. These include DiffuLLaMA (34), which adapts the LLaMA model for DLMs; EDLM (35), which introduces an energy function to enhance sequence-level cor- relation; Dream 24. which utilizes the Qwen-2.5- 7B model for initialization; and dI-LLaDA [26].
GSMBK
MATH500
Model/Length
256
512
256
512
Daff
63.1
Pream-7B-Instrue
77.2
EDLM
78.1
80.0
34.4
35.4
dI-LLaDA
81.1
82.1
38.6
40.2
which trains the LLaDA-SB-Instruuc' model using GRPO. The results are presented in Table 7 Note that, except for EDLM, the results for the other baselines are taken directly from the original pa- pers. For EDLM, we replicate its autoregressive energy function version in our codebase. First, compared DiffuLLaMA and Dream, our MRO achieves competitive results. Moreover, we ob- serve that although EDLM incorporates sequence- level correlation (i.e., inter- inter-sequence correlation described in this paper), our MRO still outperforms it. We attribute this the lack of consideration for intra-sequence correlation in EDLM. However, we observe that our model performs slightly worse than the dI-LLaDA model. We argue that this comparison is not entirely fair, as dI-LLaDA is trained using task-specific training set for GRPO. This potentially gives it an advantage by benefiting from
LLaDA-MRO-RS
79.6
82.6
34.2
36.2
LLaDA-MRS -HRO-RL
80.9
81.8
35.2
37.4
LLaDA-MRO-TS
82.5
82.9
39.4
42.6
Table 7: Performance comparison of MRO with other reasoning-enhannn models and ap- proaches. "-TS" indicates that use task- specific training data to optimize the DLM in reinforcement learning.
28

---

MATH500
GPQA
Model/L.ength
256
512
256
512
LLaDA-MRO LLaDA-MRO w/o SGRO
34.2
36.2
32.1
34.3
32.8
35.4
31.3
33.3
Table 6: Ablation study of SGRO.
task-specific training. To validate this, we also conduct reinforcement learning using a task-specific training set. More specifically, during the reinforcement learning training, use the GSM8K and MATH50Q training sets to perform the MRO, respectively. We find that the results from this training approach surpass the performance of dI-LLaDA.
D.2 Scaling Training with Different Group Sizes
We scale training in MRO by using different group sizes. Specifically, we keep the total number of training samples constant while testing various group sizes (1, 4, 8, 16, 32,
Rejection Sampling -
Reinforcement Leamiing
39.0
38.0
Accuracy
37.0
64) during both rejection sampling and re- inforcement learning. The results are shown in Figure 7 From the results, we can find that increasing the group size provides some benefits, but these benefits diminish after a group size of 2, with minimal improvements beyond that. Notably, performance becomes unstable after a group size of 16. Considering both performance and training costs, we choose group size of 2 for our experiments.
36.0
35.0
64
Group Size
Figure 7: Performance of MRO with different temper- ature settings, evaluated on the MATH500 benchmark.
D.3 Performance with Different Perplexity Reward Upper Bounds
We conduct ablation study by varying the upper bound of the perplexity reward over the range (50, 80, 100, 130, 150), imple- mented on the We use the LLaDA-8B-Insten
MATH150
CPQA
38.0
36.0
model. model as the baseline. The results show that model per- formance remains relatively stable when the bound lies between 80 and 130, demonstrat- ing the robustness of our method as well as the insensitivity of this hyperparameter. In contrast, smaller bound tends to compress reward differences and weaken the learning
34.0
32.0
30.0
50
80
100
130
150
Upper Bound
Figure 8: Performance with different perplexity reward upper bounds. The solid line denotes the performance of MRO, while the dashed line represents the baseline.
signal, whereas a larger bound introduces excessive reward variance, both of which can negatively affect training stability.
29

---

Kylar went the buy glasses for his new apartment. One glass costs $5, but every second glass costs only 60% of the price. Kylar wants to buy 16 glasses. How much does he need to pay for them?
Prompt
Answer
64
determine how much Kylar needs to pay for 16 glasses, we need to consider the cost of the first glass and the discounted of every second glass.
first glass costs $5.
2. Every second glass costs 60% of the price price of the first glass, which is 16 x 0.60 9.60 dollars.
Let's calculate the total cost:
The first glass costs $5,
LLaDA
The second glass costs $9.60.
The third glass costs $5.
The fourth glass costs $9.60.
The sixteenth glass costs $5,
So, the total cost is: 5 9.60 5 9.60 + 5 +5000055505++00++5+50000000 9.60 5 9.60 + = 120
The answer is $120.
determine much Kylar needs to pay 16 glasses, we need to consider the cost of the first glass and the discounted of the second glass.
1. The cost of the first glass $5.
2. The cost of the second glass is 60% of the first glass's price, which is 0.6 x 3 dollars.
Since Kylar wants to buy 16 glasses, he can pair them follows:
First glass: $5
Second glass:
Third glass: $5
Fourth glass: $3
LLaDA MRO
Sixteenth glass: $5
Each pair of glasses consists of $5 + S3 = $8.
There are pairs of glasses (since 64 divided by is 8).
So, the total cost is: 8 pairs x 8 per pair $64
Therefore, Kylar needs pay S64 for 16 glasses.
The answer is $64.
Table A case from the test-time scaling experiments. MRO provides accurate guidance for token correlation and helps identify a more precise CoT-style reasoning path. Note that we remove the "sthink><</hhnkk=" and "canswer></answer>> format tags to improve the readability of this case.
30

---

Two sides of triangle are each 8 units long. the third side has whole number length, what is the greatest possible perimeter, units, for the triangle?
Prompt
Answer
31
x is positive number, this condition is always satisfied
From conditions and 2, we get:
16 > r x < 0
LLaDA
Since a must be a whole number, the largest possible value for æ is satisfies both conditions is z = 7.
Now, calculate the perimeter:
Perimeter = 8 + 8 23
Thus, the greatest possible perimeter of the triangle is 23.
To solve:
1. We know form a triangle, the sum of the lengths of any two sides must be greater than the length of the third side.
Let the denote the length of the third side as
3. We to satisfy the triangle inequality theorem:
88 8 > x
-8 + x > 8
x 8>
4. Simplifying these inequalities:
16 >
x >
LLaDA-MRO-RS
5. Combining these inequalities, get: 16
Since x must be a whole number, the possible values for x 1, 3, 4, 5, 8, 9, 10, 11, 12, 13, 14, and 15.
7. The find the greatest possible perimeter, we choose the largest possible value for r: 15
The perimeter of the triangle the sum of its sides:
Perimeter = + 8 + 15 = 31
Therefore, the greatest possible perimeter for the triangle is 31.
Table 9: A case from the rejection sampling experiments. We find that using MRO to search for better and more accurate reasoning paths, and then fine-tuning the DLM with these paths, helps the model tend toward better reasoning paths during the decoding process.
31

---

Artic has a flower stand at the Farmers Market. He sells three kinds of flowers: marigolds, petunias and begonias. He usually sells marigolds for $2.74 per pot, petunias for $1.87 per pot and begonias for $2.12 12 per pot. Artie has change today, he has decided to round all his prices to the nearest dollar. If Artie sells 12 pots of marigolds, pots of petunias and 17 pots of begonias, how much will he make?
Prompt
Answer
88
To determine how much Artie will make, we need to calculate the total revenue from each type of flower and then sum these amounts.
** Mariioolss:**
Price per pot: $2.74
Number of pots sold: 12
Revenue from marigolds: 2.74 x 12 32.88
2. **Peuniaas:**
Price per pot: $1.87
Number of pots sold: 9
LLaDA
Revenue from petunias: 1.87 x 9 16.83
** Begonias:
Price per pot: $2.12
Number of pots sold: 17
Revenne from begonias: 2.12 x 17 36.04
Now, we sum the revenues:
32.88 + 16.83 + 36.04 85.75
The answer is $86.
First, we's round the prices to the nearest dollar:
Marigolds: $2.74 rounds to
Petunias: $1.87 rounds to $2
Begonias: $2.12 rounds to $2
Next, we calculate the total revenue for each type of flower:
Marigolds: 12 pots $3/pot $36
LLaDA-MRO-RL
Petunias: pots * $2/pot = $18
Begonias: 17 pots $22pot S34
Finally, we add up the revenues together: $36 + $18 $34 $88
The answer is $88.
Table 10: A case from the reinforcement learning experiments. Interestingly, after applying reinforce- ment learning for MRO optimization, we find that the DLM generates shorter and more accurate CoT-style reasoning paths. We conjecture that this improvement is enabled by the optimization against the inter-sequence correlation reward, which encourages the model to achieve better format and accuracy rewards with fewer denoising steps.
32


Reviews:
[{"review_id": "wf3HssQ4pK", "reviewer": "Reviewer_3ZHw", "summary": "The 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.", "questions": "1. 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?", "limitations": "Yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\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.", "was_revised": true, "final_justification": "I have carefully read the author's rebuttal. They have added the experiments I suggested in my review. However, I am not an expert in diffusion language models, and after reading the other reviews, I agree with reviewer Sm7r that the computational cost is high. Since the authors incorporated additional experiments I recommended and solve my concern, I will maintain my original score but lower my confidence to 2.", "rebuttal": "Dear Reviewer 3ZHw,\n\nWe sincerely appreciate your constructive and thoughtful feedback. \n\nWe greatly appreciate your recognition that the main claim of our analysis on token correlation can \"offer valuable insights and could inspire future research\", as well as our reward design, empirical improvements across benchmarks, and the clarity of our writing.\n\nBelow, we explain the main points you are concerned about.\n\n---\n>*W1&S1: 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.* \n\nThanks for your helpful feedback! Below, we address the concern regarding decoding efficiency and provide empirical evidence to support our claims. **Our results demonstrate that the proposed MRO method is both practical and computationally efficient, making it suitable for real-world applications**.\n\n1. **We would like to clarify that in both the rejection sampling and RL fine-tuning settings, multi-reward computation and policy gradient updates are performed only during training**. Once the model parameters are updated, decoding proceeds identically to the standard LLaDA baseline [22], with no additional reward computation or forward passes involved. Thus, the end-to-end decoding latency remains unaffected. This is further supported by the latency comparison presented in the table below, which shows that our rejection sampling and RL setups do not introduce noticeable overhead during decoding. Note that all tests are conducted on eight A800 GPUs with a length of 256 under data parallelism, and each setting is run five times to report a range of latency values.\n\n    |Method|MATH500 |  |GPQA|  |\n    |-|:-|:-:|:-|:-:|\n    ||Decoding Time |Score|Decoding Time|Score|\n    |Vanilla Decoding (LLaDA)| 0.35h~0.39h |33.2 | 0.16h~0.21h |29.2 |\n    |LLaDA-RS + MRO| 0.34h~0.37h |34.2 | 0.14h~0.18h |32.1 |\n    |LLaDA-RL + MRO| 0.36h~0.39h |35.2 | 0.16h~0.20h |33.8 |\n\n2. To further address the concern regarding computational overhead during test-time scaling, we provide a detailed latency comparison between vanilla decoding (i.e., LLaDA) and our proposed MRO, as shown in the table below. Given that TTS requires additional sampling and thus introduces time overhead, we also include a comparison with a confidence-based reward TTS baseline, which involves the same sampling cost as our MRO but does not require extra reward computation. From the results, we observe that **while reward computation does introduce some overhead, it remains within a reasonable range and is justified by the substantial performance gains it brings**. Compared to the confidence-based reward TTS baseline, we further observe that **when excluding the time cost introduced by sampling, our reward design incurs only minimal additional overhead while delivering substantial performance gains**. \n\n    |Method|MATH500|  |GPQA|  |\n    |-|:-|:-:|:-|:-:|\n    ||Decoding Time |Score|Decoding Time|Score|\n    |Vanilla (LLaDA)| 0.35h~0.39h |33.2 | 0.16h~0.21h |29.2|\n    |LLaDA-TTS + Confidence-based Reward| 0.73h~0.81h |35.2 | 0.27h~0.34h |30.6| \n    |LLaDA-TTS + MRO| 0.84h~0.85h |36.0 | 0.31h~0.44h |34.6|\n\n>*W2: This proposed approach generally mitigates the decoding problem of DLMs and could benefit from further verification on other tasks.*\n\nThanks for the thoughtful suggestion! We have extended our evaluation to a broader range of tasks, including **general knowledge question answering (MMLU)**, **code generation (HumanEval)**, and **instruction following (AlpacaEval2 and Arena-Hard)**, to verify the effectiveness of our proposed MRO further. The results demonstrate that **our MRO achieves strong robustness and can deliver consistent improvements (albeit to varying degrees) across these diverse tasks**. The results and details are as follows.\n\n|Model|MMLU|HumanEval|AlpacaEval2|Arena-Hard|\n|-|:-:|:-:|:-:|:-:|\n|LLaMA-3-8B-Instruction|68.4|59.8|25.3|22.3\n|Mistral-7B-Instruct|60.1|30.5|14.7|12.6\n|Deepseek-LLM-7b-Chat|48.2|26.2|11.2|10.3\n|Qwen2.5-7B-Instruct|71.9|56.7|27.8|25.2\n|||\n|LLaDA|65.5 |47.6 |16.3 |10.0 |\n|LLaDA-RS + MRO|67.5 | 48.1 | **20.2** | 12.3 |\n|LLaDA-RL + MRO|**68.2** | **50.0** | 19.4 | **15.7**|\n\n1. On general benchmarks such as MMLU and MMLU-Pro, which span multiple domains including biology and physics, our method continues to demonstrate effectiveness. This suggests that MRO is not limited to mathematical reasoning tasks. Specifically, we replace Countdown and Sudoku in our training set with a subset of the MMLU training set for both rejection sampling and reinforcement learning.\n\n2. On the code generation task, we train the model using the CodeAlpaca-20K dataset. Our method still yields notable improvements under both RL and rejection sampling settings.\n\n3. To evaluate instruction-following capabilities, we train the model on the Alpaca dataset. Note that during the computation of inter-sequence rewards for this task, we replace the standard answer-matching metric with a learned reward model, which allows for more nuanced assessment of response quality.\n\n\n>*W3: The paper could benefit from adding more baselines that require training on the dataset in order to verify the proposed method better.*\n\nThanks for your valuable feedback! In response, we have achieved several additional baselines, all trained on the same dataset (DeepScaleR + Countdown/Sudoku) using identical hyperparameters. Our experimental results show that **LLaDA combined with our proposed MRO consistently achieves the best performance**. Specifically, we designed the following experiments:\n\n1. **Reward variants within the LLaDA model**: We have achieved baselines using different reward functions for both rejection sampling and reinforcement learning to the LLaDA model, including LLaDA-$R^{tv}\\_{t}$, LLaDA-$R^{ppl}\\_{t}$, and LLaDA-$R^{q}\\_{0}$. _Note that these results are already present in our initial submission and are used as baselines in training, as shown in Figure 6._\n\n2. **Baselines using other foundation models trained on the dataset**: We have achieved RL fine-tuning baselines using other popular foundation models, such as LLaMA-3-8B-Instruction, Deepseek-LLM-7B-Chat, and Mistral-7B-Instruct, trained on the same dataset using our reward designs and the REINFORCE++ algorithm. For these models, we use a delayed reward at the end of the sampled reasoning trajectory based on the final output, following the same reward evaluation protocol.\n\n3. **Rejection sampling with other foundation models**: We also perform rejection sampling experiments with other foundation models, where multiple reasoning paths are sampled and the one with the highest reward score is selected using our proposed reward design.\n\n\nThe results are summarized below. For consistency, the maximum length is set to 512 for MATH500 and GPQA, and to 128 for Countdown.\n\n|Model||Rejection Sampling|||Reinforcement Learning||\n|-|:-:|:-:|:-:|:-:|:-:|:-:|\n||MATH500|GPQA|Countdown|MATH500|GPQA|Countdown|\n|_**Training with Other Foundation Models**_|\n|LLaMA-3-8B-Instruction|18.8(+0.4)|26.8(+1.2)|4.5(+4.5)|17.4(-1.0)|28.7(+3.1)|0.4(+0.4)\n|Mistral-7B-Instruct|13.6(+0.5)|24.2(-0.5)|3.2(+3.2)|14.0(+0.9)|23.2(-1.5)|3.4(+3.4)\n|Deepseek-LLM-7B-Chat|11.2(+5.2)|20.7(+1.2)|3.8(+3.8)|10.0(+4.0)|22.0(+2.8)|2.6(+2.6)\n|_**Reward Variants within the LLaDA**_|||\n|LLaDA-$R^{tv}_{t}$|36.0(+1.6)|33.2(+2.9)|**23.4(+12.2)**|36.2(+1.8)|32.7(+2.4)|25.3(+11.2)|\n|LLaDA-$R^{ppl}_{t}$|34.0(-0.4)|30.5(+0.2)|14.1(+2.9)|33.6(-0.8)|30.8(+0.5)|18.9(+4.8)|\n|LLaDA-$R^{q}_{0}$|35.0(+0.6)|32.8(+2.5)|19.8(+8.6)|34.8(+0.4)|31.2(+0.9)|23.5(+9.4)|\n|LLaDA + MRO|**36.2(+1.8)**|**34.3(+4.0)**|22.0(+7.9)|**37.4(+3.0)**|**33.8(+3.5)**|**27.2(+13.1)**|\n\nFrom the results, we can observe the following:\n\n- When comparing different reward variants, our proposed MRO consistently achieves the best performance. The results confirm the validity of our design approach, which focuses on improving intra-/inter-sequence correlation through carefully crafted rewards. **This further demonstrates that the improvements are not primarily attributable to the dataset itself, but rather stem from our method's ability to address the overlooked issue of correlation in decoder language models.**\n\n- Applying the same-scale RL fine-tuning or rejection sampling on other foundation models (such as LLaMA-3 and Mistral) also leads to performance improvements on this dataset. However, the gains (_as indicated in parentheses in the table_) are generally smaller compared to those achieved by LLaDA. This can be attributed to the fact that these models have already undergone RLHF training and are relatively well-aligned. In contrast, LLaDA has not been exposed to any prior RL training, leaving more room for improvement.\n\n\n>*S2: Could the authors provide analysis on the intra-/inter- sequence correlations with and without the proposed method?*\n\nThanks for this helpful suggestion! We have conducted a quantitative analysis of intra- and inter-sequence correlations with and without our proposed MRO. The results show that **MRO can significantly enhance both intra- and inter-sequence correlations**. Specifically, we randomly sampled 200 examples from the MATH500 and GPQA datasets, and computed the differences in intra- and inter-sequence reward scores across five decoding runs with different seeds. The detailed comparison is presented below.\n\n| Model | MATH500 | | GPQA||\n| - | :- | :-: | :- | :-: |\n|   | intra-corr | inter-corr | intra-corr | inter-corr |\n| LLaDA | 3.44±0.18  |  1.02±0.14 | 2.76±0.21  |  1.02±0.15 |\n| LLaDA-RS-MRO | 3.79±0.16 | **1.58**±0.12 | 3.34±0.19  |  1.27±0.13 |\n| LLaDA-RL-MRO | **4.12**±0.24 | 1.42±0.11 | **3.55**±0.17 | **1.43**±0.12 |\n\nWe promise to add these related discussions in the revised version.\n\n---\nWe sincerely thank you for your positive feedback on our paper!\n\nBest,\n\nAuthors"}, {"review_id": "azJllP8GY7", "reviewer": "Reviewer_9WDq", "summary": "This 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.", "questions": "1.\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”", "limitations": "YES", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths:\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.", "was_revised": true, "final_justification": "I appreciate the authors’ detailed responses, which effectively address my concerns. I have accordingly raised my score. I strongly recommend that the authors include the additional experiments and analyses in the final version of the paper.", "rebuttal": "Dear Reviewer 9WDq,\n\nWe sincerely thank the reviewer for your positive and insightful feedback. \n\nWe greatly appreciate your recognition of our paper’s \"thorough analysis of the limitations of current DLMs in reasoning tasks\", particularly the focus on \"the lack of token correlation across denoising steps\". We are also pleased that our \"detailed explanations of the proposed MRO framework\", including the design of intra-/inter-sequence correlation rewards, and SGRO, are found to be clear. Finally, we are grateful for your recognition that the paper presents \"extensive experiments on multiple reasoning tasks across different models and settings\".\n\nWe will provide explanations for the main points that you are concerned about.\n\n---\n\n> *W1: The main contributions of the paper can be clearly articulated in the introduction.*\n\nThanks for your helpful suggestion! We intend to convey the core contributions throughout the introductory section. Specifically, as noted in line 73, we introduce MRO as an approach to enhance token correlation in DLMs without introducing additional architectural burdens. In line 81, we describe SGRO, our sparse reward optimization strategy, designed to further reduce computational overhead. Finally, in line 86, we summarize the empirical findings from extensive experiments across test-time scaling, rejection sampling, and RL, which collectively demonstrate the effectiveness of our proposed method.\n\nWe appreciate your suggestion that the main contributions could be made more explicit and easily identifiable. We promise to revise the introduction in the revision to include a clear bullet-point summary of our key contributions at the end of the section. \n\n> *W2: It is recommended that the paper include brief pseudocode for algorithms such as SGRO.*\n\nThank you for your valuable suggestion! We promise to add a concise Python-style pseudocode illustration of our RL training procedure with SGRO in the revised version, as shown below:\n\n**Input**: Pre-trained DLM $π\\_θ$; Group Size $w$; Discount $\\lambda$; Learning Rate $\\eta$; Response Length $L$   \n**Output**: Fine-tuned DLM $π\\_{θ^*}$\n\n```Python\nfor each prompt p_0 in dataset:\n    # 1. Roll-out full denoising trajectory\n    r_T = FULL_MASK(L)\n    states, actions = [], []\n    for t = T downto 1:\n        r_t_1 ~ pi_theta(·|p_0, r_t)\n        states.append((p_0, r_t, t))   # s_t\n        actions.append(r_t_1)        # a_t\n\n    # 2. SGRO: compute rewards and returns per group\n    grad = 0.0\n    for g in range(0, T, w):\n        start, end = g, min(g + w, T)\n        # 2-a Intra-sequence rewards inside this group\n        R_intra = 0.0\n        for k in range(start, end):\n            R_intra += TokenVerificationReward(p_0, actions[k])\n            R_intra += PerplexityReward(actions[k])\n        # 2-b Inter-sequence reward (only if the group ends at t=0)\n        R_q = 0.0\n        if end == T:  # covers final step t=0\n            R_q = TaskAccuracyAndFormatReward(actions[end-1])\n        # 2-c Potential-based shaping\n        potential_start = potential_function(states[start])\n        potential_end   = potential_function(states[end-1])\n        shaping = _lambda * potential_end - potential_start\n\n        # group return\n        R_group = R_intra + R_q + shaping\n\n        # 2-d REINFORCE update for steps in this group\n        for k in range(start, end):\n            log_prob = log pi_theta(actions[k] | states[k])\n            grad += ∇_θ log_prob * R_group\n\n    # 3. Parameter update\n    fine_tuned_theta = _theta + _eta * grad\n```\n\n> *W3&Q4: There are some typos in this paper.*\n\nWe appreciate your comment and have carefully proofread the manuscript to correct all typos in the revision.\n\n> *W4&Q1: 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?*\n\n| Upper Bound      | MATH500  | GPQA     |\n|:---------------- |:--------:|:--------:|\n| LLaDA (Baseline) | 34.4     | 30.3     |\n| LLaDA-MRO-RL-50  | 35.6     | 31.3     |\n| LLaDA-MRO-RL-80  | 37.0     | 33.3     |\n| LLaDA-MRO-RL-100 | **37.4** | **33.8** |\n| LLaDA-MRO-RL-130 | 36.6     | **33.8** |\n| LLaDA-MRO-RL-150 | 35.2     | 32.3     |\n\nWe have conducted an ablation study by varying the upper bound of the perplexity reward over a wide range [50,80,100,130,150] and observed the following:\n\n- Performance remains relatively stable within the range of 80 to 130, indicating the robustness of our method and the ease of tuning this hyperparameter.\n- A lower bound may suppress meaningful differences in reward, while a higher bound can introduce excessive variance, both of which can negatively affect training stability. \n\nWe promise to add this discussion to the revision.\n\n> *Q2: In Equation 6, what is the rationale for considering only $r_0$ in the inter-sequence correlation rewards? 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.*\n\nThanks for your insightful comment! **We would like to clarify that the inter-sequence correlation reward is designed not to measure similarity between sequences at different denoising steps, but rather to assess whether sequences generated under different denoising steps can collaborate to produce a high-quality final rationale**. To capture this cooperative effect, we adopt a **delayed reward** formulation, using only the final reward $r_0$ as a proxy for overall reasoning quality. This approach is inspired by established practices in reinforcement learning literature, where delayed rewards are commonly used to encourage an agent to explore effective an action sequence that lead to a desirable long-term outcome. In our case, the use of $r_0$ allows us to encourage early-stage actions that contribute positively to the final outcome, without being distracted by intermediate, potentially noisy rewards. \n\n> *Q3: In 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?*\n\nThanks for your insightful comment! We provide both theoretical analysis and empirical evidence showing that SGRO is more effective than reward normalization in reducing reward variance within our MRO.\n\n1. Our analysis in Appendix B establishes two key results that jointly explain why SGRO is more effective than per-dimension reward normalization in mitigating reward variance.\n   \n   - **Variance Amplification from Potential-Based Shaping (Property 1)**  \n       Lemma 1 proves that under potential-based reward shaping  \n       $R(s_t,a_t)=\\hat R(s_t,a_t)+\\lambda\\Phi(s_{t+1})-\\Phi(s_t)$  \n       the expected reward remains unchanged, yet the variance strictly increases:  \n       $\\text{Var}(R(s_t,a_t))>\\text{Var}(\\hat R(s_t,a_t))$  \n       The increased variance originates from the covariance term    \n       $-2\\lambda\\times \\text{Cov}(\\Phi(s_{t+1}),\\Phi(s_t))$  \n       which remains **invariant under any affine transformation of the rewards, including min-max normalization, z-score scaling, or simple rescaling. Therefore, reward normalization alone cannot mitigate this covariance-induced variance**.\n   \n   - **Variance Reduction via SGRO (Property 2)**  \n       Lemma 2 shows that SGRO mitigates this variance by grouping every $w$ consecutive steps and applying the potential difference only once per group:  \n       $R^{(w)}(s_t,a_t)=\\sum_{i=0}^{w-1}\\hat R(s_{t+i},a_{t+i})+\\lambda\\Phi(s_{t+w})-\\Phi(s_t)$  \n       By increasing the temporal gap between evaluations of $\\Phi$, SGRO effectively reduces the dominant covariance term:  \n       $\\bigl|\\text{Cov}(\\Phi(s_{t+w}),\\Phi(s_t))\\bigr|<\\bigl|\\text{Cov}(\\Phi(s_{t+1}),\\Phi(s_t))\\bigr|$   \n       leading to a strict reduction in reward variance:  \n       $\\text{Var}(R^{(w)}(s,a))<\\text{Var}(R(s,a))$\n     \n     In summary, while reward normalization only rescales individual reward components and leaves the covariance structure unchanged, SGRO directly targets and reduces the source of variance introduced by potential-based shaping.\n\n2. We have further conducted empirical comparisons between SGRO and reward normalization on the MATH500 and GPQA datasets. **As shown below, SGRO consistently yields superior performance, further validating its effectiveness in variance mitigation**:\n   \n   | Model/Length      | MATH500  |          | GPQA     |          |\n   | ----------------- |:-------- |:--------:|:-------- |:--------:|\n   |                   | 256      | 512      | 256      | 512      |\n   | LLaDA             | 33.2     | 34.4     | 29.2     | 30.3     |\n   | LLaDA-MRO w/ Norm | 32.8     | 35.0     | 30.0     | 32.8     |\n   | LLaDA-MRO w/ SGRO | **34.2** | **36.2** | **32.1** | **34.3** |\n\nWe promise to add this analysis and the empirical comparison in the revision.\n\n---\n\nWe sincerely thank you for your positive feedback on our paper!  \nThank you once again for your time.\n\nBest,\n\nAuthors"}, {"review_id": "cZ4s8AAstE", "reviewer": "Reviewer_CEHF", "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.", "questions": "See weakness above.", "limitations": "Yes.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Pros:\n1. 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. \n2. Extensive experiments across five reasoning benchmarks (e.g., GSM8K, MATH500, GPQA) demonstrate that MRO substantially improves performance and sample efficiency.\n3. The definition of the problem is clear. The authors pinpoint token-correlation deficiencies as a root cause of poor DLM reasoning.\n\nCons:\n1. 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?\n2. 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.", "was_revised": true, "final_justification": "I decide to keep my score.", "rebuttal": "Dear Reviewer CEHF,\n\nWe appreciate the reviewer’s constructive and thoughtful feedback.\n\nThanks for your recognition that our paper \"studied an important problem\" and clearly defines the core issue by \"pinpointing token-correlation deficiencies as a root cause of poor DLM reasoning\". We are also grateful for your findings that our MRO addresses this limitation by \"designing multiple reward signals tailored to different aspects of reasoning\", and that our \"extensive experiments across five reasoning benchmarks\" demonstrate that MRO \"substantially improves performance and sample efficiency\".\n\nWe provide explanations for the main points that you are concerned about below.\n\n---\n>*W1: The reliance on repeated forward passes, especially in test-time scaling (TTS) and rejection sampling, increases inference and training overhead. Can the authors comment on this?*\n\nThanks for your insightful comment! Below, we clarify the issue related to inference/training time and present empirical results to support our claims. **Taken together, we demonstrate that our MRO is highly practical and well-suited for real-world applications**.  \n\n- Regarding your concern about inference-time overhead, we first would like to clarify that in both the rejection sampling and RL fine-tuning settings, multi-reward computation and policy gradient updates are performed **only during training**. Once the model parameters are updated, decoding proceeds identically to the standard LLaDA baseline \\[22], with **no additional reward computation or repeated forward passes** involved during inference.\n\n- Additionally, to further address your concern about the overhead introduced by our reward computation, we provide a detailed breakdown of the associated computational cost in the table below, where $\\Theta$ denotes the model size and $L$ denotes the sequence length. To reduce the overall burden, we adopt several efficiency-oriented strategies. Specifically, **we leverage SGRO to sparsely compute rewards across decoding steps, significantly limiting the number of evaluations needed**. For the Perplexity Reward, we employ a lightweight model (GPT-2-small) to minimize FLOPs. We also employ GPU-parallelized computation to process multiple reward evaluations concurrently during training, substantially reducing wall-clock latency.\n\n    |Reward Type| Model Used | Formula (Per Occurrence) | Occurrence Count (w/ SGRO) |Total FLOPs (Estimated)|\n    | - | :-: | :-: | :-: | :-: |\n    |Token Verification Reward| LLaDA-8B  | $2 \\cdot \\Theta_{\\text{LLaDA}} \\cdot L \\cdot (N + 1)$ | $\\frac{T}{w} = \\frac{128}{w}$    | $\\frac{128}{w} \\cdot 1.23 \\times 10^{13}$ |\n    |Perplexity Reward| GPT-2-Small  | $2 \\cdot \\Theta_{\\text{GPT-2}} \\cdot L$ | $\\frac{T}{w} = \\frac{128}{w}$    | $\\frac{128}{w} \\cdot 6.0 \\times 10^{10}$  |\n\n\n- We have also conducted a comparison of both inference and training performance with and without our reward design, as shown in the tables below. The runtime is measured on eight A800 GPUs with a length of 256 under data parallelism, with each configuration run five times to report a consistent range of values. The results demonstrate that **our reward introduces only a modest computational overhead while delivering substantial performance improvements**, confirming its practicality and effectiveness.\n\n    |Method|Inference Time (MATH500)|Training Time|Score (MATH500)|\n    |-|:-:|:-:|:-:|\n    |LLaDA|0.35h~0.39h|-|33.2|\n    ||\n    |LLaDA-RS w/o MRO|0.35h~0.38h|10.3h~11.0h|33.0|\n    |LLaDA-RS w/ MRO|0.34h~0.37h|12.7h~13.1h|34.2|\n    ||\n    |LLaDA-RL w/o MRO|0.33h~0.37h|8.2h~8.4h|33.4|\n    |LLaDA-RL w/ MRO|0.36h~0.39h|9.6h~10.1h|35.2|\n    ||\n    |LLaDA-TTS w/o MRO|0.73h~0.81h|-|35.2|\n    |LLaDA-TTS w/ MRO|0.84h~0.85h|-|36.0|\n\n\n>*W2: 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.*\n\n\n|Model|AIME2024|AMC2023|\n|-|:-:|:-:|\n|LLaMA-3-8B-Instruction|6.7|30.0|\n|Qwen2.5-7B-Instruct|10.0|52.5|\n||\n|LLaDA|3.3|17.5|\n|LLaDA-TTS + MRO|6.7|22.5|\n|LLaDA-RS + MRO|**10.0**|**25.0**|\n|LLaDA-RL + MRO|6.7|20.0|\n\n- The results show that our proposed MRO remains effective on long-form reasoning tasks, consistently enhancing the reasoning capabilities of LLaDA even at a length of 768. Notably, MRO helps close the performance gap between DLMs and AR-style models. These findings further support that MRO can effectively address a fundamental weakness of DLMs, the lack of intra-/inter-sequence correlation modeling.\n\n- We observe that DLMs continue to underperform on long-form reasoning tasks, which we attribute to their relatively early stage of development compared to AR-based models. This observation underscores the significance of our work: MRO directly tackles one of the key challenges in DLM reasoning (i.e., the lack of token-level correlation modeling) and represents a meaningful step toward enhancing the reasoning capabilities of this emerging class of models.\n\n---\nWe sincerely appreciate your positive feedback and thank you again for your time.\n\nBest,\n\nAuthors"}, {"review_id": "MnqE1YFFPf", "reviewer": "Reviewer_9xuo", "summary": "The 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.", "questions": "1. 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?", "limitations": "Yes", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "## 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.", "was_revised": true, "final_justification": "I believe my concern was not adequately addressed in the rebuttal. As detailed in Replying to Rebuttal by Authors, issues such as imprecise mathematical notation and potentially misleading writing still remain. Therefore, I am keeping my score as Borderline Reject.", "rebuttal": "Dear Reviewer 9xuo,\n\nWe appreciate your recognition that \"our paper is the first to investigate how intra-sequence correlation affects reasoning performance in diffusion language models\", and that this investigated issue is \"crucial\".\n\nWe will provide explanations for the main points that you are concerned about.\n\n---\n\n>*W1: The mathematical definition of the x-axis in Figure 1 is not specified.*\n\nThanks for your helpful feedback! The x-axis in Figure 1 represents the intra- and inter-sequence correlation rewards defined in Section 4.2: the left part shows the cumulative intra-sequence rewards $R^{tv}\\_{t} + R^{ppl}\\_{t}$, and the right part shows the inter-sequence reward $R^{q}_{0}$. As the reward scales differ, we apply standardization for visualization. We will add this clarification to the introduction for better accessibility.\n\n>*W2: In Equation 3, it is unclear why the left-hand side of the first equality represents the marginal distribution of $r_0$, while the right-hand side corresponds to the joint distribution of $\\{r_0, r_1, \\dots, r_{T-1}\\}$. The same issue applies to the second equality.*\n\nWe apologize for the confusion caused by Equation 3. Below, we clarify the roles of the two distributions presented in the equation.\n\n- Left-hand side:  \n   The symbol $\\mathrm{Pr}'\\_{\\theta}(r_{0}\\mid p_{0},r_{T})$ is the decoding distribution, i.e., the marginal distribution over the final response $r_{0}$ after $T$ denoising steps.  This is the quantity we actually sample at inference time.\n\n- Right-hand side:  \n   The product\n\n   $\\prod\\_{t=T}^{1}\\mathrm{Pr}\\_{\\theta}(r_{t-1}\\mid p_{0},r_{t})$\n\n   is the joint denoising distribution over the entire latent sequence $\\{r_{T},r_{T-1},\\dots,r_{0}\\}$.\n   - Each factor $\\mathrm{Pr}\\_{\\theta}(r_{t-1}\\mid p_{0},r_{t})$ conditions only on the current masked sequence $r_{t}$ and predicts the corresponding masked tokens.  \n   - Marginalizing the latent variables $\\{r_{1},\\dots,r_{T-1}\\}$ yields the decoding distribution on the left-hand side.  \n\n- Second equality:  \n   The further factorization  \n   \n   $\\prod\\_{i=1}^{L_{r}}\\mathbf{1}[r_{t-1}^{i}=M]\\,\\mathrm{Pr}\\_{\\theta}(r_{t-1}^{i}\\mid p_{0},r_{t})$\n\n   explicitly reflects that, within each denoising step, masked tokens are predicted in parallel and independently, conditioned on the current input $r_t$. The indicator function $\\mathbf{1}[r_{t-1}^i = M]$ ensures that only masked tokens are predicted.\n\nIn essence, Equation 3 expresses that the decoding distribution over \\$r\\_0\\$ is constructed through a chain of \\$T\\$ denoising steps, where each step factorizes over masked tokens. This formulation is equivalent to that used in Sahoo et al. (2024), with the key difference that we employ an indicator function instead of an explicit conditional to present the token-wise structure more concisely. A similar formulation also appears in the original LLaDA paper, which we adopt here for consistency.\n\nWe promise to add this detailed explanation and revise the presentation in the revision to prevent potential confusion.\n\n>*W3: Token Verification Reward (TVR) lacks theoretical grounding.*\n\nWe provide a theoretical proof for the TVR by analyzing it from the perspective of mutual information, showing that it effectively promotes intra-sequence correlation.\n\nRecall the definition of the TVR at the denoising step $t$:\n\n$\nR^{\\text{tv}}\\_{t}= \\frac{1}{N}\\sum\\_{n=1}^{N}\\Pr\\_{\\theta}(r\\_{t}^{m\\_{n}}\\mid p\\_{0},r\\_{t-1}/\\{r\\_{t-1}^{m\\_{n}}\\}),\n$\n\nwhere $M = \\{m_1, \\dots, m_N\\}$ is the set of masked token indices, $r_{t-1}^{M}$ is the set of the predicted tokens at these positions, and $r_{t-1} / r_{t-1}^{M}$ is the set of unmasked token set. The joint probability over the masked positions can be factorized autoregressively (within the masked set) as:\n\n$\n\\Pr_{\\theta}(r_{t-1}^{M} \\mid p_0, r_t)\n= \\prod_{n=1}^{N} \\Pr_{\\theta}(r_{t-1}^{m_n} \\mid p_0, r_{t-1} / r_{t-1}^{M}, r_{t-1}^{<m_n}),\n$\n\nwhere the rest of the sequence is fixed and only the masked tokens are predicted.\n\nTVR approximates this joint modeling by computing **leave-one-out log-probabilities**:\n\n$\nR^{\\text{tv}}\\_{t}= \\frac{1}{N} \\sum_{n=1}^{N} \\log \\Pr_{\\theta}(r_{t}^{m_n} \\mid p_0, r_{t-1} / r_{t-1}^{m_n})\n$\n\nNow define the empirical average pairwise mutual information (PMI) among the masked tokens:\n\n$\n\\mathrm{PMI}\\_{\\text{avg}} = \\frac{2}{N(N-1)} \\sum_{1 \\le i < j \\le N} \\mathrm{I}(r_{t-1}^{m_i}; r_{t-1}^{m_j} \\mid p_0, r_{t-1} / r_{t-1}^{M})\n$\n\nUsing a standard second-order Taylor expansion around the independence assumption (as in energy-based models), this quantity can be approximated as:\n\n$\n\\mathrm{PMI}\\_{\\text{avg}} \\approx \\frac{2}{N(N-1)} \\sum_{i < j} \\log \\frac{\\Pr_{\\theta}(r_{t-1}^{m_i}, r_{t-1}^{m_j} \\mid \\cdot)}{\\Pr_{\\theta}(r_{t-1}^{m_i} \\mid \\cdot) \\Pr_{\\theta}(r_{t-1}^{m_j} \\mid \\cdot)}\n$\n\nImportantly, the leave-one-out log-probabilities computed in TVR serve as sufficient statistics for estimating these pairwise interactions. Therefore, **maximizing the average leave-one-out log-probability is first-order equivalent to maximizing $\\mathrm{PMI}_{\\text{avg}}$, thereby promoting stronger intra-sequence correlation**. We promise to add the theoretical grounding in the revision.\n\n>*W4: 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\nThanks for your insightful comment! We would like to clarify that **our reward categorization is intended to support the modeling and optimization of token correlations in DLMs, rather than to define independent reward signals**. Specifically, the intra-sequence rewards are designed to capture token-level correlations within a single sequence. In contrast, the inter-sequence reward is only observable after the entire denoising process is complete and is intended to assess the correlation among sequences generated at different denoising steps, i.e., whether they collectively produce an accurate final answer.\n\n>*W5: Computational overhead is unreported in Figure 2 and Figure 5.*\n\nBelow, we provide clarifications regarding the computational overhead reported in Figures 2 and 5. **Our results demonstrate that MRO is computationally efficient and well-suited for real-world deployment**.\n\n1. For the test-time scaling (TTS) experiments in Figure 2, we provide a detailed latency comparison between vanilla decoding and MRO, as shown in the table below. Given that TTS requires additional sampling and thus introduces time overhead, we also include a comparison with a confidence-based reward TTS baseline, which involves the same sampling cost as our MRO but does not require extra reward computation. From the results, we see that **while reward computation does introduce some overhead, it remains within a reasonable range and is justified by the substantial performance gains it brings**. Compared to the confidence-based reward TTS baseline, we further see that **when excluding the time cost introduced by sampling, our reward design incurs only minimal additional overhead while delivering substantial performance gains**.  \n\n    |Method|MATH500|  |GPQA|  |\n    |-|:-|:-:|:-|:-:|\n    ||Decoding Time |Score|Decoding Time|Score|\n    |Vanilla (LLaDA)| 0.35h~0.39h |33.2 | 0.16h~0.21h |29.2|\n    |LLaDA-TTS + Confidence-based Reward| 0.73h~0.81h |35.2 | 0.27h~0.34h |30.6| \n    |LLaDA-TTS + MRO| 0.84h~0.85h |36.0 | 0.31h~0.44h |34.6|\n\n2. **In Figure 5, we would like to clarify that in both the rejection sampling and RL fine-tuning settings, multi-reward computation and policy gradient updates are performed only during training**. Once the model parameters are updated, decoding proceeds identically to the standard LLaDA [22], with no additional reward computation or forward passes involved, and thus, the end-to-end decoding latency remains unaffected. This is further demonstrated in the table below, where the comparison confirms that our rejection sampling and RL setups do not introduce noticeable decoding-time latency on eight A800 GPUs.\n\n    |Method|MATH500 |  |GPQA|  |\n    |-|:-|:-:|:-|:-:|\n    ||Decoding Time |Score|Decoding Time|Score|\n    |Vanilla Decoding (LLaDA)| 0.35h~0.39h |33.2 | 0.16h~0.21h |29.2 |\n    |LLaDA-RS + MRO| 0.34h~0.37h |34.2 | 0.14h~0.18h |32.1 |\n    |LLaDA-RL + MRO| 0.36h~0.39h |35.2 | 0.16h~0.20h |33.8 |\n\n>*Q1: 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?*\n\nWe respectfully clarify that **policy gradient updates occur only during training, not at test time**. The TTS procedure serves purely as a diagnostic tool to evaluate reward signal quality without any backpropagation or parameter updates, as also noted in line 273. For clarity, we will move this explanation to a more prominent location. Indeed, our approach aligns with OpenAI's RLHF evaluation protocol \\[39], where reward models are first assessed via TTS before guiding methods like rejection sampling or RL. **Lines 304–305 aim to highlight that our TTS analysis validates the reward signal’s effectiveness for downstream policy optimization.**\n\n\n>*Q2: When evaluating the LLaDA-Instruct, did you enable semi-autoregressive sampling? If so, what block length did you use?*\n\nYes, we follow the official evaluation code released by LLaDA, which includes semi-autoregressive sampling by default. The block length is set according to the recommended configuration in the original LLaDA paper, as shown below.\n\n||GSM8K|MATH500|GPQA|Countdown|Sudoku|\n|-|:-:|:-:|:-:|:-:|:-:|\n|Block Length|8|64|64|64|64\n\n---\nWe truly appreciate your positive feedback on our paper!\n\nBest,\n\nAuthors"}, {"review_id": "cRAEpQcA3b", "reviewer": "Reviewer_Sm7r", "summary": "This 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.", "questions": "see weakness", "limitations": "yes", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "Overall, my main concern remains—the theoretical analysis is not sufficient to demonstrate why the approach works or why this design is an effective choice. The authors acknowledge that the reward design is not optimal and defer improvements to future work, which undermines the claimed contribution. After reading the discussion between the authors and other reviewers, I also agree that the structure of the paper is unclear and needs further refinement. Taking everything into consideration, I have decided to keep my score.", "rebuttal": "Dear Reviewer sm7r,\n\nWe appreciate the reviewer’s constructive and thoughtful feedback.\n\nWe would like to thank the reviewer for the positive feedback regarding \"the consistent performance improvements over strong DLM baselines across diverse reasoning benchmarks\". We especially appreciate the recognition of \"our method’s effectiveness in logic-intensive tasks such as Sudoku and Countdown, where token correlation plays a crucial role\".\n\nBelow, we explain the main points you are concerned about.\n\n---\n\n>*W1: 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\nThanks for your insightful feedback! We would like to clarify that the primary goal of our work is to improve DLM reasoning from the perspective of token correlation, rather than to design an optimal reward function. In this work, we formally define the token correlation problem and introduce intuitive formulations of intra-/inter-sequence correlation rewards as a first step toward addressing it. By optimizing these reward signals, we demonstrate that token correlation plays a crucial role in enhancing the reasoning capabilities of DLMs. As also mentioned by Reviewer3ZHw, we believe our contribution serves as a direction-setting effort, and the design of more optimal reward functions remains an open direction for future research.\n\n>*W2: What is the setup of AR models? Are they evaluated with the same prompt formats, decoding strategies, or number of shots?*\n\nOur evaluation setup for AR models follows the configuration reported in the LLaDA paper. We ensure the use of identical prompt formats and number of shots to maintain fair comparisons. The evaluation is conducted using the *lm-evaluation-harness* framework, which provides standard decoding strategies and shot settings for the benchmarks we use. Additionally, the resulting scores are consistent with those reported in prior literature.\n\n>*W3: MRO introduces considerable computational complexity due to repeated model evaluations, especially for the token verification reward that requires re-inference of multiple masked variants.*\n\nThanks for your insightful comment! Below, we clarify the issue related to inference/training time and present empirical results to support our claims. **Taken together, we demonstrate that our MRO is highly practical and well-suited for real-world applications**. \n\n- Regarding your concern about inference-time overhead, we first would like to clarify that in both the rejection sampling and RL fine-tuning settings, **multi-reward computation and policy gradient updates are performed only during training**. Once the model parameters are updated, decoding proceeds identically to the standard LLaDA baseline \\[22], with **no additional reward computation or repeated forward passes involved during inference**.\n\n- Additionally, to further address your concern about the overhead introduced by our reward computation, we provide a detailed breakdown of the associated computational cost in the table below, where $\\Theta$ denotes the model size and $L$ denotes the sequence length. To reduce the overall burden, we adopt several efficiency-oriented strategies. Specifically, **we leverage SGRO to sparsely compute rewards across decoding steps, significantly limiting the number of evaluations needed. For the Perplexity Reward, we employ a lightweight model (GPT-2-small) to minimize FLOPs. We also employ GPU-parallelized computation to process multiple reward evaluations concurrently during training, substantially reducing wall-clock latency**.\n\n    |Reward Type| Model Used | Formula (Per Occurrence) | Occurrence Count (w/ SGRO) |Total FLOPs (Estimated)|\n    | - | :-: | :-: | :-: | :-: |\n    |Token Verification Reward| LLaDA-8B  | $2 \\cdot \\Theta_{\\text{LLaDA}} \\cdot L \\cdot (N + 1)$ | $\\frac{T}{w} = \\frac{128}{w}$    | $\\frac{128}{w} \\cdot 1.23 \\times 10^{13}$ |\n    |Perplexity Reward| GPT-2-Small  | $2 \\cdot \\Theta_{\\text{GPT-2}} \\cdot L$ | $\\frac{T}{w} = \\frac{128}{w}$    | $\\frac{128}{w} \\cdot 6.0 \\times 10^{10}$  |\n\n\n- We have also conducted a comparison of both inference and training performance with and without our reward design, as shown in the tables below. The runtime was measured on eight A800 GPUs under data parallelism, with each configuration run five times to report a range of values. The results demonstrate that **our reward introduces only a modest computational overhead while delivering substantial performance improvements**, confirming its practicality and effectiveness.\n\n    |Method|Inference Time (MATH500)|Training Time|Score (MATH500)|\n    |-|:-:|:-:|:-:|\n    |LLaDA|0.35h~0.39h|-|33.2|\n    ||\n    |LLaDA-RS w/o MRO|0.35h~0.38h|10.3h~11.0h|33.0|\n    |LLaDA-RS w/ MRO|0.34h~0.37h|12.7h~13.1h|34.2|\n    ||\n    |LLaDA-RL w/o MRO|0.33h~0.37h|8.2h~8.4h|33.4|\n    |LLaDA-RL w/ MRO|0.36h~0.39h|9.6h~10.1h|35.2|\n    ||\n    |LLaDA-TTS w/o MRO|0.73h~0.81h|-|35.2|\n    |LLaDA-TTS w/ MRO|0.84h~0.85h|-|36.0|\n\n>*W4: The paper only compares the performance with AR models. They didn't compare MRO with other RL methods.*\n\nTo the best of our knowledge, our work is among the earliest to explore RL for DLMs. In the paper, we have made direct comparisons between MRO and other RL methods in the context of DLMs.\n\n- As of the NeurIPS2025 submission deadline (May 15), the only concurrent work that conducted RL training on DLMs is *d1-LLaDA* (released on arXiv on April 16). In Table 5 (Appendix), we provide a detailed comparison with d1-LLaDA *under the same training settings*. As shown in the results (copied below), our proposed MRO achieves consistently better performance than d1-LLaDA.\n    |Model/Length|GSM8K||MATH500||\n    |-|:-|:-:|:-|:-:|\n    ||256|512|256|512|\n    |d1-LLaDA|81.1|82.1|38.6|40.2|\n    |LLaDA-MRO | **82.5** | **82.9** | **39.4** | **42.6** |\n\n- We also conducted ablation studies using various reward designs. Specifically, we evaluate LLaDA with individual reward signals—namely, LLaDA-$R^{tv}\\_{t}$, LLaDA-$R^{ppl}\\_{t}$, and LLaDA-$R^{q}\\_{0}$, under RL training. These variants can serve as RL baselines to assess the contribution of each reward component. The results are summarized in the table below and further illustrated in Figure 6 (Appendix), with the partial scores copied here for reference (_as indicated in parentheses in the table_). From the results, we observe that our multi-reward design consistently outperforms the single-reward baselines, highlighting the effectiveness of combining multiple reward signals.\n    |Model||Reinforcement Learning||\n    |-|:-:|:-:|:-:|\n    ||MATH500|GPQA|Countdown|\n    |LLaDA-$R^{tv}_{t}$|36.2(+1.8)|32.7(+2.4)|25.3(+11.2)|\n    |LLaDA-$R^{ppl}_{t}$|33.6(-0.8)|30.8(+0.5)|18.9(+4.8)|\n    |LLaDA-$R^{q}_{0}$|34.8(+0.4)|31.2(+0.9)|23.5(+9.4)|\n    |LLaDA + MRO|**37.4(+3.0)**|**33.8(+3.5)**|**27.2(+13.1)**|\n\n>*W5: The study only looks at STEM tasks; it doesn't assess how well the suggested MRO applies to tasks involving common sense, instruction following, and so on.*\n\nThanks for the thoughtful suggestion! We have extended our evaluation to a broader range of tasks, including general knowledge question answering (MMLU), code generation (HumanEval), and instruction following (AlpacaEval2 and Arena-Hard), to verify the effectiveness of our proposed MRO further. The results demonstrate that **our MRO achieves strong robustness and can deliver consistent improvements (albeit to varying degrees) across these diverse tasks**. The results and details are as follows.\n\n|Model|MMLU|HumanEval|AlpacaEval2|Arena-Hard|\n|-|:-:|:-:|:-:|:-:|\n|LLaMA-3-8B-Instruction|68.4|59.8|25.3|22.3\n|Mistral-7B-Instruct|60.1|30.5|14.7|12.6\n|Deepseek-LLM-7b-Chat|48.2|26.2|11.2|10.3\n|Qwen2.5-7B-Instruct|71.9|56.7|27.8|25.2\n|||\n|LLaDA|65.5 |47.6 |16.3 |10.0 |\n|LLaDA-RS + MRO|67.5 | 48.1 | **20.2** | 12.3 |\n|LLaDA-RL + MRO|**68.2** | **50.0** | 19.4 | **15.7**|\n\n>*W6: 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\nAs you said, \"the experiments are limited to models up to 8B parameters\", **we would like to clarify that current open-source diffusion language models are only available up to this scale**. Nevertheless, to further demonstrate the robustness of our method, we additionally develop LLaDA-s1, a reasoning-specialized diffusion language model. As shown in Table 1, experimental results show that MRO consistently improves performance even on this customized model.\n\n>*W7: The paper only considers the experiments with a few shots. It is unclear about the performance of zero-shot.*\n\nThanks for the helpful feedback! In our main experiments, we adopted the same shot setting as our baseline methods to ensure a fair comparison. To address your concern, we have additionally conducted experiments under the zero-shot setting, and the results are provided below. **As shown, our method continues to deliver effective improvements even without in-context demonstrations (i.e., zero-shot), highlighting its robustness across different evaluation regimes**.\n\n|Model/Length|GSM8k (0-shot)||MATH500 (0-shot)||\n|-|:-|:-:|:-|:-:|\n||256|512|256|512|\n|LLaDA|75.7|78.4|28.2|30.0|\n|LLaDA-MRO-RS|77.5|79.6|29.8|32.4|\n|LLaDA-MRO-RL|**79.2**|**81.0**|**30.2**|**33.2**|\n\n\n--- \nWe sincerely thank you for your positive feedback on our paper!\n\nBest,\n\nAuthors"}]

Author rebuttal:

Meta-review: This paper introduces MRO, a multi-reward optimization framework to address token-correlation deficits in diffusion language models. The method is validated across diverse reasoning and general tasks, showing consistent gains and meaningful ablations. Reviewers raise concerns about mathematical precision, clarity of reward taxonomy, limited theory, and scalability, but these are presentation issues rather than fatal flaws. With camera-ready fixes (notation, compute reporting, clearer reward framing), the work makes a timely and solid contribution.
Decision: Accept (poster)