论文

$S^3$:扩散语言模型中测试时间的分层缩放搜索

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

模型推理测试时计算扩展

摘要

测试时间缩放研究固定扩散语言模型 (DLM) 在给予更多推理计算时是否可以生成更好的输出,而无需额外训练。然而,朴素的 best-of-K$ 采样从根本上是有限的,因为它重复地从相同的基础扩散分布中抽取,其高概率区域通常与高质量输出不一致。我们提出$S^3$(分层缩放搜索),这是一种经典的验证者引导搜索方法,通过在去噪过程中而不是仅在最终输出阶段重新分配计算来改进生成。在每个去噪步骤中,$S^3$ 扩展多个候选轨迹,使用轻量级无参考验证器对其进行评估,并有选择地对有希望的候选者进行重新采样,同时保留搜索前沿内的多样性。该过程有效地近似了奖励倾斜的采样分布,有利于更高质量的输出,同时保持锚定于模型先验。在 MATH-500、GSM8K、ARC-Challenge 和 TruthfulQA 上使用 LLaDA-8B-Instruct 进行的实验表明,$S^3$ 在基准测试中持续提高性能,在数学推理任务上实现最大收益,同时保持底层模型和解码时间表不变。这些结果表明,对去噪轨迹的经典搜索为 DLM 中的测试时间缩放提供了实用的机制。