Source-linked AI summary
Diffusion of Thoughts: Chain-of-Thought Reasoning in Diffusion Language Models
Jiacheng Ye, Shansan Gong, Liheng Chen, Lin Zheng, Jiahui Gao, Han Shi, Chuan Wu, Xin Jiang, Zhenguo Li, Wei Bi, Lingpeng Kong
TL;DR
CoT can support complex reasoning in autoregressive models, but intermediate-step errors, self-correction difficulties, and inefficiency remain concerns; whether diffusion language models can use CoT-style reasoning is an open question. DoT diffuses latent reasoning steps over time and adds techniques for self-correction and inference control. The paper reports advantages on multiplication, boolean logic, and grade-school math, including performance against larger autoregressive models.
Problem
The paper asks whether diffusion language models can use CoT-style reasoning to gain complex reasoning abilities, against a background of intermediate-step errors, self-correction difficulties, and inefficiency in CoT.
Source: §1 · p. 1Method
DoT progressively updates latent variables representing thoughts so reasoning steps diffuse over time, with multi-pass generation and training-time sampling for error recovery.
Source: §1 · p. 2Results
Across multiplication, boolean logic, and grade-school math, DoT is reported to outperform autoregressive baselines, including a 4.6× larger model on grade-school math problems.
Source: §1 · p. 2Takeaways & Limitations
The study presents an initial exploration of reasoning with diffusion language models and emphasizes their potential for complex reasoning beyond the autoregressive paradigm.
Source: §1 · p. 2, §6 · p. 10The authors note that current pre-trained diffusion models face challenges competing with autoregressive LLMs hundreds of times larger.
Source: §1 · p. 2
Abstract
from arXiv · showhide
Recently, diffusion models have garnered significant interest in the field of text processing due to their many potential advantages compared to conventional autoregressive models. In this work, we propose Diffusion-of-Thought (DoT), a novel approach that integrates diffusion models with Chain-of-Thought, a well-established technique for improving the reasoning ability of autoregressive language models. In contrast to autoregressive language models that make decisions in a left-to-right, token-by-token manner, DoT allows reasoning steps to diffuse over time through a diffusion language model and offers greater flexibility in trading-off computation for reasoning performance. Our experimental results demonstrate the effectiveness of DoT in multi-digit multiplication, boolean logic, and grade school math problems, with a small diffusion model outperforming a much larger autoregressive model in both efficiency and accuracy. In addition to that, DoT showcases promising self-correction abilities and benefits from existing reasoning-enhancing techniques like self-consistency decoding. Our findings contribute to the understanding and development of reasoning with diffusion language models.
1 Introduction
The introduction asks whether diffusion language models can use Chain-of-Thought reasoning to improve complex reasoning. It presents DoT as a method for diffusing reasoning over time, with self-correction and adjustable inference computation as central advantages.
1 Introduction
- The introduction frames CoT as vulnerable to errors in intermediate steps, self-correction difficulties, and inefficiency, while diffusion models offer potential benefits including global planning and efficiency.It asks whether diffusion language models can also gain complex reasoning ability from CoT-style techniques.Source: §1 · p. 1
- DoT combines diffusion models with Chain-of-Thought, progressively updating latent thoughts so reasoning steps diffuse in parallel over time.The introduction also describes a multi-pass variant that generates one thought at a time and uses classifier-free guidance for exact-token conditioning.Source: §1 · p. 2
- DoT integrates training-time sampling to learn recovery from errors in prior or current reasoning steps, and adapts a conditional ODE solver to accelerate continuous-model inference.The authors describe reasoning time and performance as a flexible trade-off.Source: §1 · p. 2
- DoT exhibits superior performance to GPT2 with CoT on grade-school math problems, enabling a small diffusion model to outperform an autoregressive model 4.6x larger.Source: §1 · p. 2
- The authors describe their study as preliminary and note that current diffusion models remain challenging to compare with autoregressive LLMs hundreds of times larger.They emphasize the potential for complex reasoning beyond the autoregressive paradigm.Source: §1 · p. 2
2 Preliminaries
The preliminaries introduce forward and reverse diffusion for text generation and distinguish continuous from discrete diffusion representations.
2 Preliminaries
- Forward diffusion progressively adds noise to a data representation, while reverse diffusion aims to recover the original representation by denoising.The text defines t = 1, 2, ..., T, with larger t corresponding to noisier data.Source: §2 · p. 3
- Continuous text diffusion maps discrete text into continuous space and applies Gaussian noise; discrete diffusion represents tokens as discrete random variables and uses transition matrices.The passage gives Diffusion-LM [33] and D3PM [1] as examples of the two categories.Source: §2 · p. 3
- For conditional sequence-to-sequence generation, DiffuSeq [15] concatenates source and target sequences and applies partial noising to target entries while keeping source entries uncorrupted.A left-aligned mask distinguishes source from target during diffusion.Source: §2 · p. 3
3 Diffusion-of-Thoughts
The method represents reasoning as latent thoughts distributed across diffusion time, with an optional multi-pass variant that generates rationales sequentially. Training and inference adapt diffusion procedures to improve error recovery, conditional generation, and control over computation.
3.1 Overview
- DoT represents the problem statement and answer as s and a, with intermediate rationales r1...n; its latent rationale representations zt correspond to noised diffusion data.Unlike implicit CoT's distribution across transformer layers, DoT distributes rationale representations over diffusion timesteps.Source: §3.1 · p. 3
3 Diffusion-of-Thoughts
- DoT progressively updates latent variables representing thoughts, distributing rationale representations over diffusion timesteps rather than transformer layers.The method section introduces DoT, its multi-pass variant, and the training and inference protocols.Source: §3.1 · p. 3, §3 · p. 3
3.2 Modeling
- DoT conditions on exact query tokens by concatenating the problem context with rationales and noising only the rationale portion; DoTMP generates one rationale at a time, conditioning later rationales on earlier ones.The multi-pass variant disentangles rationales and introduces a causal inductive bias.Source: §3.2 · pp. 3–4, Figure 2 · p. 4
3.3 Training
- Scheduled sampling exposes the denoiser during training to its own predicted errors, while coupled sampling adds noise to prior thoughts so later predictions can handle errors in them.The shared training objective minimizes a negative variational lower bound; its rounding loss regularizes embedding learning and its diffusion loss sums timestep KL divergences.Source: §3.3 · pp. 4–5
3.4 Inference
- Inference can trade reasoning time for computation by setting a larger backward timestep T; conditional ODE solving accelerates continuous diffusion, and self-consistency aggregates answers by majority vote across sampled pathways.Self-consistency marginalizes over rationale paths by selecting the most-voted answer.Source: §3.4 · p. 5
4 Experiments
Across arithmetic, logic, and GSM8K, DoT combines exact-answer performance with adjustable inference effort. It reaches 100% accuracy with up to 27× speedup on simpler tasks, while GSM8K performance improves as reasoning steps increase.
4 Experiments
- The experiments test diffusion-based reasoning on simple arithmetic and logic tasks as well as more complex grade-school math problems.The paper frames these evaluations as an exploration of the reasoning paradigm in diffusion models.Source: §4 · p. 5
4.1 Experimental Setup
- Experiments cover multiplication, boolean logic, and GSM8K, comparing exact-match answer accuracy and batch-size-one inference throughput across autoregressive and diffusion baselines.The setup includes GPT-2 scales, Plaid, SEDD-small and SEDD-medium, plus a five-shot ChatGPT baseline.Source: §4.1 · p. 5
4.2 Results on Digit Multiplication and Boolean Logic
- On multiplication and boolean logic, DoT reaches 100% accuracy with sampling steps set at 1 or 2, achieving up to 27× speedup over GPT-2.Neither ChatGPT nor distilled Implicit CoT reaches 100%, while fine-tuned GPT-2 sacrifices throughput during CoT.Source: §4.2 · p. 6
4.3 Results on Grade School Math
- DoT based on SEDD-medium outperforms similar-sized GPT2-medium with CoT by around 10% on GSM8K.Training from scratch reaches only around 5%; the results also report that multi-pass DoT performs slightly better than single-pass DoT on Plaid, while single-pass is more efficient.Source: §4.3 · p. 7
- Ablations find gradient-based conditioning unsuitable for downstream fine-tuning; training-time sampling aids self-correction, and conditional ODE speeds Plaid inference with decent performance at 8 generation timesteps.The authors attribute query-token recovery differences to soft gradient guidance and report that the two proposed sampling strategies are effective.Source: §4.3 · p. 7, Figure 3 · p. 7
4.4 Reasonability-efficiency Trade-off
- DoT-Plaid and DoT-SEDD reach 100% accuracy on simpler 5×5 tasks in 1–2 steps, while DoT-SEDD-medium beats CoT-GPT2-medium at 32 timesteps on GSM8K and keeps improving with more.The paper contrasts this with autoregressive CoT and Implicit CoT, which are hard to make more efficient because of token-by-token prediction.Source: §4.4 · p. 8
4.5 Self-consistency in DoT
- Self-consistency improves DoT and DoTMP on Plaid and also significantly improves SEDD DoT.The paper attributes this to diverse diffusion-generated reasoning paths, which provide cross-validation when selecting the most consistent answer.Source: §4.5 · p. 8
4.6 Self-correction in DoT
- Examples show DoT can establish a solution framework early, revise an incorrect intermediate thought using later thoughts, and maintain predictions after an answer is complete.The cases span questions of increasing difficulty and illustrate reasoning that does not follow a left-to-right sequence.Source: §4.6 · pp. 8–9, Table 3 · p. 9
5 Related Work
Diffusion models have been used across text-generation and sequence-to-sequence tasks, while pre-training and fine-tuning for diffusion language models remains underexplored. This work places DoT within efforts to address CoT’s inefficiency and difficulty with self-correction.
5.1 Diffusion Models for Text
5.2 Pre-train and fine-tune Diffusion LMs
- Prior diffusion-LM work includes masked-LM initialization [20] [59] [66], GENIE paragraph denoising [35], and from-scratch Plaid and SEDD [18]; this paper is the first to explore fine-tuning a pre-trained diffusion LM for reasoning.The passage describes pre-training and fine-tuning as relatively under-explored for diffusion language models.Source: §5.2 · p. 9
5.3 Reasoning Paradigms
- CoT supports deliberate reasoning, but errors in intermediate steps can yield inaccurate answers and CoT can be inefficient; diffusion is explored as an alternative.Self-consistency [52] samples reasoning paths and selects the most consistent answer, while tree-of-thought [57] uses tree search.Source: §5.3 · pp. 9–10
6 Conclusion and Limitation
The paper presents DoT, which combines CoT reasoning with continuous diffusion models, and evaluates it across mathematical reasoning tasks. It identifies additional training for accurate reasoning as a limitation and outlines directions for broader models and tasks.
6 Conclusion and Limitation
- DoT integrates CoT reasoning with continuous diffusion models and is evaluated on mathematical tasks for efficiency control, self-correction, and reasoning-path diversity.The paper frames this as an initial exploration of current diffusion language models, which remain early in development, particularly in model scale relative to autoregressive models.Source: §6 · p. 10
- Accurate reasoning with DoT requires additional training, a stated limitation of the method.Source: §6 · p. 10
- The authors anticipate that more powerful pre-trained diffusion models may improve generalization and remove the need for specialized training; other Transformer variants may improve inference efficiency.Source: §6 · p. 10
- The diffusion training techniques are described as applicable beyond mathematical reasoning, with multi-task instruction tuning and other modalities proposed as future extensions.Source: §6 · p. 10
A Derivations
The appendix derives training and sampling components for continuous and discrete diffusion models used in sequence-to-sequence generation and DoT inference. It covers latent construction, forward and reverse processes, model-specific objectives, and an adapted conditional ODE solver.
A.1 Seq2Seq Modeling in DiffuSeq
- In DiffuSeq, source and target text are embedded jointly, but forward noising perturbs the target while preserving the source; reverse denoising aims to recover the original latent sequence.The training objective includes denoising and embedding-related terms, with embedding regularization in the simplified objective.Source: §A.1 · pp. 15–16
A.2 Pre-trained Plaid
- For pre-trained Plaid, the VDM formulation is adapted to sequence-to-sequence modeling with partial noise and an anchored source; its continuous-time diffusion loss uses the standard formulation rather than the reweighted objective.The Plaid variance schedule is parameterized by a scalar-to-scalar neural network.Source: §A.2 · pp. 16–17
A.3 Pre-trained SEDD
- SEDD learns a concrete score rather than directly modeling the data distribution, sparsifying probability ratios to nearby states to avoid an intractable set of all ratios.Its forward process uses categorical noise, and its reverse process derives token transitions from a reverse transition-rate matrix.Source: §A.3 · pp. 18–19
- SEDD’s score-entropy objective is extended with denoising score entropy to handle an unknown data-distribution ratio.Source: §A.3 · p. 18
A.4 Conditional ODE solver
- The conditional ODE solver adapts continuous-diffusion sampling to partial-noising conditions and approximates the solution at each step using the previous latent state and denoiser derivatives.The solution is expressed using a change of variable for logSNR, with higher-order derivatives estimated through a Taylor expansion.Source: §A.4 · pp. 19–20
B.1 Dataset Statistics
The dataset statistics passage identifies the data sources for digit multiplication, GSM8K, and Boolean logic, and reports their test-set sizes.
B.1 Dataset Statistics
- Digit multiplication and GSM8K use processed datasets from Implicit CoT [7].Source: §B.1 · p. 20
- The Boolean logic training and test datasets are constructed using the method from DyVal.Source: §B.1 · p. 20
- Each dataset has 1,000 test examples except GSM8K, which has 1,319.Source: §B.1 · p. 20
B.2 Details of Baselines
The baseline setup specifies GPT-2 training and decoding, a 5-shot ChatGPT prompt, and dataset statistics; ChatGPT throughput measures response speed rather than model generation speed.
B.2 Details of Baselines
- GPT-2 was fine-tuned for 40 epochs, with learning rates of 1e-4 for boolean logic and 5e-4 for other tasks; inference used greedy decoding or self-consistency sampling.Self-consistency sampling used temperature 0.5 and top-k 40.Source: §B.2 · p. 20
- Table 4 reports training-set size, validation-set average input, intermediate, and output token counts, and average rationale counts using the Plaid tokenizer.The supplied passage describes the table’s contents but does not provide its values.Source: Table 4 · p. 21
- Plaid 1B was compared with fine-tuned GPT-2 small because their perplexities were similar, though Plaid’s nearly 10× parameter count could disadvantage its inference speed.The comparison uses Plaid 1B [18].Source: §B.2 · p. 21
- The Transformer-scratch baseline used six encoder and six decoder layers, a BERT-base-uncased tokenizer, learning rate 1e-5, and 60k training steps.Training used batch size 128.Source: §B.2 · p. 21
- ChatGPT was prompted with 5-shot examples, and its reported throughput measures response speed, not the model’s actual generation speed.As a black-box product, it may use optimization techniques to accelerate responses.Source: §B.2 · p. 21, Table 5 · p. 21
B.3 DoT Implementation Details
Implementation details cover the hardware and datasets, from-scratch training and sampling settings, and fine-tuning configurations for Plaid and SEDD.
B.3 DoT Implementation Details
- Experiments used eight NVIDIA V100-32G GPUs, with fp16; from-scratch DoT was trained on multiplication and GSM8K, and Plaid-1B was also fine-tuned on GSM8K.The V100 hardware does not support bf16.Source: §B.3 · p. 21
- From-scratch DoT used a 12-layer Transformer, learning rate 1e-4, 60k steps, batch size 128, and maximum sequence length 128; sampling used T = 1 for multiplication and T = 2 for boolean logic.The four-digit and five-digit multiplication data were preprocessed for the multi-path variant.Source: §B.3 · pp. 21–22
- Plaid fine-tuning used 120k DoT steps and 30k multi-pass DoT steps, while SEDD used 200k steps for both; default inference used T = 64.For Plaid, the authors report that more training steps degraded performance; all experiments were run multiple times to verify statistical significance.Source: §B.3 · p. 22
B.4 Additional Results
Additional results compare DoT with answer-only fine-tuning, examine throughput and reasoning paths, and report further attempts, base-model constraints, and broader impacts.
B.4 Additional Results
- Answer-only fine-tuning of diffusion models performed worse than DoT, paralleling the degradation of autoregressive models without CoT.Table 7 presents the DoT versus answer-only comparison.Source: §B.4 · p. 22, Table 7 · p. 22
- Throughput for Plaid DoTMP appears to vary nearly linearly with the number of timesteps T.The comparison is presented in Table 8.Source: §B.4 · p. 22, Table 8 · p. 23
- DoTMP produces more correct reasoning paths than DoT, while DoT slightly exceeds DoTMP in reasoning-path diversity.The passage points to Figure 4(b) and gives examples of DoTMP paths succeeding where DoT fails.Source: §B.4 · pp. 22–23
B.5 Other Attempts
- Replacing Plaid 1B with GPT-2 large in diffusion-style fine-tuning did not produce the correct answer format, suggesting the need for a pretrained diffusion model.The altered setup removed the causal mask while keeping the other diffusion training settings.Source: §B.5 · p. 23
- Mixing four-digit and five-digit multiplication data allowed the model to learn when to stop computation and attain 100% accuracy.The authors mixed the datasets because their rationales differ in number.Source: §B.5 · p. 23
C Discussion about base models
- DoT is constrained to pre-training and fine-tuning because its base models are not strong enough; the study is preliminary amid a lack of scaled pretrained diffusion models.The authors describe stronger pretrained models and post-instruction tuning as future directions for generalization.Source: C Discussion about base models · p. 23
D Boarder Impacts
- DoT achieves competitive performance against autoregressive CoT, but remains far behind modern large autoregressive models such as ChatGPT.The authors also note that generative models may facilitate misinformation and perpetuate biases reflected in training data.Source: D Boarder Impacts · pp. 23–24