$S^3$: Stratified Scaling Search for Test-Time in Diffusion Language Models

Ahsan Bilal; Muhammad Ahmed Mohsin; Muhammad Umer; Asad Aali; Muhammad Usman Khanzada; Muhammad Usman Rafique; Zihao He; Emily Fox

cs.LGApr 7, 2026

$S^3$ : Stratified Scaling Search for Test-Time in Diffusion Language Models

Authors: Ahsan Bilal, Muhammad Ahmed Mohsin, Muhammad Umer, Asad Aali, Muhammad Usman Khanzada, Muhammad Usman Rafique, Zihao He, Emily Fox, +1 more

Organizations: University of Oklahoma · Stanford University · University of Göttingen · Zoox · Meta

Abstract

Test-time scaling investigates whether a fixed diffusion language model (DLM) can generate better outputs when given more inference compute, without additional training. However, naive best-of- $K$ sampling is fundamentally limited because it repeatedly draws from the same base diffusion distribution, whose high-probability regions are often misaligned with high-quality outputs. We propose $S^3$ (Stratified Scaling Search), a classical verifier-guided search method that improves generation by reallocating compute during the denoising process rather than only at the final output stage. At each denoising step, $S^3$ expands multiple candidate trajectories, evaluates them with a lightweight reference-free verifier, and selectively resamples promising candidates while preserving diversity within the search frontier. This procedure effectively approximates a reward-tilted sampling distribution that favors higher-quality outputs while remaining anchored to the model prior. Experiments with LLaDA-8B-Instruct on MATH-500, GSM8K, ARC-Challenge, and TruthfulQA demonstrate that $S^3$ consistently improves performance across benchmarks, achieving the largest gains on mathematical reasoning tasks while leaving the underlying model and decoding schedule unchanged. These results show that classical search over denoising trajectories provides a practical mechanism for test-time scaling in DLMs.

Figures & tables

Panel of Figure 1 — Figure 1: Density–quality mismatch and trajectory search under $S^{3}$ for MATH-500 . Panels (a)-(c) illustrate the distributional mismatch between $p_{0}$ , the verifier landscape $f(x)$ , and the reward-tilted target $\tilde{p}_{0}$ . Panels (d)-(f) show the corresponding inference behavior of direct sampling versus $S^{3}$ trajectory search and the resulting shift in verifier score distribution.

Figure 2:

S^{3}

N

particles are initialized at

t{=}T

[P1], then at each step expanded into

Nb

candidates, scored via one-step clean predictions

\hat{x}_{0}^{(i,j,t)}

[P2], and resampled to

N

particles via SSP [P3]. Final output is selected by majority voting at

t{=}0

Dataset	Baseline Diffusion	Best-of- $K$	$S^{3}$
GSM-8K	68.16	69.56	70.21
MATH-500	25.60	28.20	30.20
TruthfulQA	46.49	49.36	49.57
ARC-Challenge	76.11	79.30	77.86

Table 1: Accuracy (%) of standard diffusion decoding, Best-of-

K

(

K{=}8

), and

S^{3}

(

N{=}4

b{=}2

\mathcal{K}{=}64

) on LLaDA-8B-Instruct.

Panel of Figure 3 — Figure 3: Inference-time scaling with $S^{3}$ across datasets. Top row: accuracy vs. compute ( $\mathrm{NFE}=\text{steps}\times N\times b$ ) across multiple $(N,b)$ settings. Bottom row: mean top-1 token confidence over denoising progress. Curves shown for Baseline, $S^{3}$ ( $N{=}4$ , $b{=}2$ , $\lambda{=}1.0$ ), and BoK ( $K{=}8$ ).

Panel of Figure 4 — Figure 4: Accuracy (%) across block lengths $\mathcal{K}\in\{2,4,8,16,32,64\}$ on four benchmarks. Results use LLaDA-8B-Instruct with $N{=}4$ , $b{=}2$ , and $K{=}8$ . Orange bars denote $S^{3}$ , blue bars denote Baseline, and green bars denote BoK. An upward arrow ( $\uparrow$ ) indicates configurations where the leading method improves over the Baseline.

Method	Look-ahead	Tilting	Accuracy (%)
Method	Look-ahead	Tilting	GSM8K	MATH-500	ARC-C	TruthfulQA
Baseline diffusion	$\times$	$\times$	68.16	25.60	76.11	46.49
Best-of- $K$	$\times$	$\times$	69.56	28.20	79.30	49.36
Tilting only	$\times$	✓	69.54	26.40	76.30	48.10
Look-ahead only	✓	$\times$	69.69	26.20	76.71	46.49
$S^{3}$ (full, ours)	✓	✓	70.21	30.20	77.86	49.57

Table 2: Ablation study decomposing the contributions of look-ahead search and Gibbs tilting in

S^{3}

(LLaDA-8B-Instruct,

N{=}4

b{=}2

K{=}N{\cdot}b{=}8

, block length 64, generation length 128). ✓/× indicate whether each mechanism is active. Scores represent task accuracy (%).

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 5 — Figure 5: Distributional shift of verifier scores under $S^{3}$ search on MATH-500. Sequential resampling during denoising concentrates particles in higher-reward regions compared to direct sampling from $p_{0}$ .

