PaperlayerView exact revision on arXiv

Source-linked AI summary

Reinforcing the Diffusion Chain of Lateral Thought with Diffusion Language Models

Zemin Huang, Zhiyang Chen, Zijun Wang, Tiancheng Li, Guo-Jun Qi

arXiv:2505.10446v3cs.CL

TL;DR

Existing diffusion-language-model RL methods overlook intermediate diffusion steps as reasoning actions, limiting how lateral reasoning is trained. DCoLT instead optimizes the full reverse-diffusion trajectory with final-answer rewards, using SEDD and LLaDA-specific policies. Across math and code benchmarks, the reinforced models achieve strong results, while the approach remains limited by data, compute, and reliance on verifiable rewards.

  • Problem

    Prior diffusion-language-model RL methods optimize final responses while ignoring intermediate diffusion steps as lateral thought.

  • Method

    DCoLT treats reverse-diffusion intermediate steps as actions and applies outcome-based reinforcement learning to the full reasoning trajectory, with policies for SEDD and LLaDA.

  • Results

    DCoLT-reinforced diffusion models outperform other diffusion models trained with SFT, RL, or both across math and code-generation benchmarks.

  • Takeaways & Limitations

    Reinforcing non-linear intermediate diffusion reasoning can improve final-answer performance without explicit intermediate reasoning supervision.

  • Takeaways & Limitations

    DCoLT is validated only on tasks with verifiable reward functions and remains constrained by limited training data and compute.

Abstract

from arXiv · show

We introduce the Diffusion Chain of Lateral Thought (DCoLT), a reasoning framework for diffusion language models. DCoLT treats each intermediate step in the reverse diffusion process as a latent "thinking" action and optimizes the entire reasoning trajectory to maximize the reward on the correctness of the final answer with outcome-based Reinforcement Learning (RL). Unlike traditional Chain-of-Thought (CoT) methods that follow a causal, linear thinking process, DCoLT allows bidirectional, non-linear reasoning with no strict rule on grammatical correctness amid its intermediate steps of thought. We implement DCoLT on two representative Diffusion Language Models (DLMs). First, we choose SEDD as a representative continuous-time discrete diffusion model, where its concrete score derives a probabilistic policy to maximize the RL reward over the entire sequence of intermediate diffusion steps. We further consider the discrete-time masked diffusion language model -- LLaDA, and find that the order to predict and unmask tokens plays an essential role to optimize its RL action resulting from the ranking-based Unmasking Policy Module (UPM) defined by the Plackett-Luce model. Experiments on both math and code generation tasks show that using only public data and 16 H800 GPUs, DCoLT-reinforced DLMs outperform other DLMs trained by SFT or RL or even both. Notably, DCoLT-reinforced LLaDA boosts its reasoning accuracy by +9.8%, +5.7%, +11.4%, +19.5% on GSM8K, MATH, MBPP, and HumanEval.

1 Introduction

The paper introduces DCoLT, which reinforces non-linear intermediate reasoning in diffusion language models using final-answer rewards. It applies this framework to SEDD and LLaDA and reports strong math and code-generation results.

  • Human lateral thinking develops ideas non-sequentially and without initially requiring intact linguistic structure, unlike vertical step-by-step reasoning.
  • DCoLT treats reverse-diffusion intermediate steps as lateral thought and reinforces the entire trajectory using a reward based only on final-response correctness.
  • SEDD uses its predicted concrete scores to define a probabilistic policy that optimizes final-answer rewards across intermediate diffusion steps.
  • LLaDA learns token-unmasking order through a Plackett–Luce ranking policy, selecting top-K tokens to retain at each diffusion step.
  • 96.2% accuracy on Sudoku 4×4 and 57.0% on GSM8K-Aug are achieved by SEDD with DCoLT.
  • Using public data and 16 H800 GPUs, DCoLT-reinforced LLaDA reaches 88.1% on GSM8K, 44.6% on MATH, 51.6% on MBPP, and 59.1% on HumanEval.

2 Related Works

Prior diffusion-language-model RL methods optimize final responses while overlooking intermediate diffusion steps as reasoning actions. This work instead reinforces the full reverse-diffusion process to train lateral reasoning.

  • Diffusion language models include continuous and discrete variants, with masked diffusion models emerging as a strong discrete approach.
  • Existing reinforcement-learning methods primarily target autoregressive language models, while diffusion-model methods have explored concrete-score policies or sampled-response probabilities.
  • Prior diffusion-model RL methods optimize only the final generated response and ignore intermediate diffusion steps as lateral thought.
  • The paper explicitly reinforces the reverse-diffusion process and reports that this can produce correct final answers while outperforming many CoT models.

3 Methods

