Test-Time Scaling with Diffusion Language Models via Reward-Guided Stitching
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
| 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 |
| 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 |
| 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 |
| Greedy | ||||||
| Qwen2.5-Math-7B | 256 | 32.0 | 2.545 | 18.0 | 4.989 | 14.4 / 14.4 |
| GSM8K | MATH500 | |||
|---|---|---|---|---|
| Model | Acc | #steps | Acc | #steps |
| 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 |
| w/ stitching | 90.8 | 87 | 53.8 | 134 |
| Dream | 82.6 | 512 | 48.2 | 512 |
| 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 | |||
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
| Setting | GSM8K | MATH500 | MBPP |
|---|---|---|---|
| Step scoring | |||
| 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 |
| 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 |
| 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 |
| MATH500 | Qwen2.5-Math-PRM-7B | Qwen2.5-Math-Instruct | strict-match |
| 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 |
| GSM8K | ||
|---|---|---|
| Model | Acc | #steps |
| RSD [ 22 ] | 94.6 | 325/75 |
| Ours | 91.5 | 80/6 |