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

• URL: http://arxiv.org/abs/2604.06260v1
• Date: Tue, 07 Apr 2026 00:51:06 GMT
• Title: $S^3$: Stratified Scaling Search for Test-Time in Diffusion Language Models
• Abstract summary: Test-time scaling investigates whether a fixed diffusion language model (DLM) can generate better outputs when given more inference compute.<n>We propose $S3$ (Stratified Scaling Search), a verifier-guided search method that improves generation by reallocating compute during the denoising process.<n>Experiments with LLaDA-8B-Instruct on MATH-500, GSM8K, ARC-Challenge, and TruthfulQA demonstrate that $S3$ consistently improves performance across benchmarks.
• Score: 10.157422365382933
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Refining Over Resampling: Test-Time Self-Correction for LLM Reasoning [3.8191765711200065]
  We propose a verifier-free breadth--depth refinement framework that uses test-time compute to both explore and improve candidate solutions.<n>Our method consistently improves over greedy decoding, majority voting, verifier-based best-of-$N$, beam search, and look decoding across multiple open-weight models.
  arXiv  Detail & Related papers  (2026-08-06T06:38:37Z)
• Flash-BoN: Instant Drafts for Inference-Time Scaling in Diffusion Models [91.19049123903137]
  Inference-time scaling for text-to-image generation has progressed from simple Best-of-$N$ sampling to guided search methods.<n>These approaches focus on when and how often to verify during denoising but largely treat the cost of generation itself as fixed.<n>We show that under wall-clock evaluation, simple BoN already matches or outperforms several guided search techniques.
  arXiv  Detail & Related papers  (2026-07-05T19:02:32Z)
• $\nabla$-Reasoner: LLM Reasoning via Test-Time Gradient Descent in Latent Space [71.23672814629448]
  $nabla$-Reasoner is an iterative generation framework that integrates differentiable optimization over token logits into the decoding loop.<n>$nabla$-Reasoner achieves over 20% accuracy improvement on a challenging mathematical reasoning benchmark.
  arXiv  Detail & Related papers  (2026-03-05T08:42:54Z)
• Free Lunch for Pass@$k$? Low Cost Diverse Sampling for Diffusion Language Models [17.37935640125399]
  We propose a training free, low cost intervention to enhance generative diversity in Diffusion Language Models.<n>Our approach modifies intermediate samples in a batch sequentially, where each sample is repelled from the feature space of previous samples.<n>Unlike prior methods that require retraining or beam search, our strategy incurs negligible computational overhead.
  arXiv  Detail & Related papers  (2026-03-05T07:35:07Z)
• Test-Time Scaling with Diffusion Language Models via Reward-Guided Stitching [66.39914384073145]
  We propose a self-consistency framework that turns cheap diffusion-sampled reasoning into a reusable pool of step-level candidates.<n>We find that step-level recombination is most beneficial on harder problems.<n>Our training-free framework improves average accuracy by up to 2 across six math and coding tasks.
  arXiv  Detail & Related papers  (2026-02-26T11:08:39Z)
• CarBoN: Calibrated Best-of-N Sampling Improves Test-time Reasoning [62.56541355300587]
  We introduce a general test-time calibration framework that adaptively modifies the model toward high-reward reasoning paths.<n>Within this framework, we propose CarBoN, a two-phase method that first explores the solution space and then learns a calibration of the logits.<n>Experiments on MATH-500 and AIME-2024 show that CarBoN improves efficiency, with up to $4times$ fewer rollouts to reach the same accuracy.
  arXiv  Detail & Related papers  (2025-10-17T14:04:37Z)
• Inference-Time Scaling of Diffusion Language Models with Particle Gibbs Sampling [70.8832906871441]
  We study how to steer generation toward desired rewards without retraining the models.<n>Prior methods typically resample or filter within a single denoising trajectory, optimizing rewards step-by-step without trajectory-level refinement.<n>We introduce particle Gibbs sampling for diffusion language models (PG-DLM), a novel inference-time algorithm enabling trajectory-level refinement while preserving generation perplexity.
  arXiv  Detail & Related papers  (2025-07-11T08:00:47Z)
• Diffusion Tree Sampling: Scalable inference-time alignment of diffusion models [13.312007032203857]
  Adapting a pretrained diffusion model to new objectives at inference time remains an open problem in generative modeling.<n>We introduce a tree-based approach that samples from the reward-aligned target density by propagating terminal rewards back through the diffusion chain.<n>By reusing information from previous generations, we get an anytime algorithm that turns additional compute into steadily better samples.
  arXiv  Detail & Related papers  (2025-06-25T17:59:10Z)
• $\texttt{SPECS}$: Faster Test-Time Scaling through Speculative Drafts [55.231201692232894]
  $textttSPECS$ is a latency-aware test-time scaling method inspired by speculative decoding.<n>Our results show that $textttSPECS$matches or surpasses beam search accuracy while reducing latency by up to $sim$19.1%.
  arXiv  Detail & Related papers  (2025-06-15T05:50:05Z)
• Test-Time Scaling of Diffusion Models via Noise Trajectory Search [10.8507840358202]
  We introduce an $epsilon$-greedy search algorithm that globally explores at extreme timesteps and locally exploits during the intermediate steps where de-mixing occurs.<n>Experiments on EDM and Stable Diffusion reveal state-of-the-art scores for class-conditioned/text-to-image generation.
  arXiv  Detail & Related papers  (2025-05-24T19:13:29Z)
• Sample, Don't Search: Rethinking Test-Time Alignment for Language Models [55.2480439325792]
  We introduce QAlign, a new test-time alignment approach.<n>As we scale test-time compute, QAlign converges to sampling from the optimal aligned distribution for each individual prompt.<n>By adopting recent advances in Markov chain Monte Carlo for text generation, our method enables better-aligned outputs without modifying the underlying model or even requiring logit access.
  arXiv  Detail & Related papers  (2025-04-04T00:41:40Z)
• Dimension-free Score Matching and Time Bootstrapping for Diffusion Models [19.62665684173391]
  Diffusion models generate samples by estimating the score function of the target distribution at various noise levels.<n>We introduce a martingale-based error decomposition and sharp variance bounds, enabling efficient learning from dependent data.<n>Building on these insights, we propose Bootstrapped Score Matching (BSM), a variance reduction technique that leverages previously learned scores to improve accuracy at higher noise levels.
  arXiv  Detail & Related papers  (2025-02-14T18:32:22Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.