Dataset	$\alpha_{1}$	$\alpha_{2}$	$\alpha_{3}$	$\alpha_{4}$	$\alpha_{5}$	$\alpha_{c}$
GSM8K	0.20	0.25	0.25	0.10	0.20	0.00
MATH-500	0.25	0.20	0.25	0.10	0.20	0.00
ARC-Challenge	0.20	0.00	0.00	0.15	0.20	0.45
TruthfulQA	0.20	0.00	0.00	0.15	0.20	0.45
Countdown	0.15	0.10	0.10	0.10	0.15	0.40
Sudoku	0.05	0.00	0.00	0.05	0.10	0.80

Appendix

Table 3: Weight profiles

(\alpha_{1},\dots,\alpha_{5},\alpha_{c})

per dataset, corresponding to structural completeness, arithmetic consistency, and answer reachability, model confidence, non-degeneracy, and constraint satisfaction respectively. All rows sum to

1.0

Dataset	Method	Avg Eff. Tokens $(\downarrow)$	Total Latency (s) $(\downarrow)$	Wall Time (s) $(\downarrow)$
ARC-C	Baseline	5.67	3648.21	1829.87
	Best-of- $K$	5.43	28803.25	14431.86
	Look-ahead only	6.08	29450.49	14756.30
	Tilting only	5.43	28804.53	14431.20
	$S^{3}$	4.34	29100.42	14580.21
GSM8K	Baseline	3.03	3919.06	1964.63

Appendix

Table 4: Average effective output tokens (

\downarrow

), total latency (

\downarrow

, seconds), and wall time (

\downarrow

, seconds) for each ablation variant and dataset (

N=4

b=2

K=8

, LLaDA-8B-Instruct, block_length =64).

Dataset	Block	Mode	Avg Tok. $\downarrow$	Latency (s) $\downarrow$	Wall (s) $\downarrow$
ARC-C	2	Base	16.92	7.42	6.10
	2	$S^{3}$	17.10	71262.09	35680.83
	4	Base	8.92	7.43	6.12
	4	$S^{3}$	8.46	71038.65	35550.89
	8	Base	8.04	7.51	6.18
	8	$S^{3}$	7.45	37840.20	18952.40

Appendix

Table 5: Average effective tokens (

\downarrow

), total latency (

\downarrow

, s), and wall time (

\downarrow

, s) across all block sizes for LLaDA-8B-Instruct. Base : single-trajectory decoding;

S^{3}

: multi-block decoding with BFS search and SSP resampling.

Explore similar work

Test-Time Scaling in the Wild: Why Exploitation, Not Exploration, Is the Bottleneck
Aug 19, 2026Davide Romano, Kanak Raj, Jerrod Parker +1Test-Time Scaling Language Model Decoding
Thomson Reuters Foundational Research, 5 Canada Square, London
Guided Trajectory Optimization with Sparse Scaling for Test-Time Diffusion
May 21, 2026Gang Dai, Yining Huang, Yiming Xia +2Test-Time Scaling Diffusion-Based Image Generation
Guangdong University of Technology · South China University of Technology · Nanyang Technological University
Test-Time Scaling via Error Localization
Jul 23, 2026Rajiv Shailesh Chitale, Rahul Madhavan, Taneesh Gupta +2Test-Time Scaling Efficient Language Model Inference
Google DeepMind

Aug 19, 2026J/K move · Enter open · S save

Test-Time Scaling in the Wild: Why Exploitation, Not Exploration, Is the Bottleneck

Davide Romano, Kanak Raj, Jerrod Parker, Daniele Giofrè

Thomson Reuters Foundational Research, 5 Canada Square, London

Test-Time Scaling Language Model Decoding Open-Ended Generation Language Model Generation Evaluation Reward Modeling

Open paper AskPDF arXiv

Test-time scaling (TTS) improves language model outputs by spending additional inference compute - generating multiple candidates, searching over partial sequences, or iteratively refining drafts. These techniques yield large gains on mathematics and code, but have been developed and stress-tested almost exclusively on tasks where verification is straightforward. We conduct the first compute-normalised comparison of five TTS families across five open-ended generation benchmarks spanning medicine, law, finance, general chat, and creative writing - grounded in a unified framework that decomposes the effectiveness of each method's token budget into exploration and exploitation. The answer depends on which side of that decomposition you examine. Scaling exploration works: the best candidate in the pool improves steadily with compute across all settings. What breaks is exploitation - the step that converts a rich candidate pool into a final output. With state-of-the-art generators, reward models correlate at only

ρ_v \approx 0.12

with true quality, rendering selection near-random regardless of budget. Tree search amplifies this failure through diversity collapse. Refinement helps on one of five benchmarks; its apparent gains elsewhere are confounded. Only synthesis across candidates (Fusion) consistently improves over single-sample baselines, yet still recovers only ~40% of available quality. The candidate pool is not the bottleneck - choosing from it is.

Abstract

Explore similar work

Test-Time Scaling in the Wild: Why Exploitation, Not Exploration, Is the Bottleneck

Guided Trajectory Optimization with Sparse Scaling for Test-Time Diffusion

Test-Time Scaling via Error Localization