Test-Time Scaling with Diffusion Language Models via Reward-Guided Stitching
作者: Roy Miles, Aysim Toker, Andreea-Maria Oncescu, Songcen Xu, Jiankang Deng, Ismail Elezi
分类: cs.CL, cs.AI
发布日期: 2026-02-28
💡 一句话要点
提出基于奖励引导拼接的扩散语言模型测试时缩放方法，提升复杂推理任务性能。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 扩散模型 语言模型 推理 奖励模型 自回归模型 步骤级拼接 测试时缩放
📋 核心要点
- 现有大语言模型推理方法通常在轨迹层面聚合结果，忽略了部分正确尝试中的有用中间步骤。
- 论文提出一种基于扩散模型的步骤级拼接框架，利用奖励模型选择高质量的中间步骤，组合成更优的推理链。
- 实验表明，该方法在数学和编码任务上显著提升了准确率，并降低了推理延迟，尤其是在困难问题上。
📝 摘要（中文）
本文提出了一种名为“拼接噪声扩散思想”（Stitching Noisy Diffusion Thoughts）的自洽框架，用于提升大型语言模型的推理能力。该框架利用廉价的扩散采样推理生成可复用的步骤级候选方案池。对于给定的问题，首先使用掩码扩散语言模型采样多个多样且低成本的推理轨迹；然后，使用现成的过程奖励模型（PRM）对每个中间步骤进行评分；最后，将这些来自不同轨迹的最高质量步骤拼接成一个组合的理由。该理由作为条件，引导自回归（AR）模型（求解器）重新计算最终答案。这种模块化流程将探索（扩散）与评估和解决方案综合分离，避免了统一混合架构，同时保留了广泛的搜索能力。在数学推理基准测试中，发现步骤级重组在较难的问题上最有益。消融实验表明，最终的AR求解器在将拼接的但不完美的理由转化为准确答案方面至关重要。该框架无需训练，通过并行、独立的rollout进行低置信度的扩散采样，在六个数学和编码任务中，平均准确率提高了高达23.8%。同时，相对于传统的扩散模型（如Dream、LLaDA）和统一架构（如TiDAR），延迟降低了高达1.8倍。
🔬 方法详解
问题定义：现有的大语言模型推理方法，例如选择最佳轨迹或对最终答案进行投票，通常在轨迹层面进行聚合。这种方法忽略了部分正确的推理尝试中可能存在的有价值的中间步骤，导致信息损失，尤其是在解决复杂问题时。
核心思路：论文的核心思路是将推理过程分解为多个步骤，并利用扩散模型生成多个不同的推理轨迹。然后，使用奖励模型对每个步骤进行评估，选择高质量的步骤进行拼接，形成一个更完整的推理链。最后，使用自回归模型基于拼接的推理链生成最终答案。这种方法能够充分利用不同推理轨迹中的信息，提高推理的准确性和效率。
技术框架：该框架主要包含三个阶段：1) 使用掩码扩散语言模型生成多个低成本的推理轨迹；2) 使用预训练的过程奖励模型（PRM）对每个中间步骤进行评分；3) 将来自不同轨迹的最高质量步骤拼接成一个组合的理由，并将其作为条件输入到自回归模型中，以生成最终答案。整个框架是模块化的，将探索（扩散）与评估和解决方案综合分离。
关键创新：该方法最重要的创新点在于步骤级的推理拼接。与传统的轨迹级聚合方法不同，该方法能够充分利用不同推理轨迹中的中间步骤信息，从而提高推理的准确性和鲁棒性。此外，使用扩散模型生成多样化的推理轨迹，并使用奖励模型进行步骤评估，也为推理过程的优化提供了新的思路。
关键设计：论文使用了掩码扩散语言模型进行推理轨迹的生成，并使用预训练的奖励模型进行步骤评估。奖励模型的设计至关重要，需要能够准确地评估每个步骤的质量。此外，自回归模型的设计也需要能够有效地利用拼接的推理链信息，生成准确的最终答案。具体的参数设置和网络结构在论文中有详细描述。
🖼️ 关键图片
📊 实验亮点
该框架在六个数学和编码任务中，平均准确率提高了高达23.8%。同时，相对于传统的扩散模型（如Dream、LLaDA）和统一架构（如TiDAR），延迟降低了高达1.8倍。实验结果表明，步骤级重组在较难的问题上最有益，并且最终的AR求解器在将拼接的但不完美的理由转化为准确答案方面至关重要。
🎯 应用场景
该研究成果可应用于各种需要复杂推理的场景，例如数学问题求解、代码生成、逻辑推理等。通过提高推理的准确性和效率，可以提升人工智能在这些领域的应用价值。此外，该方法还可以用于教育领域，帮助学生更好地理解和解决问题。
📄 摘要（原文）
Reasoning with large language models often benefits from generating multiple chains-of-thought, but existing aggregation strategies are typically trajectory-level (e.g., selecting the best trace or voting on the final answer), discarding useful intermediate work from partial or "nearly correct" attempts. We propose Stitching Noisy Diffusion Thoughts, a self-consistency framework that turns cheap diffusion-sampled reasoning into a reusable pool of step-level candidates. Given a problem, we (i) sample many diverse, low-cost reasoning trajectories using a masked diffusion language model, (ii) score every intermediate step with an off-the-shelf process reward model (PRM), and (iii) stitch these highest-quality steps across trajectories into a composite rationale. This rationale then conditions an autoregressive (AR) model (solver) to recompute only the final answer. This modular pipeline separates exploration (diffusion) from evaluation and solution synthesis, avoiding monolithic unified hybrids while preserving broad search. Across math reasoning benchmarks, we find that step-level recombination is most beneficial on harder problems, and ablations highlight the importance of the final AR solver in converting stitched but imperfect rationales into accurate answers. Using low-confidence diffusion sampling with parallel, independent rollouts, our training-free framework improves average accuracy by up to 23.8% across six math and coding tasks. At the same time, it achieves up to a 1.8x latency reduction relative to both traditional diffusion models (e.g., Dream, LLaDA) and unified architectures (e.g., TiDAR). Code is available atthis https URL.