Reinforcing the Diffusion Chain of Lateral Thought with Diffusion Language Models

Q1: Have authors considered integrating auxiliary rewards alongside outcome-based RL?

In this work, we focus exclusively on outcome-based rewards to ensure a fair comparison with existing models. We agree that incorporating auxiliary rewards, such as step-level correctness, could potentially enhance performances. We will explore it in future work.

Q2: Provide concrete examples of failure cases

Thank you for the suggestion. We provide an example in Fig. 12(b), where DCoLT generates incorrect reasoning due to a miscalculation in an intermediate step. This type of error often arises when complex problems are not adequately decomposed into trackable steps.

As shown in Fig. 12(b), the failure case was corrected by increasing the generation length to $L=512$, which enables finer-grained decomposition of reasoning steps. This is also supported by the improved performance with longer generations shown in Tab 4 in the submission. We will add more failure cases and show how they can be remedied by the DCoLT.

Q3: Can you provide theoretical analysis of when and why DCoLT should outperform traditional approaches?

One of DCoLT's unique features is it learns the order of generated tokens based on the Plackett–Luce model, grounded in Luce’s Axiom of Choice [a] and Random Utility Model [b]. This enables the model to explicitly learn an optimal token generation order, rather than sticking to a fixed left-to-right order as in autoregressive models.

This flexibility is particularly advantageous in tasks where the natural problem-solving order differs from the left-to-right text order, such as in Sudoku. For example, as shown in Fig. 3, a natural order to complete the task is to fill in easier cells first, and then use them as the context to infer more difficult ones.

[a] Mollica C, et al. Remarkable properties for diagnostics and inference of ranking data modelling. British Journal of Mathematical and Statistical Psychology, 2022.

[b] Ma J, et al. Learning-to-rank with partitioned preference: Fast estimation for the Plackett-Luce model. International Conference on Artificial Intelligence and Statistics, 2021.

Q4: How does the computational cost compare to autoregressive methods with similar performance?

We provide an inference-time comparison between DCoLT and a recent autoregressive reasoning model, Deepseek-R1-Distill-Llama-8B. Specifically, we randomly sample 500 problems from GSM8K and report the average wall-clock inference time required to complete a full response. As shown in Tab. I, DeepSeek-R1-Distill-Llama-8B exhibits higher inference latency. In contrast, DCoLT requires less inference time, with even better accuracy.

Table I: Computation cost to solve one problem in GSM8K. The inference latency is averaged on 500 problems.

Q5: Can you demonstrate that the "lateral thinking" behavior is consistent and beneficial across a broader range of problems?

DCoLT demonstrates consistent improvements across both math and code generation tasks, as shown in our results on MATH, GSM8K, MBPP and HumanEval. These various domains involve complex multi-step reasoning, and the observed gains suggest that lateral exploration during reasoning is broadly beneficial.

While this work focuses on tasks with verifiable answers, we believe that the principles behind DCoLT generalize to broader settings. In future, we will extend DCoLT to open domains(e.g., scientific inquiry and creative generation) with trained reward models to guide lateral reasoning where verifiable answers are not available.

Q6: How does the approach scale to larger models and longer sequences?

During RL training, LLaDA+DCoLT that is optimized with a generation length of $L=256$ has already outperformed LLaDA and other RL algorithms (e.g., d1) under the same length. As shown in Tab.4, we showed that scaling the generation length further improves the model performance.

Regarding scalability with increasing model sizes, our approach demonstrates consistent gains on both smaller SEDD models (400M) and larger LLaDA models (8B). To our best knowledge, LLaDA 8B model is the largest open-source diffusion language model. We will test the ability of DCoLT scaling to larger DLMs once they are available.

Response to Strengths And Weaknesses

We thank the reviewers for acknowledging our contributions in modeling lateral thinking through reverse diffusion, achieving competitive results using only public data and moderate compute, and demonstrating versatility across different diffusion language models.

