Reinforcing the Diffusion Chain of Lateral Thought with Diffusion Language Models

TL;DR

This work introduces Diffusion Chain of Lateral Thought (DCoLT), a framework that treats intermediate reverse-diffusion steps as latent thinking actions and trains them with final-outcome reinforcement learning to encourage nonlinear, bidirectional reasoning. It is instantiated in continuous-time (SEDD) and discrete-time (LLaDA) diffusion models, employing a concrete-score policy and a Plackett-Luce-based Unmasking Policy Module, respectively. Empirical results on math and code benchmarks show that DCoLT-enhanced diffusion models outperform traditional SFT/RL baselines and even compete with autoregressive models trained on substantially more data, illustrating improved reasoning capabilities with limited supervision. The findings suggest diffusion-based lateral thinking as a promising avenue for scalable, reasoning-focused AI, with practical impact in domains requiring verifiable, complex problem solving.

Abstract

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.

Figures (14)

• Figure 1: Comparison between CoT and DCoLT. (a) A typical CoT performs vertical thinking by following an auto-regressive convention that generates responses token by token from left to right in a linear way. (b) DCoLT performs lateral thinking that generates the responses in a non-linear way without following the auto-regressive order; moreover, at each step, it can generate multiple tokens at chosen positions. We focus on the lateral thinking in this paper by reinforcing the chain of such lateral thought as an entirety in Diffusion Language Models (DLMs).
• Figure 2: The structure of LLaDOU. It first predicts the token set to unmask $\mathcal{U}_n$ according to the ranking score $h_{\theta,n}$ by the UPM, and then samples those unmasked tokens in $\mathcal{U}_n$ by LLaDA blocks.
• Figure 3: This figure shows the model's inference patterns on Sudoku dataset, with the left side displaying the pattern of SEDD + SFT model and the right side showing that of our SEDD + DCoLT model. We plot the total number of generated tokens for these three cell types over diffusion steps on Sudoku $4\times4$ dataset, with the number of sampling steps set to 32.
• Figure 4: Visualization of predicted tokens $\hat{x}_0$ by SEDD + DCoLT on Sudoku $4\times4$: those that are still masked appear in light blue, with unmasked ones in dark blue, incorrect predictions in red, and corrected ones in green. Black borders indicate the given cells.
• Figure 5: Average generation step for each token position in GSM8K-Aug. Different from CoT, SEDD + DCoLT generates in a non-linear way.