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

Dear Reviewer sm7r,

We appreciate the reviewer’s constructive and thoughtful feedback.

We 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".

Below, we explain the main points you are concerned about.

---

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.

Thanks 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.

W2: What is the setup of AR models? Are they evaluated with the same prompt formats, decoding strategies, or number of shots?

Our 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.

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.

Thanks 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.

• 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.

• 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.

  Reward Type	Model Used	Formula (Per Occurrence)	Occurrence Count (w/ SGRO)	Total FLOPs (Estimated)
  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}$
  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}$

• 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.

  Method	Inference Time (MATH500)	Training Time	Score (MATH500)
  LLaDA	0.35h~0.39h	-	33.2
  LLaDA-RS w/o MRO	0.35h~0.38h	10.3h~11.0h	33.0
  LLaDA-RS w/ MRO	0.34h~0.37h	12.7h~13.1h	34.2
  LLaDA-RL w/o MRO	0.33h~0.37h	8.2h~8.4h	33.4
  LLaDA-RL w/ MRO	0.36h~0.39h	9.6h~10.1h	35.2
  LLaDA-TTS w/o MRO	0.73h~0.81h	-	35.2
  LLaDA-TTS w/ MRO	0.84h~0.85h	-	36.0

W4: The paper only compares the performance with AR models. They didn't compare MRO with other RL methods.

To 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.

• 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.

  Model/Length	GSM8K		MATH500
  	256	512	256	512
  d1-LLaDA	81.1	82.1	38.6	40.2
  LLaDA-MRO	82.5	82.9	39.4	42.6

• 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.

  Model		Reinforcement Learning
  	MATH500	GPQA	Countdown
  LLaDA-$R^{tv}_{t}$	36.2(+1.8)	32.7(+2.4)	25.3(+11.2)
  LLaDA-$R^{ppl}_{t}$	33.6(-0.8)	30.8(+0.5)	18.9(+4.8)
  LLaDA-$R^{q}_{0}$	34.8(+0.4)	31.2(+0.9)	23.5(+9.4)
  LLaDA + MRO	37.4(+3.0)	33.8(+3.5)	27.2(+13.1)

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.

Thanks 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.

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.

As 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.

W7: The paper only considers the experiments with a few shots. It is unclear about the performance of zero-shot.

Thanks 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.

---

We sincerely thank you for your positive feedback on our paper!

Best,

Authors

Dear Reviewer sm7r,

Thank you again for your thoughtful feedback and for continuing to engage with our paper. Below, we address your remaining concerns regarding reward design, RL baseline, and computational overhead.

---

The paper frames intra- and inter-sequence correlation rewards fail to provide any theoretical analysis or empirical ablation to justify why these are the right or even effective choices. ... The work thus lacks both theoretical depth and genuine novelty in reward design.

Thank you for your feedback. We would like to clarify our response to the concern that "the intra- and inter-sequence correlation rewards lack theoretical analysis or empirical ablation to justify their effectiveness".

• Empirical Ablation. As shown in Figure 6 of the appendix, we have conducted extensive ablation studies to evaluate the effectiveness of our reward design. Specifically, we tested all combinations of the proposed reward components under both TTS and RL settings. These experiments consistently demonstrate that our full multi-reward design outperforms all single-reward baselines, confirming its empirical effectiveness. A subset of these results was also included in the rebuttal for reference.

• Theoretical Analysis. In addition to the empirical results, we also provide a theoretical justification for one of the core components, the Token Verification Reward, based on its relationship to pairwise mutual information, showing that it approximates intra-sequence correlation.

  Recall the definition of the TVR at the denoising step $t$:

  $R^{\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}}\}),$

  where $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:

  $\Pr_{\theta}(r_{t-1}^{M} \mid p_0, r_t) = \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}),$

  where the rest of the sequence is fixed and only the masked tokens are predicted.

  TVR approximates this joint modeling by computing leave-one-out log-probabilities:

  $R^{\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})$

  Now define the empirical average pairwise mutual information (PMI) among the masked tokens:

  $\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})$

  Using a standard second-order Taylor expansion around the independence assumption (as in energy-based models), this quantity can be approximated as:

  $\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)}$

  Importantly, 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.

