Can reward-guided stitching improve test-time scaling with diffusion language models?

Reward-guided stitching reuses high-scoring steps from many cheap diffusion traces, lifting accuracy up to 23.8% with lower latency.

Direct answer

Yes, within limits. The paper's training-free pipeline samples many cheap reasoning traces with a masked diffusion language model, scores each intermediate step with an off-the-shelf process reward model, stitches the highest-scoring steps across traces into one composite rationale, and lets an autoregressive solver recompute only the final answer. The authors report average accuracy gains of up to 23.8% across six math and coding tasks and up to 1.8x lower latency than diffusion and unified baselines, with the largest benefit on harder problems.

1sources cited

This article was generated with WisPaper-powered search and paper analysis.

Trajectory-level aggregation throws away partial work

Most ways of combining multiple reasoning attempts work at the level of the whole trace: pick the single best chain, or vote over final answers. The paper argues that this discards useful intermediate work from attempts that are partly right but fail later. A single incorrect step can derail an otherwise promising derivation, so spending all compute on one long trajectory is risky.

The paper's alternative is to treat diffusion sampling as a source of reusable step-level candidates rather than a set of complete solutions. The claim is that cheap diffusion-sampled chains often contain locally correct sub-results, and that identifying and recombining those steps can recover much of the benefit of a broad search without paying for one long derivation from scratch.

Three stages separate exploration, scoring, and synthesis

The pipeline has three stages. First, a masked diffusion language model generates many diverse, low-cost reasoning trajectories for a problem, and each trajectory is split into steps. Second, an off-the-shelf process reward model (PRM) assigns every intermediate step a score in [0,1]. Third, steps at or above a confidence threshold are collected across traces and sorted chronologically, with ties broken by higher score; the full best-scoring trace is also kept as an anchor. This stitched evidence list is passed to a generative solver, which recomputes only the final answer.

The design deliberately keeps the components modular: diffusion handles exploration, the PRM handles evaluation, and an autoregressive model handles solution synthesis. The authors contrast this with unified hybrid architectures that couple exploration, scoring, and final generation, which they say limits flexibility and makes targeted improvements harder. Because the traces are independent, they can be dispatched in parallel across devices, which is where the latency argument comes from.

Stitching beats voting and best-trace selection

The ablation on GSM8K and MATH500 shows a clear progression. A plain LLaDA baseline reaches 78.8 and 37.6. Majority voting over final answers reaches 85.1 and 42.0. Adding a solver but no stitching, conditioning on all sampled chains reaches 86.6 and 46.0, and conditioning on the single best chain reaches 90.1 and 49.2. Confidence-filtered stitching reaches 90.4 and 52.0, and stitching plus the best chain as an anchor reaches 91.5 and 54.2.

Two things stand out. The solver matters: recomputation converts stitched but imperfect rationales into correct answers, and the authors call this out as a key ablation finding. And selection matters: keeping only above-threshold steps beats feeding the solver everything, so the PRM's filtering is doing real work rather than just adding context.

The final solver does more than a diffusion decoder

A separate comparison holds the stitched evidence fixed and varies who writes the final answer. With four traces on GSM8K, an autoregressive solver reaches 91.5 accuracy using 96.9 steps, while a diffusion solver at confidence 0.7 reaches 90.1 using 106.3 steps and at 0.9 reaches 90.3 using 141.6 steps. On MATH500 the gap widens: 54.2 for the autoregressive solver versus 50.2 and 52.0 for the diffusion solver.

The pattern repeats with eight traces, where the autoregressive solver reaches 91.4 on GSM8K and 55.0 on MATH500. The authors read this as evidence that the final solver is what converts stitched but imperfect rationales into accurate answers, and that diffusion decoding alone, even at higher confidence and more steps, slightly underperforms it.

Gains hold across diffusion backbones

The framework is model-agnostic in the sense that it only assumes a generator that can produce candidate trajectories and a PRM that can score steps. Applying it to other diffusion backbones produces consistent gains. LLaDA-1.5 rises from 38.0 to 53.8 on MATH500 while inference steps fall from 163 to 134, because stitching allows a lower confidence threshold. Dream rises from 82.6 to 85.4 on GSM8K and from 48.2 to 53.2 on MATH500 with essentially unchanged step counts. LLaDA 2.0 rises from 44.6 to 61.4 on MATH500 with compute roughly comparable at 256 to 267 steps.

The headline aggregate is an average accuracy improvement of up to 23.8% across six math and coding tasks, alongside up to 1.8x latency reduction relative to both traditional diffusion models such as Dream and LLaDA and unified architectures such as TiDAR. In the 8B setting the paper reports 72.38 versus 48.51 average points against an LLaDA baseline decoded at a high confidence threshold of 0.9, despite using a much lower threshold of 0.7.

Harder problems benefit most from recombination

