Test-Time Scaling with Diffusion Language Models via Reward-Guided Stitching

TL;DR

Stitching Noisy Diffusion Thoughts, a self-consistency framework that turns cheap diffusion-sampled reasoning into a reusable pool of step-level candidates and highlights the importance of the final AR solver in converting stitched but imperfect rationales into accurate answers.

Abstract

Reasoning with large language models often benefits from generating multiple chains-of-thought, but existing aggregation strategies are typically trajectory-level (e.g., selecting the best trace or voting on the final answer), discarding useful intermediate work from partial or "nearly correct" attempts. We propose Stitching Noisy Diffusion Thoughts, a self-consistency framework that turns cheap diffusion-sampled reasoning into a reusable pool of step-level candidates. Given a problem, we (i) sample many diverse, low-cost reasoning trajectories using a masked diffusion language model, (ii) score every intermediate step with an off-the-shelf process reward model (PRM), and (iii) stitch these highest-quality steps across trajectories into a composite rationale. This rationale then conditions an autoregressive (AR) model (solver) to recompute only the final answer. This modular pipeline separates exploration (diffusion) from evaluation and solution synthesis, avoiding monolithic unified hybrids while preserving broad search. Across math reasoning benchmarks, we find that step-level recombination is most beneficial on harder problems, and ablations highlight the importance of the final AR solver in converting stitched but imperfect rationales into accurate answers. Using low-confidence diffusion sampling with parallel, independent rollouts, our training-free framework improves average accuracy by up to 23.8% across six math and coding tasks. At the same time, it achieves up to a 1.8x latency reduction relative to both traditional diffusion models (e.g., Dream, LLaDA) and unified architectures (e.g., TiDAR). Code is available at https://github.com/roymiles/diffusion-stitching.

Figures (8)

• Figure 1: Accuracy vs. wall-clock latency on Math500 comparing autoregressive baselines ($\mathbin{}$) using early stopping, diffusion baselines ($\color{gray} \circ$), and our stitching pipeline ($\color{blue} \blacksquare$, $\color{blue} \bullet$) using 4 reasoning traces. All models are 7B in size, except for LLaDA 2.0. All diffusion models are run without KV caching, so enabling caching wu2025fastdllmtrainingfreeaccelerationdiffusion would likely provide additional practical speedups.
• Figure 2: Diffusion Stitching pipeline. We first use a diffusion model to efficiently explore diverse reasoning paths, then score each intermediate step with a PRM, and finally stitch the highest-quality steps into a single rationale that conditions an AR solver to recompute the final answer.
• Figure 3: Parallisable inference cost (gen length 512). We report the number of parallelisable diffusion steps and the number of AR solver decoding steps for math and coding tasks.
• Figure 4: Pareto fronts of forward passes vs. accuracy on MBPP, GSM8K, and MATH500, showing that diffusion stitching bridges the gap between diffusion and AR models while enabling efficient test-time scaling for low-latency reasoning. All models are of the same 7-8B size, except for DeepSeek Coder which is 14B. We report our result for generation length 128, 256, and 512.
• Figure 5: Qualitative examples of step stitching. Each column shows a stitched evidence trace built from multiple diffusion trajectories (colors denote retained high-scoring steps). A lightweight AR solver then conditions on this trace to output only the final answer. Column 1 illustrates resolving disagreements; column 3 shows using confidence as to guide the final prediction.