There is no comparison with a broader set of RL methods or alternative reward functions.

Thank you for the valuable comment. In addition to the results presented in our initial rebuttal (including baselines using individual reward components such as $R^{\text{tv}}_t$ and $R^{\text{ppl}}_t$), we have further incorporated a broader set of RL methods and reward functions to better evaluate the generality and effectiveness of our proposed approach. These new baselines include:

• RM-Baseline-1: A REINFORCE-based method using rule-based rewards derived from format correctness and answer accuracy, following prior setups such as [1].

• RM-Baseline-2: A REINFORCE-based method that uses a general-purpose pretrained reward model (Skywork-Reward-LLaMA-3.1-8B-v0.2) to compute delayed scalar rewards.

• RL-Baseline-1: A variant inspired by GRPO, which samples multiple trajectories and computes an advantage estimate to replace the raw reward signal during policy optimization.

• RL-Baseline-2: A baseline based on REINFORCE++ [2], which employs more stable advantage estimation for improved gradient signals.

  Model		Reinforcement Learning
  	MATH500	GPQA	Countdown
  RM-Baseline-1	35.5(+1.1)	31.6(+1.3)	20.2(+6.1)
  RM-Baseline-2	33.8(-0.6)	27.1(-3.2)	13.1(-1.0)
  RL-Baseline-1	35.6(+1.2)	32.1(+1.8)	21.4(+7.3)
  RL-Baseline-2	36.4(+2.0)	31.7(+1.4)	22.7(+8.6)
  LLaDA + MRO	37.4(+3.0)	33.8(+3.5)	27.2(+13.1)

Due to the limited time available during the rebuttal phase, we focused on reproducing a representative subset of widely used RL baselines. Nevertheless, we believe these additional experiments are sufficient to demonstrate that our proposed reward design consistently outperforms not only d1-LLaDA but also a broader range of RL approaches and reward strategies. We will include these results and detailed discussions in the revised version for completeness.

The method introduces significant computational overhead, especially for the token verification reward, which requires repeated model evaluations. Authors' claim that this is only a training-time cost does not address the fundamental scalability issue, especially as model and sequence sizes grow.

To further address the concern regarding training-time scalability, we conducted additional experiments using 1K RL training samples under varying generation lengths: 64, 128, 256, 512, 1024, and 2048 tokens. The wall-clock training times are summarized in the table below:

As shown above, although the total training time increases with sequence length, the relative overhead introduced by the TV reward remains consistently modest, even for sequences as long as 2048 tokens (only 1.15× compared to a no-TV baseline).

This non-linear scaling is primarily due to two key optimizations:

• Using SGRO: Instead of computing the TV reward at every denoising step, we sub-sample a fixed number of steps during training, significantly reducing computational cost while preserving effective training signals.

• Using Batch-parallelized Masked Evaluation: Our implementation jointly batches all masked tokens for each denoising step, allowing for a single forward pass to evaluate multiple token positions in parallel. For example, for a denoising sequence like "ABC", we construct a batch such as [_BC, A_C, AB_] to efficiently evaluate the correlation of each token with just one pass through the model.

Together, these optimizations ensure that the TV reward remains computationally practical and scalable, even as model size and sequence length grow. We promise to clarify these implementation details and runtime behaviors in the revised version.

Reference:
[1] Deepseekmath: Pushing the limits of mathematical reasoning in open language models
[2] Reinforce++: A simple and efficient approach for aligning large language models

---

We sincerely thank you for your positive feedback on our paper!

Best,

Authors

Dear Reviewer Sm7r,

As the discussion period draws to a close, we would like to follow up once more to check if there are any remaining concerns or questions we can help clarify. We greatly value your feedback and have made a concerted effort to address each of your comments in detail.