The paper reports that step-level recombination is most beneficial on harder problems. That is consistent with the mechanism: on easy problems a single trace is often already correct, so there is little partial work to salvage, while on harder problems traces diverge and complementary correct steps become more likely to exist across attempts.

The failure analysis sharpens this. Among 50 incorrect MATH500 predictions, adding a single missing intermediate step and rerunning the solver repaired only 1 case attributable to a PRM false negative and 7 attributable to a step that was never generated; 40 of 50 were not repaired by one missing step. The authors conclude that incomplete diffusion exploration, rather than isolated PRM errors, is the dominant failure mode. Separately, among 26 examples whose stitched evidence contained at least one incorrectly selected step, the solver still recovered the correct answer in 8 cases, or 30.8%.

Verifier quality sets the ceiling on stitched evidence

The PRM is not a neutral filter. On 2,016 candidate steps from 60 MATH500 examples, labelled by three blinded judge passes, Qwen-Math-PRM-7B reached 88.97% precision, 89.23% recall, and 89.10% F1, while Skywork-1.5B reached 82.81%, 78.17%, and 80.42%. Among steps selected exclusively by each verifier, 150 of 250 Skywork-only selections were incorrect, compared with 79 of 331 Qwen-only selections.

The documented failure modes follow from this. Stitching fails when sampled trajectories share the same early error, so the pool contains no reliable evidence; when the PRM scores a crucial correct step too low and filters it out; or when the PRM is over-confident in a fluent but incorrect late step that then dominates the stitched evidence. The authors also note that overall performance still depends on the quality of the underlying diffusion backbone and reward model.

Latency gains come with memory costs

The efficiency claim is about forward evaluations and wall-clock latency, not about total resource use. The paper measures inference steps as the number of model forward evaluations needed to answer a problem, counting the diffusion model, the PRM, and the autoregressive solver. On that measure, stitching sits on a better accuracy-versus-steps frontier than diffusion-only baselines, which are efficient but less accurate, and autoregressive baselines, which are accurate but need many sequential decoding steps.

The trade-off is explicit in the limitations: the speedup comes from parallel diffusion sampling, which increases memory usage and hardware requirements as the number of simultaneous trajectories grows. The paper positions the method as practical for latency-constrained reasoning when parallel inference is available, not as a way to reduce total compute.

When stitching is the right choice

The decision rule that emerges from the evidence is about the scenario, not the average. If latency matters and parallel devices are available, and the task is hard enough that individual diffusion traces diverge and fail in different places, stitching plus a single autoregressive solve is the configuration the paper supports: it beat majority voting, all-chain conditioning, and best-chain conditioning on both GSM8K and MATH500, and it improved every diffusion backbone tested.

If the problem is easy enough that one trace usually succeeds, or if memory is the binding constraint rather than latency, the reported gains do not establish an advantage. And because 40 of 50 analyzed failures were not repaired by adding one missing step, the more useful lever is broadening exploration or improving the PRM, not tuning the stitching threshold alone.

Training-free, but not backbone-free

The framework requires no fine-tuning, which the paper contrasts with d1, an RL fine-tuning approach that optimizes a model directly for specific evaluation tasks. Across generation lengths of 256 and 512, stitching outperformed d1-LLaDA on GSM8K (89.4 and 91.5 versus 81.1 and 82.1), MATH500 (49.4 and 54.2 versus 38.6 and 40.2), and Countdown (37.5 and 45.7 versus 32.0 and 42.2).

That comparison is about a generalist inference-time pipeline versus a task-specific fine-tuned specialist, and it does not mean the pipeline is free of dependencies. It still needs a diffusion generator that produces diverse traces, a PRM that ranks steps well, and an autoregressive solver, and the paper's own failure analysis shows the first two are where the remaining errors concentrate.

The priority is exploration, not threshold tuning

The most decision-relevant result in this paper is the failure attribution, not the headline accuracy number. Because 40 of 50 analyzed failures could not be repaired by adding a single missing step, and only 1 of 50 traced to a PRM false negative, the binding constraint on this pipeline is whether the diffusion sampler ever produces the needed sub-derivation. That reframes where effort should go: increasing trace diversity and coverage matters more than fine-tuning the stitching threshold or swapping in a marginally better verifier.

The second priority is the solver. Holding stitched evidence fixed, the autoregressive solver beat diffusion decoding on both benchmarks at both trace counts, which is why the paper treats the final recomputation step as essential rather than incidental. For readers deciding whether to adopt this approach, the paper's evidence supports it as a latency-oriented configuration for hard reasoning tasks with parallel hardware, with the caveat that its accuracy ceiling is set by exploration breadth and verifier quality rather than by the stitching mechanism itself.

Sources used in this answer

1

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

A training-free pipeline that samples many cheap diffusion reasoning traces, scores each step with a process reward model, stitches high-scoring steps into one rationale, and recomputes only the final answer improves average accuracy by up to 23.8% across six math and coding tasks with up to 1.8x lower latency, with the largest benefit on harder problems.