We agree that performance can be further improved with access to larger datasets and increased computational resources. We also recognize the importance of extending DCoLT to open-ended natural language tasks. To this end, we are exploring the integration of learned reward models to tackle these tasks, and will scale DCoLT in future works to further showcase the generality and potential of DCoLT across diverse reasoning challenges.

Q1: How does DCoLT's RL approach (e.g., GRPO) compare to other RL algorithms (e.g., PPO) in terms of training stability and performance?

In the paper, we adopt GRPO in DCoLT for a fair comparison with existing models (e.g., d1). Prior studies [a] have noted that PPO could suffer an unstable training process to learn a value network particularly in long reasoning chains. In contrast, GRPO completely eliminates the need for learning a value network, thereby demonstrating more stable training with better performances as evidenced by existing works [b].

[a] Yuan Y, et al. What’s Behind PPO’s Collapse in Long-CoT? Value Optimization Holds the Secret. 2025.

[b] Shao Z, et al. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. 2024.

Q2: What is the optimal length for tasks with different degrees of difficulty, and how does it balance computation and model performance?

We evaluate LLaDA+DCoLT 8B on MATH subsets grouped by different difficulty levels, following the original dataset annotations[c]. As shown in Tab. I, the accuracy on level-1 problems saturates with length 256, suggesting that shorter generations are sufficient for simpler questions. In contrast, the performance on harder problems (levels 2–5) continues to improve with longer generations, indicating that complex reasoning benefits from longer responses. This complies with our intuition that optimal generation length is dependent on difficulty levels — shorter for easier problems, and longer for more difficult ones. This also reflects a trade-off between computational cost and model performance, allowing generation length to be adapted to varying difficulty levels.

Table I: Accuracy of LLaDA+DCoLT 8B on MATH subsets across difficulty levels (5: the hardest and 1: the easiest) and generation lengths.

[c] Hendrycks D, et al. Measuring Mathematical Problem Solving With the MATH Dataset. NeurIPS 2021.

Q1 & Q2: Based on Appendix C.3, it appears that d1 trains separate models for GSM8K and MATH, each using its own corresponding training set. However, DCoLT seems to combine both datasets during training. This may not be a fair comparison. Can you please clarify whether the same training data is used for both models? I may be misunderstanding, but it seems like d1 only finetunes LoRA adapters, whereas DCoLT does full parameter updates, making the results difficult to compare.

A direct comparison with d1 using the same LoRA structure and dataset still reveals the superior performance by DCoLT. Specifically, we train DCoLT on GSM8K only, using LoRA with rank $r=128$ and scaling factor $\alpha=64$. During inference, we set the generation length to 256, the number of diffusion steps to 128, and the block length to 32 — mirroring d1’s configuration. In the same setting, DCoLT with LoRA achieves an accuracy of 84.7%, which outperforms d1 with both diffu-GRPO (79.8%) and d1-LLaDA (81.1%) setups, as shown in Tab. I. These results indicate that DCoLT still performs better than d1 with the same LoRA structure and the training data. We will add it in the final revision.

Table I: Comparison between d1 and DColT with LoRA on GSM8K. Both methods only use the GSM8K to train the LLaDA model with their respective RL approach.

Q3: The term “nonlinear generation” is ambiguous. For instance, in line 131, I don’t quite understand what it means. If the intent for it is to mean “non-left-to-right,” then it might be better to write non-left-to-right or something similar.

In our paper, "nonlinear generation" intends to mean "non-left-to-right". We will clarify this in the revision.

Q4: Could you comment on the coherence of the complete answer sequence when the answer is correct? I realize that a final-answer-based 0/1 reward may not encourage, or perhaps even discourage, coherent text sequence in the middle of the answer. It will be nice to present some concrete analysis on this front.

Although DCoLT is trained using only final-answer-based 0/1 rewards, we observe that it often produces coherent text sequences as shown in Fig. 8-9 of Appendix B. To further quantify this, we decompose the generated text responses on GSM8K into individual reasoning steps and evaluate their text coherency quality using a process reward model (Qwen2.5-Math-PRM-7B [a]). On correctly answered examples, LLaDA+DCoLT achieves an average step-wise reward of 0.96 (very close to the maximum reward of 1.0), better than LLaDA with an average reward of 0.94. This suggests that DCoLT produces not only correct answers but also coherent text sequence.