We noticed that the current rating remains negative. If you feel that our responses have adequately addressed your concerns, we would be sincerely grateful if you could consider updating your score. Please feel free to let us know if there is anything that still requires clarification.

Thank you again for your time and thoughtful engagement.

Best regards,
The Authors

Dear Reviewer Sm7r,

As the discussion period draws to a close, we would like to check whether our rebuttal has adequately addressed your concerns. If you have any remaining questions or points that require clarification, please don't hesitate to let us know. Thank you once again for your time and engagement.

Best regards,
The Authors

Dear Reviewer 9xuo,

We 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".

We will provide explanations for the main points that you are concerned about.

---

W1: The mathematical definition of the x-axis in Figure 1 is not specified.

Thanks 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.

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.

We apologize for the confusion caused by Equation 3. Below, we clarify the roles of the two distributions presented in the equation.

• Left-hand side:
  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.

• Right-hand side:
  The product

  $\prod\_{t=T}^{1}\mathrm{Pr}\_{\theta}(r_{t-1}\mid p_{0},r_{t})$

  is the joint denoising distribution over the entire latent sequence $\{r_{T},r_{T-1},\dots,r_{0}\}$.

  • 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.
  • Marginalizing the latent variables $\{r_{1},\dots,r_{T-1}\}$ yields the decoding distribution on the left-hand side.

• Second equality:
  The further factorization

  $\prod\_{i=1}^{L_{r}}\mathbf{1}[r_{t-1}^{i}=M]\,\mathrm{Pr}\_{\theta}(r_{t-1}^{i}\mid p_{0},r_{t})$

  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.

In 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.

We promise to add this detailed explanation and revise the presentation in the revision to prevent potential confusion.

W3: Token Verification Reward (TVR) lacks theoretical grounding.

We provide a theoretical proof for the TVR by analyzing it from the perspective of mutual information, showing that it effectively promotes intra-sequence correlation.

Recall the definition of the TVR at the denoising step $t$:

$R^{\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}}\}),$

where $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:

$\Pr_{\theta}(r_{t-1}^{M} \mid p_0, r_t) = \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}),$

where the rest of the sequence is fixed and only the masked tokens are predicted.

TVR approximates this joint modeling by computing leave-one-out log-probabilities:

$R^{\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})$

Now define the empirical average pairwise mutual information (PMI) among the masked tokens:

$\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})$

Using a standard second-order Taylor expansion around the independence assumption (as in energy-based models), this quantity can be approximated as:

$\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)}$

Importantly, 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.

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.

Thanks 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.

W5: Computational overhead is unreported in Figure 2 and Figure 5.

Below, 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.

1. 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.

   Method	MATH500		GPQA
   	Decoding Time	Score	Decoding Time	Score
   Vanilla (LLaDA)	0.35h~0.39h	33.2	0.16h~0.21h	29.2
   LLaDA-TTS + Confidence-based Reward	0.73h~0.81h	35.2	0.27h~0.34h	30.6
   LLaDA-TTS + MRO	0.84h~0.85h	36.0	0.31h~0.44h	34.6

2. 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.

   Method	MATH500		GPQA
   	Decoding Time	Score	Decoding Time	Score
   Vanilla Decoding (LLaDA)	0.35h~0.39h	33.2	0.16h~0.21h	29.2
   LLaDA-RS + MRO	0.34h~0.37h	34.2	0.14h~0.18h	32.1
   LLaDA-RL + MRO	0.36h~0.39h	35.2	0.16h~0.20h	33.8

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?

We 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.

Q2: When evaluating the LLaDA-Instruct, did you enable semi-autoregressive sampling? If so, what block length did you use?

Yes, 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.

---

We truly appreciate your positive feedback on our paper!

Best,

Authors

Dear Reviewer 9xuo,

Thank you once again for your thoughtful response and continued engagement. Please allow me to offer a further clarification regarding your concern.

---

