$S^3$: 扩散语言模型测试时的分层缩放搜索
机器学习
2026-04-09 v1 人工智能
摘要
测试时缩放研究一个固定的扩散语言模型（DLM）在不进行额外训练的情况下，给定更多推理计算时能否生成更好的输出。然而，朴素的最佳-采样从根本上受到限制，因为它反复从相同的基扩散分布中采样，而该分布的高概率区域往往与高质量输出不一致。我们提出 （分层缩放搜索，Stratified Scaling Search），一种经典的验证器引导搜索方法，通过在去噪过程中重新分配计算而非仅在最终输出阶段来改进生成。在每个去噪步骤中， 扩展多个候选轨迹，用一个轻量级的无参考验证器评估它们，并选择性地重新采样有前景的候选，同时保持搜索前沿的多样性。该过程有效地近似了一个奖励倾斜采样分布，有利于高质量输出，同时锚定在模型先验上。在 LLaDA-8B-Instruct 上对 MATH-500、GSM8K、ARC-Challenge 和 TruthfulQA 的实验表明， 在各个基准上持续提升性能，在数学推理任务上取得最大增益，同时保持底层模型和解码调度不变。这些结果表明，对去噪轨迹的经典搜索为扩散语言模型中的测试时缩放提供了实用机制。
引用
@article{arxiv.2604.06260,
title = {$S^3$: Stratified Scaling Search for Test-Time in Diffusion Language Models},
author = {Ahsan Bilal and Muhammad Ahmed Mohsin and Muhammad Umer and Asad Aali and Muhammad Usman Khanzada and Muhammad Usman Rafique and Zihao He and Emily Fox and Dean F. Hougen},
journal= {arXiv preprint arXiv:2604.06260},
year = {2026}
}
备注
Submitted to COLM 2026