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

Roy Miles; Aysim Toker; Andreea-Maria Oncescu; Jiankang Deng; Ismail Elezi

cs.CLFeb 26, 2026

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

Authors: Roy Miles, Aysim Toker, Andreea-Maria Oncescu, Jiankang Deng, Ismail Elezi

Organizations: Huawei London Research Center

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 is then used 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 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). The code is available https://github.com/roymiles/diffusion-stitching

Figures & tables

Figure 1: 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 a generative solver to recompute the final answer.

Figure 2: Accuracy vs. wall-clock latency on Math500 comparing autoregressive baselines (

\mathbin{\hbox to5.62pt{\vbox to4.94pt{\pgfpicture\makeatletter\hbox{\hskip 0.25pt\lower-0.25pt\hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke=#000000} \lxSVG@begingroup@{fill=#000000} \lxSVG@setlinewidth{\the\pgflinewidth}\lxSVG@begingroup@{stroke-width=0.4pt} \lx@inpgf@ignorespaces\nullfont\hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} {{\lx@inpgf@ignorespaces}}{{}{{}}{} {}{} {}{} {\lx@inpgf@ignorespaces}\lxSVG@begingroup@{_scopebegin=1} \color[rgb]{0.6914,0.8164,0.8789}\lxSVG@fill\lxSVG@drawpath@unclipped{M 0 0 L 7.09 0 L 3.54 6.14 Z}{stroke:none} \lx@inpgf@ignorespaces \lxSVG@closescope {}{{}}{} {}{} {}{} {\lx@inpgf@ignorespaces}\lxSVG@begingroup@{_scopebegin=1} \color[rgb]{0.5,0.5,0.5}\lxSVG@setlinewidth{\the\pgflinewidth}\lxSVG@begingroup@{stroke-width=0.5pt} \lx@inpgf@ignorespaces{}\lxSVG@stroke\lxSVG@drawpath@unclipped{M 0 0 L 7.09 0 L 3.54 6.14 Z}{fill:none} \lx@inpgf@ignorespaces \lxSVG@closescope } \lxSVG@closescope {\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}\hss}\lxSVG@discardpath\lxSVG@closescope \hss}}\lxSVG@closescope\endpgfpicture}}}

) using early stopping, diffusion baselines (

\color[rgb]{0.5,0.5,0.5}\circ

), and our stitching pipeline (

\color[rgb]{0.1211,0.4648,0.707}\blacksquare

\color[rgb]{0.1211,0.4648,0.707}\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 Wu et al. (2025) would likely provide additional practical speedups.

		Coding				Math		Avg
Model Arch	Size	HumanEval	HumanEval+	MBPP	MBPP+	GSM8K	Minerva Math
Qwen3	4B	57.32	50.61	67.00	80.69	77.48	47.10	63.37
Qwen3	8B	64.63	56.71	69.40	83.07	81.80	52.94	68.09
LLaDA	8B	32.32	27.44	40.80	51.85	70.96	27.30	41.78
$\Lsh$ w/ conf sampling	8B	40.24	37.80	41.89	55.82	78.77	36.54	48.51
Dream	7B	54.88	49.39	56.80	74.60	77.18	39.60	58.74

Table 1: Evaluation Results: We evaluate the performance of Diffusion Stitching over several coding and math tasks. Baseline LLaDA and Dream numbers are taken from the TiDAR paper. The confidence-sampling variants were reproduced using

\gamma=0.9

and 512 denoising steps, respectively. For our pipeline ⋆ the 8B size refers to the diffusion backbone and all runs were evaluated in a zero-shot setting.

Panel of Figure 3 — Figure 3: 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.

Greedy
		GSM8K		MATH500
Setting	Len.	Acc.	Lat.	Acc.	Lat.	Mem.
		(%)	(s)	(%)	(s)	(GB)
Ours (Qwen PRM)	–	91.3	2.585	54.3	4.735	42.8 / 42.7
Ours (Skywork PRM)	–	89.4	2.421	51.8	3.910	32.3 / 32.2
Qwen2.5-Math-7B	256	32.0	2.545	18.0	4.989	14.4 / 14.4

Table 2: Accuracy–latency comparison against AR baselines. We compare against Qwen2.5-Math-7B-Instruct using greedy decoding, self-consistency, and PRM best-of-4 under the same hardware setup.

Figure 4: Qualitative examples of step stitching . Each column shows a stitched evidence trace built from multiple diffusion trajectories (colors denote retained high-scoring steps). A generative solver then conditions on this trace to output only the final answer. Column 1 illustrates resolving disagreements; column 3 shows confidence scores guide the final prediction.

Table 7

Table 8

Figure 5: 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. The results covering only diffusion generation (reasoning and answer) are provided in table 4 .

Model	Acc	#steps	Acc	#steps
	GSM8K		MATH500
Qwen-Math-Instruct	94.3	440	79.6	1504
Phi-4	92.4	295	76.8	791
LLaDA-1.5	84.0	101	38.0	163
$\Lsh$ w/ stitching	90.8	87	53.8	134
Dream	82.6	512	48.2	512

Table 7: dLLM/AR model ablation . AR models are accurate but require many decoding steps, whereas our method improves dLLM performance while reducing inference cost by enabling lower-confidence sampling.

Failure attribution
Failure source	Count	Percent
PRM false negative	1	2.0%
Diffusion exploration failure	7	14.0%
Stitching / solver failure	2	4.0%
Not repaired by one missing step	40	80.0%
Step-level verifier quality

Table 8: PRM reliability and failure analysis on Math500 . Top: attribution of 50 incorrect predictions after adding one candidate missing reasoning step and rerunning the solver. Bottom: step-level verifier quality on 2,016 candidate steps from 60 examples.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 6 — Figure 6: Ablation on sampling temperature . Accuracy (%) improves with higher temperature by promoting diverse reasoning paths. Although strong performance is achieved across a range of values, the optimal temperature is dataset-specific.

Step scoring
Setting	GSM8K	MATH500	MBPP
Geometric mean (ours)	91.5	54.2	74.6
Arithmetic mean	91.6	54.4	73.6
Anchor choice
Strong anchor (ours)	91.5	54.2	74.6
Weakest anchor	89.7	48.8	73.1

Appendix

Table 9: Robustness of stitching design choices. We ablate the trace-scoring function, anchor selection, step ordering, and segmentation strategy across math and coding benchmarks.

Method	Mode	Token limit	AIME25	Latency (s)
LLaDA2.0-mini	–	8192	33.33	91.43
Qwen3 solver-only	Non-thinking	4096	23.33	40.35
Qwen3 solver-only	Thinking	8192	16.70	148.75
Qwen3 solver-only	Thinking	16392	56.70	233.13
Ours (LLaDA2.0-mini + Qwen3)	Non-thinking	8192 / 256	46.67	112.54

Appendix

Table 10: Results on AIME25. For our method, the token budget is reported as diffusion generator / final solver. The Qwen3 solver in our pipeline is used in non-thinking mode to consolidate the stitched reasoning evidence and output the final answer.

Task Category	Task Name	PRM	Solver	Score Metric
Coding	HumanEval	Qwen2.5-Math-PRM-7B	QWEN2.5-CODER-7B	Pass@1
	HumanEval+	Qwen2.5-Math-PRM-7B	QWEN2.5-CODER-7B	Pass@1
	MBPP	ACECODERRM	QWEN2.5-CODER-7B	Pass@1
	MBPP+	ACECODERRM	QWEN2.5-CODER-7B	Pass@1
Math	GSM8K-CoT	Qwen2.5-Math-PRM-7B	Qwen2.5-Math-Instruct	strict-match
Math	MATH500	Qwen2.5-Math-PRM-7B	Qwen2.5-Math-Instruct	strict-match

Appendix

Table 11: Metrics used for each evaluation task . All tasks were evaluated zero-shot, using 4 generations, and with a generation length of 512. Strict match is implemented using the math-verify library 3 3 3 github.com/huggingface/Math-Verify and exact match uses string matching equality.

Figure 7: Qualitative examples . Left: resolving split decisions and right: over-ruling uniformly low-confidence predictions.

Figure 8: Representative failure cases for stitching . Each column shows the stitched evidence list and the final autoregressive solver prediction. Failures occur when the sampled trajectories contain no correct sub-derivation (left, third), when a key correct step is scored too low by the PRM and is omitted (second), or when the PRM is over-confident in an incorrect late step (or anchor trace), causing the evidence set to be dominated by a wrong conclusion (right).

	GSM8K (0-shot)		MATH500 (0-shot)		Countdown (0-shot)
Model / Seq Len	256	512	256	512	256	512
d1-LLaDA	81.1	82.1	38.6	40.2	32.0	42.2
Ours	89.4	91.5	49.4	54.2	37.5	45.7

Appendix

Table 12: Model performance on GSM8K, MATH500, and Countdown Benchmarks: All models are evaluated zero-shot with generation lengths of 256–512. Diffusion stitching is training-free, whereas d1 requires task-specific RL fine-tuning.

	GSM8K
Model	Acc	#steps
RSD [ 22 ]	94.6	325/75
Ours	91.5	80/6

Appendix

Table 13: Comparison to speculative decoding on GMS8K. RSD steps are reported as drafter steps / verifier steps . For ours, the steps are reported as diffusion steps / solver steps .

Explore similar work

Reasoning with Continuous Latent Diffusion
Sep 28, 2026Xiang ChengNumerical Reasoning in Language Models Code Generation
Duke University Department of Electrical and Computer Engineering
DreamReasoner-8B: Block-Size Curriculum Learning for Diffusion Reasoning Models
Jun 17, 2026Zirui Wu, Lin Zheng, Jiacheng Ye +5Blockwise Diffusion Language Models Efficient Language Model Reasoning
The University of Hong Kong · Peking University
Procedural Knowledge at Scale Improves Reasoning
Apr 1, 2026Di Wu, Devendra Singh Sachan, Wen-tau Yih +1LLM Reasoning Knowledge Augmentation for Language Models
University of California, Los Angeles · Meta FAIR

Sep 28, 2026J/K move · Enter open · S save

Reasoning with Continuous Latent Diffusion

Xiang Cheng

Duke University Department of Electrical and Computer Engineering

Numerical Reasoning in Language Models Code Generation Continuous Latent Reasoning Diffusion Language Models Continuous Diffusion Language Models

Open paper AskPDF arXiv

Continuous diffusion generates complete reasoning solutions through iterative refinement in latent space. We introduce the Continuous Embedding Diffusion Reasoner (CEDR), an ELF-based training and inference recipe. Our experiments show that accurate decoding alone does not ensure strong reasoning performance. We therefore learn compact representations from multiple layers of a strong autoregressive teacher. Their decomposition also enables asynchronous denoising at different rates. We show that prompt encodings need only preserve the information required for the correct text-conditional score, rather than exactly match teacher features, and use a staged curriculum to learn a compact prompt encoder that replaces the teacher Transformer at inference. We adapt DiffusionNFT to learned self-conditioning guidance and incorporate gold-solution endpoints to supplement sparse rewards. Our supervised models outperform reported results from recent continuous-diffusion baselines at comparable backbone scales on mathematical reasoning and HumanEval code generation. With a 638M-parameter denoising backbone and learned prompt conditioning, post-NFT CEDR-L achieves 63.74% pass@1 on GSM8K and 24.6% on MATH500 at 64 denoising steps, and 32.85% on HumanEval and 30.18% on HumanEval+ at 128 denoising steps. Code will be available at: https://github.com/chengxiang/CEDR.

Abstract

Explore similar work

Reasoning with Continuous Latent Diffusion

DreamReasoner-8B: Block-Size Curriculum Learning for Diffusion Reasoning Models

Procedural Knowledge at Scale Improves Reasoning