The authors appear to have misunderstood a critical issue I raised. ... I strongly encourage the authors to carefully compare the equation in the submission with what they claim to be “the same formulation” (i.e., Equation 6 in [1]) and reflect on the differences.

Thank you for your continued feedback. With this discussion, we now better recognize the source of the confusion and would like to offer a more precise clarification.

We fully agree that the expression you suggested,

$p(r_0) = \int p(r_0, r_1, \dots, r_T) \, d(r_1, \dots, r_T)$

is mathematically sound and correctly defines the decoding distribution as a marginalization over the full joint distribution. This formulation provides a rigorous foundation and allows for an equality to hold.

However, in practice, such integrals over the full denoising trajectory are computationally intractable during decoding. Instead, DLMs approximate the decoding distribution through a sequential denoising process:

$\prod_{t=T}^{1} p(r_{t-1} \mid r_t)$

This expression reflects the fact that generating $r_0$ requires iteratively denoising intermediate sequences $\{r_{T-1}, r_{T-2}, \dots, r_1\}$, and is a commonly adopted approach in practice.

That said, we now understand that the first equation in our original submission may have been misleading. Although we introduced a separate symbol $\mathrm{Pr}'(\cdot)$ to distinguish the decoding distribution from the learned denoising distributions $\mathrm{Pr}(\cdot)$, this notation was insufficient to clearly communicate the relationship between the marginal and the compositional approximation.

To address this, we will revise the formulation more explicitly in the final version. Specifically, we define:

$q(r_{T-1}, \dots, r_0 \mid r_T) = \prod_{t=T}^{1} p(r_{t-1} \mid r_t)$

where $q(\cdot)$ denotes the distribution over the entire denoising trajectory leading to $r_0$, and $p(\cdot)$ corresponds to the denoising distribution at each time step.

We truly appreciate your suggestion to view decoding from an integral-based perspective—it provides a more complete mathematical grounding. We will explore such a formulation in future work, and we believe it could lead to promising directions in improving the theoretical rigor and practical efficiency of DLM decoding.

We promise to revise the text accordingly to clearly differentiate between the marginal decoding distribution and its sequential approximation, based on this discussion.

In addition, the paper explicitly categorizes reward signals as Intra-sequence Correlation and Inter-sequence Correlation, placing the final-answer reward under the latter. However, based on both my understanding and the authors’ rebuttal, the final-answer reward in fact optimizes both intra- and inter-sequence aspects. This inconsistency is confusing, and the rebuttal did not attempt to resolve it. Instead, it confirms that my concern was well-founded.

Through this discussion, we recognize that the current naming may have caused confusion, and we appreciate your feedback. In the revised version, we will reorganize the section headings and explicitly clarify the relationship between the different rewards, as well as their distinct optimization objectives, based on the insights from this discussion.

---

We sincerely appreciate your positive feedback and thank you again for your time.

Best,

Authors

Dear Reviewer CEHF,

We appreciate the reviewer’s constructive and thoughtful feedback.

Thanks 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".

We provide explanations for the main points that you are concerned about below.

---

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?

Thanks 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.

• 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.

• 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.

  Reward Type	Model Used	Formula (Per Occurrence)	Occurrence Count (w/ SGRO)	Total FLOPs (Estimated)
  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}$
  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}$

• 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.

  Method	Inference Time (MATH500)	Training Time	Score (MATH500)
  LLaDA	0.35h~0.39h	-	33.2
  LLaDA-RS w/o MRO	0.35h~0.38h	10.3h~11.0h	33.0
  LLaDA-RS w/ MRO	0.34h~0.37h	12.7h~13.1h	34.2
  LLaDA-RL w/o MRO	0.33h~0.37h	8.2h~8.4h	33.4
  LLaDA-RL w/ MRO	0.36h~0.39h	9.6h~10.1h	35.2
  LLaDA-TTS w/o MRO	0.73h~0.81h	-	35.2
  LLaDA-TTS w/ MRO	0.84h~0.85h	-	36.0

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.

