通过分层搜索提升扩散语言模型推理质量,无需重新训练。
$S^3$: Stratified Scaling Search for Test-Time in Diffusion Language Models
- 在去噪过程中分步扩展候选路径,用轻量验证器评估并重采样优质路径。
- 在MATH-500等数据集上显著提升数学推理性能,最高增益达18.7%。
- 适合追求高推理质量但无法重训练的场景,兼容现有模型与解码策略。
测试时缩放研究固定扩散语言模型(DLM)是否能在不额外训练的情况下,通过增加推理计算量生成更优输出。然而,朴素的best-of-K采样存在根本局限:反复从同一基础扩散分布中取样,而该分布的高概率区域常与高质量输出不一致。我们提出S³(分层缩放搜索),一种基于轻量级无参考验证器的古典验证引导搜索方法,将计算资源重新分配至去噪过程而非仅在最终输出阶段。每一步去噪中,S³扩展多个候选轨迹,以轻量验证器评估其质量,选择性地重采样前景良好的路径,同时保持搜索前沿的多样性。此过程有效逼近奖励倾斜的采样分布,更倾向高质量输出,且仍锚定于模型先验。在LLaDA-8B-Instruct模型上对MATH-500、GSM8K、ARC-Challenge和TruthfulQA的实验表明,S³在各基准上均持续提升性能,尤其在数学推理任务上增益最大,最高达18.7%,且不改变原始模型与解码调度。结果表明,对去噪轨迹进行经典搜索为DLM提供了实用的测试时缩放机制。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。