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-26
🔗 代码/项目: GITHUB
💡 一句话要点
提出基于奖励引导拼接的扩散语言模型测试时缩放方法，提升复杂推理任务性能。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 扩散模型 语言模型 推理 奖励模型 自回归模型 步骤级拼接 测试时缩放
📋 核心要点
- 现有大型语言模型推理方法通常在轨迹层面聚合结果，忽略了部分正确推理过程中的有用信息。
- 论文提出一种基于扩散模型的步骤级拼接方法，通过奖励模型选择高质量步骤，组合成完整推理链。
- 实验表明，该方法在数学和编码任务上显著提升了准确率，并降低了推理延迟，尤其在困难问题上效果更佳。
📝 摘要（中文）
本文提出了一种名为“拼接噪声扩散思想”（Stitching Noisy Diffusion Thoughts）的自洽框架，用于提升大型语言模型的推理能力。该框架利用廉价的扩散采样推理生成可复用的步骤级候选方案。对于给定的问题，首先使用掩码扩散语言模型采样多个多样且低成本的推理轨迹；然后，使用现成的过程奖励模型（PRM）对每个中间步骤进行评分；最后，将这些来自不同轨迹的最高质量步骤拼接成一个组合的推理链。该推理链作为条件，引导自回归（AR）模型（求解器）重新计算最终答案。这种模块化流程将探索（扩散）与评估和解决方案综合分离，避免了统一混合架构，同时保留了广泛的搜索能力。在数学推理基准测试中，发现步骤级重组在更困难的问题上最有益。消融实验表明，最终的AR求解器在将拼接的但不完美的推理链转换为准确答案方面至关重要。通过并行、独立的rollout进行低置信度的扩散采样，该免训练框架在六个数学和编码任务中平均准确率提高了高达23.8%。同时，相对于传统的扩散模型（例如Dream，LLaDA）和统一架构（例如TiDAR），实现了高达1.8倍的延迟降低。
🔬 方法详解
问题定义：现有的大型语言模型在进行复杂推理时，通常采用生成多个推理轨迹然后选择最佳轨迹或对最终答案进行投票的方法。这种方法忽略了那些“几乎正确”的轨迹中包含的有用中间步骤信息，造成了计算资源的浪费，并且限制了模型利用不同推理路径中的优势信息。
核心思路：本文的核心思路是将推理过程分解为多个步骤，利用扩散模型生成大量低成本的推理步骤候选，然后使用奖励模型对这些步骤进行评估，选择高质量的步骤进行拼接，最终形成一个完整的推理链。这样可以充分利用不同推理轨迹中的优势信息，提高推理的准确性和效率。
技术框架：该方法包含三个主要阶段：(1) 扩散采样：使用掩码扩散语言模型生成多个低成本的推理轨迹，每个轨迹包含多个推理步骤。(2) 步骤评估：使用预训练的过程奖励模型（PRM）对每个推理步骤进行评分，评估其质量。(3) 拼接与求解：将来自不同轨迹的最高质量步骤拼接成一个组合的推理链，然后使用自回归（AR）模型（求解器）以该推理链为条件，重新计算最终答案。
关键创新：该方法最重要的创新点在于步骤级的重组。与传统的轨迹级选择或投票方法不同，该方法能够充分利用不同推理轨迹中的优势信息，将高质量的步骤组合成一个更完整的推理链。此外，该方法将探索（扩散）与评估和解决方案综合分离，避免了统一混合架构的局限性，同时保留了广泛的搜索能力。
关键设计：该方法的关键设计包括：(1) 使用掩码扩散语言模型生成多样化的推理轨迹，保证步骤的多样性。(2) 使用预训练的过程奖励模型（PRM）对步骤进行评分，保证步骤的质量。(3) 使用自回归（AR）模型作为求解器，将拼接的推理链转化为最终答案，提高答案的准确性。论文中没有明确提及具体的参数设置、损失函数或网络结构等技术细节，这些可能依赖于具体的扩散模型和自回归模型选择。
🖼️ 关键图片
📊 实验亮点
实验结果表明，该方法在六个数学和编码任务中平均准确率提高了高达23.8%。同时，相对于传统的扩散模型（例如Dream，LLaDA）和统一架构（例如TiDAR），实现了高达1.8倍的延迟降低。尤其是在困难问题上，步骤级重组带来的收益更加明显。
🎯 应用场景
该研究成果可应用于需要复杂推理能力的各种场景，例如数学问题求解、代码生成、知识图谱推理等。通过提高推理的准确性和效率，可以提升AI在这些领域的应用价值，并为开发更智能的AI系统提供新的思路。
📄 摘要（原文）
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 at https://github.com/roymiles/diffusion-stitching.