• 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.

• 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.

---

We sincerely appreciate your positive feedback and thank you again for your time.

Best,

Authors

Dear Reviewer 9WDq,

We sincerely thank the reviewer for your positive and insightful feedback.

We 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".

We will provide explanations for the main points that you are concerned about.

---

W1: The main contributions of the paper can be clearly articulated in the introduction.

Thanks 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.

We 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.

W2: It is recommended that the paper include brief pseudocode for algorithms such as SGRO.

Thank 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:

Input: Pre-trained DLM $π\_θ$; Group Size $w$; Discount $\lambda$; Learning Rate $\eta$; Response Length $L$
Output: Fine-tuned DLM $π\_{θ^*}$

for each prompt p_0 in dataset:
    # 1. Roll-out full denoising trajectory
    r_T = FULL_MASK(L)
    states, actions = [], []
    for t = T downto 1:
        r_t_1 ~ pi_theta(·|p_0, r_t)
        states.append((p_0, r_t, t))   # s_t
        actions.append(r_t_1)        # a_t

    # 2. SGRO: compute rewards and returns per group
    grad = 0.0
    for g in range(0, T, w):
        start, end = g, min(g + w, T)
        # 2-a Intra-sequence rewards inside this group
        R_intra = 0.0
        for k in range(start, end):
            R_intra += TokenVerificationReward(p_0, actions[k])
            R_intra += PerplexityReward(actions[k])
        # 2-b Inter-sequence reward (only if the group ends at t=0)
        R_q = 0.0
        if end == T:  # covers final step t=0
            R_q = TaskAccuracyAndFormatReward(actions[end-1])
        # 2-c Potential-based shaping
        potential_start = potential_function(states[start])
        potential_end   = potential_function(states[end-1])
        shaping = _lambda * potential_end - potential_start

        # group return
        R_group = R_intra + R_q + shaping

        # 2-d REINFORCE update for steps in this group
        for k in range(start, end):
            log_prob = log pi_theta(actions[k] | states[k])
            grad += ∇_θ log_prob * R_group

    # 3. Parameter update
    fine_tuned_theta = _theta + _eta * grad

W3&Q4: There are some typos in this paper.

We appreciate your comment and have carefully proofread the manuscript to correct all typos in the revision.

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?

We 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:

• Performance remains relatively stable within the range of 80 to 130, indicating the robustness of our method and the ease of tuning this hyperparameter.
• 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.

We promise to add this discussion to the revision.

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.

Thanks 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.

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?

Thanks 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.

1. 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.

   • Variance Amplification from Potential-Based Shaping (Property 1)
     Lemma 1 proves that under potential-based reward shaping
     $R(s_t,a_t)=\hat R(s_t,a_t)+\lambda\Phi(s_{t+1})-\Phi(s_t)$
     the expected reward remains unchanged, yet the variance strictly increases:
     $\text{Var}(R(s_t,a_t))>\text{Var}(\hat R(s_t,a_t))$
     The increased variance originates from the covariance term
     $-2\lambda\times \text{Cov}(\Phi(s_{t+1}),\Phi(s_t))$
     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.

   • Variance Reduction via SGRO (Property 2)
     Lemma 2 shows that SGRO mitigates this variance by grouping every $w$ consecutive steps and applying the potential difference only once per group:
     $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)$
     By increasing the temporal gap between evaluations of $\Phi$, SGRO effectively reduces the dominant covariance term:
     $\bigl|\text{Cov}(\Phi(s_{t+w}),\Phi(s_t))\bigr|<\bigl|\text{Cov}(\Phi(s_{t+1}),\Phi(s_t))\bigr|$
     leading to a strict reduction in reward variance:
     $\text{Var}(R^{(w)}(s,a))<\text{Var}(R(s,a))$

     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.