The reason can be that: unlike auto-regressive models that are pretrained to predict next token from readable text prefix, DLMs are pretrained [b] to recover coherent text from randomly masked inputs that are incomplete and unreadable. This training strategy endows DLMs with the unique robustness to generate coherent text outputs from corrupted text sequences amid the reverse diffusion process. This trained ability of DLMs allows the outcome-based RL to maintain the coherence of the final output sequences without relying on explicit rewards over intermediate diffusion steps to recover masked tokens.

[a] Zhang Z, et al. The Lessons of Developing Process Reward Models in Mathematical Reasoning. 2025.

[b] Nie S, et al. Large language diffusion models. 2025.

Q5: Please provide more details regarding the final answer-based reward for the question-answering datasets. Specifically, what is the filter used to extract the final answer for GSM8K-Aug, GSM8K, and MATH datasets for reward generation?

For GSM8K-Aug experiments with SEDD, we follow the approach used in DoT [c]: if the model outputs contain ####, we extract the number that follows as the predicted answer. Example responses following this pattern are provided in Fig. 6.

For GSM8K and MATH experiments with LLaDA, we follow DeepseekMath [d] and require the final answer to be put in the right response format (e.g., be enclosed in a LaTeX-style box \box{}). During reward computation, we extract the content inside the box and attempt to parse it into a number or symbolic expression. Then, we compare it against the ground-truth answer using symbolic math equivalence checks, following the evaluation code from the MATH dataset by Hendrycks et al [e]. Examples of accepted answer formats are shown in Fig. 8 and Fig. 9 of Appendix B.

[c] Ye J, et al. Diffusion of thought: Chain-of-thought reasoning in diffusion language models. NeurIPS 2024.

[d] Shao Z, et al. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. 2024.

[e] Hendrycks D, et al. Measuring Mathematical Problem Solving With the MATH Dataset. NeurIPS 2021.

Q6: Do you train the model to predict [object Object] tokens? In Figure 8, in step 64, I see that there are [object Object] tokens. Are these predicted by the model, or are they inserted automatically after the model predicts the first <|eot_id|>?

The <eos> tokens are predicted by the model. These predictions are considered as part of the trajectory to reinforce during training.

Q7: Have you tried DCoLT on Dream 7B? It would make the paper stronger if the method shows similar gains across the two models.

We appreciate the reviewer’s suggestion and conduct an additional experiment applying DCoLT to Dream 7B on GSM8K, using 64 denoising steps and a generation length of 256. As shown in Tab. II, DCoLT improves Dream 7B's accuracy to 80.53%, representing a +30.42% gain over the original model — comparable to the improvement observed on LLaDA.

Due to limited time during the rebuttal period, we have not yet conducted a experiment with 256 denoising steps, but we plan to include those results in the revised version of the paper.

Table II: Experimental results by applying DCoLT to Dream 7B on GSM8K with 64 denoising steps.

Q8: Since RL training is generally unstable, it would be good if you could provide average reward vs training progress curves.

Thanks for your constructive advice. During training, our model exhibits a stable increase in reward. Due to formatting constraints, we are unable to include these figures in the rebuttal. However, we will include training curves in the revised version to illustrate the training process.

W1: Some claims in this paper are problematic, e.g. for the generation length, I don’t think it plays an important role here. The performances with length from 256-512 are very close.

During RL training, LLaDA+DCoLT 8B is only optimized with a generation length of $L=256$, and it already outperforms LLaDA and other RL algorithms (e.g., d1) under the same length setting. In the submission, inspired by prior works on auto-regressive models, we also investigate the effect of generation length on performances in Tab. 4 as well.

Particularly, although the overall performance with varying lengths may be close, when they are applied to MATH subsets with various difficulty levels, LLaDA+DCoLT 8B with longer generations could perform much better on harder problems than that with shorter ones as shown in Tab. I. For example, on the hardest category with difficulty level 5, DCoLT improves the accuracy from 13.2% with a generation length of 128 to 21.5% with that of 512, a remarkable gain of 8.3% in accuracy.

