arXiv:2507.08390cs.LG2025-07被引 2

通过轨迹精炼提升扩散语言模型推理时的生成质量

Inference-Time Scaling of Diffusion Language Models via Trajectory Refinement

  • 引入粒子吉布斯采样,可在完整生成轨迹上迭代优化
  • 2.9个粒子平均下达90.07%准确率,24.8个粒子达94.47%准确率
  • 支持按需分配计算资源,适合高效奖励引导生成任务

离散扩散模型作为自回归语言模型的有力替代,通过大规模训练已达到相当性能。然而,推理时控制仍相对未被充分探索。本文研究如何在不重新训练模型的前提下,引导生成结果向目标奖励靠拢。以往方法通常在单条去噪轨迹内进行重采样或过滤,逐步优化奖励,但缺乏对轨迹层面的精炼。为此,本文提出粒子吉布斯采样(PG-DLM),一种面向扩散语言模型的推理时算法,通过构建全去噪轨迹上的马尔可夫链,并应用条件顺序蒙特卡洛核实现轨迹重采样。该方法引入新的可扩展维度——精炼迭代次数,而此前方法不具备此能力。即使增加并行样本数带来的收益趋于饱和,提升迭代次数仍能持续改进效果。此外,PG-DLM支持自适应计算分配,仅在需要时执行额外迭代,进一步提升效率。本文提供了收敛性与方差界理论保证,并分析了不同扩展轴间的权衡。实验表明,在各类计算预算下,PG-DLM均优于现有方法。在GSM8K任务中,平均使用2.9个粒子时取得90.07%准确率,平均使用24.8个粒子时达94.47%准确率。

原文摘要 · Abstract (English)

Discrete diffusion models have recently emerged as strong alternatives to autoregressive language models, matching their performance through large-scale training. However, inference-time control remains relatively underexplored. In this work, we study how to steer generation toward desired rewards without retraining the models. Prior methods typically resample or filter within a single denoising trajectory, optimizing rewards step-by-step without trajectory-level refinement. We introduce particle Gibbs sampling for diffusion language models (PG-DLM), an inference-time algorithm enabling trajectory-level refinement. PG-DLM constructs a Markov chain over full denoising trajectories and applies a conditional sequential Monte Carlo kernel to resample them. By doing so, PG-DLM introduces a new scaling axis, the number of refinement iterations, which is unavailable to prior methods. Increasing iterations remains effective even as gains from adding more parallel samples saturate. Furthermore, PG-DLM enables adaptive compute allocation by performing additional iterations only when needed, leading to further efficiency gains. We derive theoretical guarantees for convergence and variance bounds, and analyze trade-offs across different scaling axes. Empirically, PG-DLM outperforms prior methods across compute budgets on reward-guided generation tasks. On GSM8K, it achieves 90.07% accuracy with 2.9 particles on average and 94.47% accuracy with 24.8 particles on average.

扩散模型推理优化轨迹精炼奖励引导

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。