2. 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:

   Model/Length	MATH500		GPQA
   	256	512	256	512
   LLaDA	33.2	34.4	29.2	30.3
   LLaDA-MRO w/ Norm	32.8	35.0	30.0	32.8
   LLaDA-MRO w/ SGRO	34.2	36.2	32.1	34.3

We promise to add this analysis and the empirical comparison in the revision.

---

We sincerely thank you for your positive feedback on our paper!   Thank you once again for your time.

Best,

Authors

Dear Reviewer 9WDq,

Please review the authors’ rebuttal at your earliest convenience. If you have further questions, use the discussion forum to engage with the authors, and kindly update your review and score as needed.

Thank you for your time and service.

AC

Dear Reviewer 9WDq,

We sincerely appreciate the time and thoughtful feedback you provided on our paper.

During the rebuttal phase, we have carefully addressed your comments and suggestions, including:

• Providing new experimental results on MRO under different upper bounds of the perplexity reward. We showed that performance remains relatively stable within the range of 80 to 130, demonstrating the robustness of our method and the ease of tuning this hyperparameter.

• Adding a comparison between SGRO and reward normalization on the MATH500 and GPQA datasets. As shown in our results, SGRO consistently outperforms reward normalization, further validating its effectiveness in mitigating reward variance.

• Offering a theoretical justification showing why SGRO is more effective than reward normalization in variance control.

• Including a pseudocode implementation of SGRO, which we promise to incorporate into the revised version for greater clarity and reproducibility.

• Addressing other clarifications as requested.

We hope that these additions and responses have adequately addressed your concerns. If so, we would be sincerely grateful if you would consider revisiting your score based on the updated content.

Thank you once again for your valuable time and constructive feedback.

Best regards,
The Authors

Dear Reviewer 9WDq,

As the discussion period draws to a close, we would like to follow up once more to check if there are any remaining concerns or questions we can help clarify. We greatly value your feedback and have made a concerted effort to address each of your comments in detail.

We noticed that the current rating remains negative. If you feel that our responses have adequately addressed your concerns, we would be sincerely grateful if you could consider updating your score. Please feel free to let us know if there is anything that still requires clarification.

Thank you again for your time and thoughtful engagement.

Best regards,
The Authors

Dear Reviewer 3ZHw,

We sincerely appreciate your constructive and thoughtful feedback.

We 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.

Below, we explain the main points you are concerned about.

---

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.

Thanks 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.

1. 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.

   Method	MATH500		GPQA
   	Decoding Time	Score	Decoding Time	Score
   Vanilla Decoding (LLaDA)	0.35h~0.39h	33.2	0.16h~0.21h	29.2
   LLaDA-RS + MRO	0.34h~0.37h	34.2	0.14h~0.18h	32.1
   LLaDA-RL + MRO	0.36h~0.39h	35.2	0.16h~0.20h	33.8

2. 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.

   Method	MATH500		GPQA
   	Decoding Time	Score	Decoding Time	Score
   Vanilla (LLaDA)	0.35h~0.39h	33.2	0.16h~0.21h	29.2
   LLaDA-TTS + Confidence-based Reward	0.73h~0.81h	35.2	0.27h~0.34h	30.6
   LLaDA-TTS + MRO	0.84h~0.85h	36.0	0.31h~0.44h	34.6

W2: This proposed approach generally mitigates the decoding problem of DLMs and could benefit from further verification on other tasks.

Thanks 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.

1. 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.

2. 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.

3. 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.

W3: The paper could benefit from adding more baselines that require training on the dataset in order to verify the proposed method better.

Thanks 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:

1. 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.

2. 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.

3. 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.

The results are summarized below. For consistency, the maximum length is set to 512 for MATH500 and GPQA, and to 128 for Countdown.

From the results, we can observe the following:

• 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.

• 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.

S2: Could the authors provide analysis on the intra-/inter- sequence correlations with and without the proposed method?

Thanks 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.

We promise to add these related discussions in the revised version.

---

We sincerely thank you for your positive feedback on our paper!

Best,

Authors