# S3: Stratified Scaling Search for Test-Time in Diffusion Language Models

arXiv: https://arxiv.org/abs/2604.06260 · HTML: https://mlsanity.com/paper/2604.06260/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion reasoning and test-time scaling, one of 17 directions in LLM reasoning.
- Builds on 23 of its 34 references that are in the map.
- First posted on arXiv 7 Apr 2026.

## Placed in

- LLM reasoning / [Diffusion reasoning and test-time scaling](https://mlsanity.com/lr/diffusion-reasoning-and-test-time-scaling/)
- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 23 of 23)

- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [TruthfulQA: Measuring How Models Mimic Human Falsehoods](https://mlsanity.com/paper/2109.07958/) (2021)
- [Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters](https://mlsanity.com/paper/2408.03314/) (2024)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [Diffusion-LM Improves Controllable Text Generation](https://mlsanity.com/paper/2205.14217/) (2022)
- [Large Language Monkeys: Scaling Inference Compute with Repeated Sampling](https://mlsanity.com/paper/2407.21787/) (2024)
- [Large Language Diffusion Models](https://mlsanity.com/paper/2502.09992/) (2025)
- [Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions](https://mlsanity.com/paper/2502.06768/) (2025)
- [Scaling up Masked Diffusion Models on Text](https://mlsanity.com/paper/2410.18514/) (2024)
- [d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learning](https://mlsanity.com/paper/2504.12216/) (2025)
- [Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-Based Decoding](https://mlsanity.com/paper/2408.08252/) (2024)
- [TFG: Unified Training-Free Guidance for Diffusion Models](https://mlsanity.com/paper/2409.15761/) (2024)
- [Energy-Based Diffusion Language Models for Text Generation](https://mlsanity.com/paper/2410.21357/) (2024)
- [Inference-time Scaling of Diffusion Models through Classical Search](https://mlsanity.com/paper/2505.23614/) (2025)
- [Sequential Controlled Langevin Diffusions](https://mlsanity.com/paper/2412.07081/) (2024)
- [Noise Hypernetworks: Amortizing Test-Time Compute in Diffusion Models](https://mlsanity.com/paper/2508.09968/) (2025)
- [Discrete Diffusion Language Model for Efficient Text Summarization](https://mlsanity.com/paper/2407.10998/) (2024)
- [Test-Time Scaling in Diffusion LLMs via Hidden Semi-Autoregressive Experts](https://mlsanity.com/paper/2510.05040/) (2025)
- [Prism: Efficient Test-Time Scaling via Hierarchical Search and Self-Verification for Discrete Diffusion Language Models](https://mlsanity.com/paper/2602.01842/) (2026)
- [Rethinking Optimal Verification Granularity for Compute-Efficient Test-Time Scaling](https://mlsanity.com/paper/2505.11730/) (2025)

## Built on by (most cited first, 2 of 2)

- [Looped Diffusion Language Models](https://mlsanity.com/paper/2605.26106/) (2026)
- [Escaping Confidence Trap: Evolutionary Decoding for Mathematical Reasoning in Diffusion LLMs](https://mlsanity.com/paper/2608.00605/) (2026)

## Abstract (arXiv)

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 S3 (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, S3 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 S3 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.