Table I: Accuracy of LLaDA+DCoLT 8B on MATH subsets across difficulty levels (5: the hardest and 1: the easiest) and generation lengths.

W2: I’m not sure whether the method really optimizes the entire reverse process jointly (line 173-174), because in Algorithm1, the optimization is processed from 1 to n, which means getting loss for each denoising step separately.

We compute and accumulate the losses for all the denoising steps in the entire reverse diffusion process, and sum them up to update the model parameters. Note that the network weights were not updated until the losses from all the steps are accumulated to optimize the entire reverse process.

Moreover, each one-diffusion-step gradient accumulated in Line 26-28 of Algorithm 1 fulfills a memory-efficient strategy. As described in Line 149, we release the computational graph after each step to reduce memory consumption, which still optimizes the entire reverse process jointly with the accumulated gradient in a memory-efficient fashion.

Q: The method only trained the model using final-answer reward signals, with no supervision or constraint on the intermediate diffusion steps. This setup has been shown to cause unreadable intermediate steps in autoregressive models (e.g. DeepSeek-R1), and this is a reason why they have an additional SFT process. Despite this, attached examples in this paper show semantically meaningful and readable CoT like reasoning. I’m curious why the case is different with autoregressive models and outcome-based RL on diffusion language models that lead to interpretable intermediate reasoning.

Unlike auto-regressive models that are pretrained to predict next token from semantically meaningful prefix, DLMs are pretrained [a] to recover coherent text outputs from randomly corrupted inputs that are incomplete and unreadable with masked tokens. This training strategy endows DLMs with the unique robustness to generate semantically readable text outputs even from corrupted text sequences amid the reverse diffusion process. This trained ability of DLMs allows outcome-based RL to generate semantically meaningful reasoning texts without relying on additional SFT to recover the corrupted tokens during intermediate diffusion steps.

[a] Nie S, et al. Large language diffusion models. 2025.

Thanks for your feedback. Here we provide a quantitative evaluation on the robustness of the Diffusion Language Model (DLM) to substantiate this claim as requested by the reviewer.

One of the key features of DLMs is the theoretical noise schedule applied for training - it trains the model to robustly recover coherent text outputs from randomly corrupted inputs at various noise levels, ensuring that the model can generate semantically readable text outputs at the end of the reverse diffusion process when the noise level reaches 0.

In Tab. II, we use Qwen2.5-Math-PRM-7B [a] to evaluate the actual quality of the model outputs at each intermediate diffusion step, and compare it with the corresponding theoretical noise level. This aims to reveal the consistency of the actual quality with the theoretical noise level by the diffusion language model. For this, we report the average rewards at each step calculated on correctly answered GSM8K samples. During the reverse diffusion process, as the theoretical text cleanness (defined as 1 - noise level) increases from 0 to 1, the average reward steadily rises from 0.08 to 0.94 (very close to the maximum reward of 1.0). We also compute the correlation between the theoretical cleanness and the reward, reaching a high value of Pearson correlation 0.95.

This verifies the noise levels structured in the DLM provides a consistent prediction on the quality of intermediate diffusion outputs, and the robustness of the DLM being able to recover the semantically meaningful reasonings comes from the scheduled various noise levels during the training and the resultant high consistency with the measured reward on the actual quality. This eventually allows the DLM to generate semantically meaningful text with increasing reward of quality as the theoretical noise level reaches 0 at the end of the reverse diffusion process.

We hope this further analysis clears your concern and provides a quantitative analysis of DLM's robustness in generating semantically meaningful reasoning texts. Please let us know if you need more information.

Table II: Relationship between theoretical cleanness (=1-noise level) and the average reward of actual quality on correctly answered GSM8K samples after our outcome-based RL training. The result shows a Pearson correlation as high as 0.95 between the noise levels and the measured rewards on the intermediate diffusion outputs.

[a] Zhang, Z, et al, The Lessons of Developing Process Reward Models in Mathematical Reasoning