Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

DiffCoT: Diffusion-styled Chain-of-Thought Reasoning in LLMs

Shidong Cao, Hongzhan Lin, Yuxuan Gu, Ziyang Luo, Jing Ma

TL;DR

DiffCoT addresses exposure bias and error accumulation in multi-step chain-of-thought reasoning by reframing reasoning as a diffusion-styled denoising process. It introduces a diffusion sliding window with a causal diffusion noise schedule to jointly enable generation and retrospective revision of intermediate steps, while preserving token-level autoregression. Training uses Direct Preference Optimization on diffusion-based thought trajectories, and data are constructed via MCTS-guided, reward-ranked step-level candidates. Across GSM8K, SVAMP, and MATH on multiple backbones, DiffCoT demonstrates robust improvements over state-of-the-art preference-learning baselines, with stronger error-correction capabilities and more stable reasoning across tasks and model sizes. This unified diffusion-plus-AR framework offers a practical and scalable path to more reliable multi-step reasoning in large language models.

Abstract

Chain-of-Thought (CoT) reasoning improves multi-step mathematical problem solving in large language models but remains vulnerable to exposure bias and error accumulation, as early mistakes propagate irreversibly through autoregressive decoding. In this work, we propose DiffCoT, a diffusion-styled CoT framework that reformulates CoT reasoning as an iterative denoising process. DiffCoT integrates diffusion principles at the reasoning-step level via a sliding-window mechanism, enabling unified generation and retrospective correction of intermediate steps while preserving token-level autoregression. To maintain causal consistency, we further introduce a causal diffusion noise schedule that respects the temporal structure of reasoning chains. Extensive experiments on three multi-step CoT reasoning benchmarks across diverse model backbones demonstrate that DiffCoT consistently outperforms existing CoT preference optimization methods, yielding improved robustness and error-correction capability in CoT reasoning.

DiffCoT: Diffusion-styled Chain-of-Thought Reasoning in LLMs

TL;DR

DiffCoT addresses exposure bias and error accumulation in multi-step chain-of-thought reasoning by reframing reasoning as a diffusion-styled denoising process. It introduces a diffusion sliding window with a causal diffusion noise schedule to jointly enable generation and retrospective revision of intermediate steps, while preserving token-level autoregression. Training uses Direct Preference Optimization on diffusion-based thought trajectories, and data are constructed via MCTS-guided, reward-ranked step-level candidates. Across GSM8K, SVAMP, and MATH on multiple backbones, DiffCoT demonstrates robust improvements over state-of-the-art preference-learning baselines, with stronger error-correction capabilities and more stable reasoning across tasks and model sizes. This unified diffusion-plus-AR framework offers a practical and scalable path to more reliable multi-step reasoning in large language models.

Abstract

Chain-of-Thought (CoT) reasoning improves multi-step mathematical problem solving in large language models but remains vulnerable to exposure bias and error accumulation, as early mistakes propagate irreversibly through autoregressive decoding. In this work, we propose DiffCoT, a diffusion-styled CoT framework that reformulates CoT reasoning as an iterative denoising process. DiffCoT integrates diffusion principles at the reasoning-step level via a sliding-window mechanism, enabling unified generation and retrospective correction of intermediate steps while preserving token-level autoregression. To maintain causal consistency, we further introduce a causal diffusion noise schedule that respects the temporal structure of reasoning chains. Extensive experiments on three multi-step CoT reasoning benchmarks across diverse model backbones demonstrate that DiffCoT consistently outperforms existing CoT preference optimization methods, yielding improved robustness and error-correction capability in CoT reasoning.
Paper Structure (28 sections, 7 equations, 5 figures, 3 tables)

This paper contains 28 sections, 7 equations, 5 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Chain-of-Thought Reasoning
  4. Diffusion Models
  5. Methodology
  6. Diffusion-Styled Noising for CoT
  7. Reformulating CoT Reasoning as Denoising
  8. Causal Diffusion Noise
  9. Experiment
  10. Settings
  11. Models and Datasets
  12. Baselines
  13. Main Results
  14. Ablation Study
  15. Analysis
  16. ...and 13 more sections

Figures (5)

  • Figure 1: Comparison of our proposed DiffCoT with existing CoT reasoning approaches: (a) Existing step-by-step CoT Reasoning methods adopt teacher-forcing training, where each step depends on the ground-truth output of the previous one. At inference time, this assumption breaks, causing exposure bias and leading to error accumulation. (b)DiffCoT performs CoT reasoning along both the noise (diffusion) and temporal (autoregressive) dimensions, enabling iterative correction of prior mistakes and effectively mitigating exposure bias.
  • Figure 2: DiffCoT Framework and Training Data Construction: (a) Step-level forward noising: MCTS-based data generation defines step-level noise by reward-ranking multiple candidates, yielding states ranging from clean to corrupted. (b) Sliding-window denoising: a diffusion sliding window refines previously generated CoT steps while producing the next step in an autoregressive manner. (c) Causal diffusion noise: a step-dependent schedule assigns stronger noise to later steps to encode the causal order of the reasoning chain.
  • Figure 3: Example illustrating how DiffCoT modifies early-stage reasoning shift steps. The steps highlighted in blue represent the diffusion sliding window.
  • Figure 4: Correction success rate under stochastic prefix corruption, where noise is injected at the midpoint of the reasoning trajectory with probability $\omega$.
  • Figure 5: Representative dataset example with step-wise reasoning annotations.