DCoLT treats the full reverse-diffusion trajectory as a lateral-thinking process, allowing bidirectional, format-free, nonlinear intermediate reasoning and training it with final-outcome rewards. The framework applies this training to continuous-time SEDD and masked discrete-time LLaDA, including learned token-unmasking order for LLaDA.

  • Diffusion Chain of Lateral Thought: DCoLT defines all intermediate reverse-diffusion results x0:N as a reasoning trajectory optimized toward the final response.Unlike conventional diffusion training, it jointly optimizes the entire reverse process using rewards associated with completed answers.
  • DCoLT features: Bidirectional reasoning lets each token attend to preceding and succeeding content, enabling global refinement during generation.This differs from causal CoT, where earlier thoughts remain fixed after generation.
  • DCoLT features: Format-free reasoning permits incomplete or incorrectly formatted intermediate steps before convergence to complete responses.The relaxation is intended to support more divergent intermediate thought patterns.
  • DCoLT features: Nonlinear generation predicts tokens at various positions rather than producing them strictly left to right.The process can begin with keywords or critical points and refine surrounding details over subsequent diffusion steps.
  • Training framework: Outcome-based reinforcement learning treats x1:N as sequential actions and assigns rewards from the verifiable correctness of the final result.The policy at each denoising step samples the next intermediate output, without explicit supervision on intermediate steps.
  • LLaDA implementation: For LLaDA, the Unmasking Policy Module ranks masked tokens with hθ,n and samples a top-K unmasking list before LLaDA blocks generate those tokens.The Plackett–Luce ranking policy makes unmasking order part of the learnable reasoning process; the resulting model is named LLaDOU.

4 Experiments

Experiments evaluate DCoLT on SEDD and LLaDA across math and code tasks, showing strong performance from outcome-based reinforcement learning and ordered unmasking. Ablations further examine the roles of the unmasking policy, adaptive conditioning, and generation length.

  • SEDD experiments: SEDD + DCoLT reaches 96.2% on Sudoku 4 × 4 and 57.0% on GSM8K-Aug, outperforming the compared CoT and DoT methods on both tasks.The Sudoku result is the best among all methods, while GSM8K-Aug also exceeds the DoT and CoT counterparts.
  • LLaDA experiments: LLaDOU consistently achieves the best performance among compared DLMs, reaching 88.1% on GSM8K and 44.6% on MATH.On MATH, it exceeds LLaDA 8B by 5.7% and Dream 7B by 1.7%.
  • LLaDA experiments: LLaDOU matches DeepseekMath’s 88.1% GSM8K performance using 15K public training samples, despite DeepseekMath using nearly two orders of magnitude more data for RL.The comparison highlights performance achieved with substantially less training data.
  • LLaDA experiments: LLaDOU achieves 51.6% on MBPP and 59.1% on HumanEval without code for supervised training, outperforming other DLMs and most compared autoregressive models.DeepseekCoder is the exception, benefiting from approximately 150× more instruction-tuning data.
  • Ablation studies: Training the UPM alone improves accuracy from 47.27% to 69.24%, while jointly training the LLaDA component with RL raises it to 81.06%.Removing adaptive layer normalization lowers accuracy to 80.53%, supporting the role of diffusion-step and mask conditioning.
  • Generation length: A generation length of 384 reaches 46.22% on MATH, 1.62% above length 256, while fine-tuning at length 512 further raises accuracy to 47.3%.The improvement from longer inference-only generations vanishes with further length increases, whereas longer-length fine-tuning provides an additional gain.

5 Conclusion

DCoLT treats reverse-diffusion intermediate steps as latent thinking actions and optimizes the full trajectory using final-outcome rewards. Implemented on SEDD and LLaDA, it outperforms SFT- and RL-based alternatives across math and code-generation benchmarks.

  • DCoLT optimizes reverse-diffusion intermediate steps as latent thinking actions using reinforcement learning rewarded by final-answer correctness.
  • The framework is implemented on two representative diffusion language models, SEDD and LLaDA.
  • DCoLT outperforms other SFT- or RL-based training algorithms across math and code-generation tasks and benchmarks.
  • DCoLT remains limited by training data and compute, leaving performance on math and code-generation tasks with substantial room for improvement.

A Analysis of Thinking Process in SEDD+DCoLT

On Sudoku, SEDD+DCoLT learns a progressive generation pattern that prioritizes easy cells before hard cells and refines initially incorrect predictions as context accumulates. On GSM8K-Aug, it generates with minimal positional bias rather than a fixed left-to-right order.

  • Sudoku: SEDD+DCoLT generates easy Sudoku cells earlier than hard cells, unlike SEDD+SFT, which shows little difference across difficulty levels.This indicates a progressive generation strategy from easy to hard.
  • Sudoku: During Sudoku denoising, initially incorrect predictions can be corrected after newly unmasked cells provide additional context.
  • GSM8K-Aug: On GSM8K-Aug, SEDD+DCoLT shows minimal positional bias, contrasting with GPT2+CoT’s consistent left-to-right generation order.

B Analysis of Thinking Process in LLaDOU (LLaDA+DCoLT)

LLaDOU’s diffusion generation reveals a non-linear reasoning process in which key numbers and symbols emerge early, while surrounding language is completed later. Its unmasking scores can favor correcting a token before it is unmasked.

  • Generation patterns: LLaDOU tends to unmask key numbers and symbols early, then fills surrounding textual elements to complete grammatical and fluent responses.
  • Generation patterns: Intermediate 256-token generations show key numbers and symbols gradually forming the thinking structure before grammatically correct sentences are completed.
  • Unmasking policy: The unmasking score estimates whether a predicted token is likely correct, allowing a lower-scored candidate to be corrected before later unmasking.In the illustrated example, “3” is corrected to “0” before unmasking.

C.2.1 Implementation Details for Math Problems

Math evaluation uses GSM8K and MATH test sets with hard rewards based on exact final-answer correctness. Training uses grouped completions for advantage calculation and runs on 16 H800 GPUs.

  • GSM8K and MATH evaluation follows common train-test splits and reports test-set accuracy.The datasets contain 7.5K training questions each, with 1.32K GSM8K and 5K MATH test questions.
  • A completion receives reward 1 only when its boxed final answer is symbolically equivalent to the ground truth; otherwise, reward is 0.
  • Training uses 64 prompts per batch and 16 completions per prompt for group-based advantage calculation, without a default KL penalty.
  • The full training runs for 140 iterations on 16 H800 GPUs, requiring about 63 GPU days.

C.2.2 Implementation Details for Code Generation

The code-generation setup filters a synthetic dataset into function-completion prompts with unit tests, then evaluates DCoLT models zero-shot on MBPP and HumanEval. Comparisons include matched LoRA settings and a Dream-7B generality experiment.

  • Training data: The training data comes from the fully synthetic KodCode-V1-SFT-R1 dataset, filtered to retain instruct-style samples with one testable function and a non-empty docstring.Additional filtering excludes class definitions and constructor methods.
  • Prompt format: 48.9K filtered samples, totaling around 12M tokens, are formatted as function-completion prompts containing target descriptions, unit tests, and code stubs.The prompt template includes the function name, parameters, and test assertions.
  • Reward: Code rewards are binary: a completion receives 1 only when its extracted code block passes all provided test cases.The code block is extracted from the response using markdown formatting.
  • Evaluation: Evaluation uses zero-shot MBPP and HumanEval with pass@1 as the reported metric.These benchmarks are used to assess code generation after post-training.
  • Matched comparison: 84.7% accuracy is achieved by LoRA-based DCoLT under matched GSM8K settings, outperforming d1 with diffu-GRPO at 79.8%.The comparison uses the same LoRA structure and dataset, with generation length 256, 128 diffusion steps, and block length 32.
  • Dream-7B generality: DCoLT raises Dream-7B accuracy on GSM8K from 50.11% to 80.53%, an absolute gain of +30.42%.The experiment applies the same reinforcement-learning configuration used in the main experiments.

C.3 Post-Training Data Used in Tab. 2

The post-training data comparison emphasizes that several competing methods use substantial proprietary math or code data, which contributes to their reported accuracies.

  • Data comparison: Several methods use considerable proprietary math or code data for pretraining and post-training, contributing crucially to their high accuracies.The passage identifies proprietary data as an important difference among methods.
  • Data comparison: Table 9 lists the post-training data used for math and code generation by the models compared in Table 2.A dash indicates that a post-training type was not used.

D Potential Scaling Law for Longer Generations

Longer generation lengths give diffusion models more room for extended reasoning, with benefits concentrated on difficult problems and sometimes allowing mistakes to be corrected later in the trajectory.

  • Generation length: Generation length in diffusion models increases with the number of mask tokens initialized before reverse diffusion, unlike autoregressive length determined by the EoT token.The EoT token can appear later when more mask tokens are initialized.
  • Scaling behavior: Longer generations can improve performance, particularly when the model is fine-tuned by DCoLT with increasing generation length.The authors suggest this may indicate a scaling law for longer generations.
  • Example: The illustrated solution derives x^4 - 7x^2 + 10 = 0, substitutes u = x^2, obtains u = 5 or 2, and concludes m - n = 3.The passage is an example of the reasoning content shown at different generation lengths.
  • Reasoning process: DCoLT examples show that some mistakes in the thinking process can be corrected in longer generations, producing better results than shorter ones.The paper presents this as a phenomenon for future study.

E Ablation on Various Block Lengths

Block length strongly affects baseline LLaDA, but DCoLT-trained LLaDOU is comparatively insensitive to the blocking choice, with only a small advantage for block length 8.

  • Experimental setup: Table 10 reports GSM8K accuracies for LLaDOU trained with generation length 256 and 64 reverse-diffusion steps across block lengths.The table organizes the ablation around alternative block-length settings.
  • Baseline LLaDA: 63.26% accuracy is reached by baseline LLaDA with block length 8 on GSM8K, versus 49.47% without dividing blocks.LLaDA generates responses blockwise from left to right, introducing a semi-autoregressive prior.
  • LLaDOU ablation: 83.17% accuracy for LLaDOU with block length 8 is only slightly above 82.27% with block length 256 and no blocking.These results suggest DCoLT-trained LLaDOU does not rely on the baseline blocking prior.
Loading 2505.10446v3…