Test Time Scaling with Diffusion Language Models via Reward Guided Stitching¶

会议: NeurIPS 2026
arXiv: 2602.22871
OpenReview: https://openreview.net/forum?id=54JyknakHB
代码: https://github.com/roymiles/diffusion-stitching
领域: LLM 推理
关键词: 扩散语言模型, 测试时扩展, 过程奖励模型, 步骤级缝合, 自回归求解器

一句话总结¶

针对传统测试时搜索在整条轨迹级聚合导致的局部有效计算浪费与高昂串行延迟，本文提出奖励引导缝合框架，利用掩码扩散语言模型低成本并行采样多样化思维链，借助 PRM 进行步骤级细粒度打分与跨轨迹置信度拼接，最后引导轻量自回归求解器快速重算最终答案，在显著降低前向传播次数的同时超越传统扩散与自回归基线。

研究背景与动机¶

大语言模型在复杂数学推导与程序合成等高难度任务中展现出强大的通用求解能力，但其高准确率严重依赖于测试时计算（Test-Time Compute）的扩展。在现有的测试时搜索范式中，最普遍的做法是通过温度采样生成多条思维链（CoT），随后依赖自洽性多数投票（Self-Consistency）或全轨迹级验证器（Verifier）挑选最佳候选。然而，推理轨迹在长程推导中天然具有高脆弱性：单步逻辑谬误或微小计算偏差便会导致后续链条全盘崩溃。轨迹级聚合将整条解答视为“全有或全无”的粗粒度单元，直接丢弃了那些即使最终失败却包含极其关键局部正确子推导的“近乎正确”轨迹，导致海量有价值的中间推理算力被白白浪费。

与此同时，传统的自回归模型在进行多路树状或轨迹级探索时面临高昂的串行自回归解码瓶颈，使得测试时扩展在低延迟应用场景中难以承受。尽管掩码扩散语言模型（dLLM，如 LLaDA、Dream）能够通过双向注意力与并行去噪在极低时间开销下生成高度多样化的推理路径，但扩散模型自身在缺乏因果单向约束时极易出现全局逻辑不连贯、晚期事实飘移以及最终答案整合困难等缺陷。近期涌现的将扩散细化与自回归生成强行绑定的统一混合架构（如 TiDAR），不仅耦合了搜索探索、置信度评估与答案生成三个正交阶段，而且往往导致整体数学推理准确率大幅落后于主流自回归基线。

面对“自回归搜索耗时过长、扩散探索容易脱轨、轨迹级粗粒度聚合浪费局部正确推导”的结构性矛盾，本文选择解耦探索、评估与整合环节。核心 idea：将低成本扩散采样转化为可复用的中间推理步骤池，利用过程奖励模型（PRM）筛选跨轨迹的高置信度局部步骤并缝合成复合证据链，最后驱动轻量自回归模型仅重算最终答案。

方法详解¶

整体框架¶

Diffusion Stitching 框架将复杂的测试时推理过程显式解耦为三阶段专业化流水线：多样化探索（扩散模型）、细粒度质量评估（过程奖励模型）以及证据重构与求解（轻量自回归求解器）。给定一个输入问题，系统首先使用掩码扩散语言模型并行采样多条独立的推理轨迹，并根据语法或语义边界将其确定性切分为步骤序列；随后利用离架的过程奖励模型（PRM）在单次前向传播中为每一个局部步骤打出标量置信度得分，构建出全局候选步骤池；最后，流水线依据全局最优轨迹构建骨干锚点，并将超过置信度阈值的跨轨迹高分步骤按时序插入，拼接成包含置信度前缀的证据列表，供轻量自回归模型快速聚焦并推导闭环答案。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入问题 x"] --> B["扩散多样性探索<br/>并行采样 N 条独立低成本轨迹"]
    B --> C["步骤划分与边界标记<br/>插入标记 token 确定性分句"]
    C --> D["PRM 步骤级评分<br/>单次前向评估局部质量 r_t"]
    D --> E["锚点选择与置信度缝合<br/>保留最佳全轨迹并拼接高分步"]
    E --> F["时序证据链构建<br/>按时间步排序与置信度标签标注"]
    F --> G["自回归求解器重算<br/>条件解码并输出最终答案 y_hat"]

关键设计¶

1. 扩散多样性探索与确定性步骤切分：解耦探索开销并构建局部步骤池

传统自回归模型在采样多条长链时面临极高的顺序前向解码延迟，而掩码扩散语言模型（dLLM）从全掩码或部分掩码模板出发，通过双向注意力网络对所有掩码位置进行迭代并行去噪。为了在极短时间内探索尽可能丰富的解题空间，系统在温度参数 $\tau > 0$ 下并行展开 $N$ 条独立的扩散轨迹： $$p_i^{(k)} = \mathrm{softmax}(\ell_i^{(k)} / \tau), \quad \hat{x}_i^{(k)} \sim p_i^{(k)}$$ 在采样过程中，模型引入基于置信度 $c_i^{(k)} = \max_j p_{i,j}^{(k)}$ 的重掩码策略，并特意采用较低的置信度阈值 $\gamma$（如 $\gamma = 0.7$），允许模型在单次去噪中容忍轻微噪声以换取最大化的探索广度与计算速度。生成完成后，系统依据任务特性对轨迹进行确定性步骤划分：数学问题按自然句法与换行符进行切分，代码任务则按独立代码行或自然语言规划段落切分，将每条轨迹 $\tau^{(n)}$ 转化为人类可读且语义自洽的离散步骤序列 $(s_1^{(n)}, \dots, s_{T_n}^{(n)})$。

2. PRM 步骤级单次前向评估：保留局部有效推导与误差隔离

为了克服传统轨迹级验证将整条推理链“一刀切”的弊端，系统利用离架过程奖励模型（PRM）为所有轨迹中的每一个中间推导赋予细粒度质量度量。常规的逐步调用 PRM 会带来巨大的计算开销，本文设计了基于特殊标记 token $\langle m \rangle$ 的单次前向打分机制：在每一步的边界处显式插入标记，使得 PRM 仅需对整个生成文本进行单次前向传播，即可直接读取该标记处的预测概率作为该步骤的条件置信度： $$r_t^{(n)} = p_\phi(\langle m \rangle \mid x, s_{1:t}^{(n)}) \in [0, 1]$$ 通过遍历 $N$ 条轨迹的所有中间步骤，系统构建起一个跨越不同思考路径的全局候选步骤池 $\mathcal{P} = \{(n, t, s_t^{(n)}, r_t^{(n)})\mid n \in [N], t \in [T_n]\}$。这种步骤级视图能够精准识别出某些轨迹中局部正确、逻辑严谨的定理推演与中间计算，哪怕该轨迹在最后几步发生了算术偏差或逻辑脱轨，其早期的高价值子结果依然能被完整保留并注入全局池。

3. 锚点选择与跨轨迹置信度缝合：兼顾全局骨架完整性与跨链互补性

纯粹从全局池中按置信度筛选步骤容易造成推导断裂或上下文脱节，而直接混拼多条链条又会引入冗余与矛盾。为此，缝合机制设计了“强锚点支撑 + 跨链增量扩充 + 时序归序”的三重策略。首先，系统计算每条轨迹中所有步骤得分的几何平均值 $\left(\prod_{t=1}^{T_n} r_t^{(n)}\right)^{1/T_n}$，选出几何平均分最高的最佳完整轨迹 $n^\star$ 作为骨干锚点。随后，系统设置置信度绝对阈值 $\delta \in [0, 1]$，将全局池中所有满足 $r_t^{(n)} \ge \delta$ 的高分步骤全部保留，并强制完整保留锚点轨迹 $n^\star$ 的所有步骤作为结构兜底。所有保留步骤被显式标注上对应的置信度分数标签（例如 [c=0.93] ...），并依据其在原始轨迹中的时间步索引 $t$ 进行时序排序： $$E = \mathrm{Sort}\left(\{(t, r_t^{(n)}, s_t^{(n)})\}, \mathrm{key}=(t, -r_t^{(n)})\right)$$ 这种构造方式不仅保证了推导过程从问题到结论的宏观递进次序，还允许多个相互印证或补充推导的局部步骤在相同时间层共存，为后续决策提供了多视角的证据支撑。

4. 证据条件下的自回归求解重算：容错调和与答案闭环

缝合得到的证据列表 $E$ 并非完美的单一连贯推演，其中可能包含不同路径之间的表述冗余、细微前提差异或边缘矛盾。因此，系统并不直接从缝合列表中截取答案，而是引入轻量级自回归语言模型（例如 Qwen2.5-Math-7B-Instruct 或 Qwen2.5-Coder-7B）充当“陪审团与法官”。将原始问题 $x$ 与置信度前缀证据列表 $E$ 拼接作为提示词，指导求解器优先参考高置信度推导，识别并调和潜在的冲突分歧，最后利用贪婪或低温解码仅生成闭环推导与最终答案： $$\hat{y} = \arg\max_y p_\psi(y \mid x, E)$$ 由于大量复杂的中间探索与分支计算已由扩散模型与 PRM 在上游并行完成，自回归求解器仅需进行极短步数的聚焦生成（如直接输出 \boxed{...} 或紧凑的目标代码块），从而以极低的时序前向开销实现了从高噪声复合证据到精确确定性答案的高质量转化。

一个完整示例¶

以 MATH500 中一道包含代数化简与根式求解的多步几何题为例： 1. 扩散并行探索：扩散模型并行生成 4 条推理轨迹（$\tau^{(1)}$ 至 $\tau^{(4)}$）。其中 $\tau^{(1)}$ 准确建立了坐标系并推导出关键中间二次方程 $x^2 - 6x + 2 = 0$，但在求根步骤算错符号；$\tau^{(2)}$ 在早期几何建模上步骤简略，但在中后期推导出了正确的判别式 $\Delta = 28$ 及简化根；$\tau^{(3)}$ 与 $\tau^{(4)}$ 探索了三角代换路径，但在中途陷入发散。 2. PRM 打分与步骤池构建：Qwen-Math-PRM 评分显示，$\tau^{(1)}$ 的前 3 步置信度高达 0.96、0.94、0.91，第 4 步求根断崖式跌至 0.12；$\tau^{(2)}$ 的几何建模步骤得分适中（0.78），但代数求根部分得分达 0.93。 3. 锚点与置信度缝合：$\tau^{(1)}$ 的整体几何平均分最高，被选为锚点骨架 $n^\star$；置信度阈值 $\delta=0.7$ 筛选出 $\tau^{(1)}$ 的前 3 步以及 $\tau^{(2)}$ 的求根步骤。按时序拼接后形成证据列表：[c=0.96] 建立笛卡尔坐标系... [c=0.94] 列出动点轨迹方程... [c=0.91] 化简得 x^2-6x+2=0... [c=0.93] 判别式为 28，求得有效正根为 3 + \sqrt{7}。 4. 自回归重算：Qwen-Math 接收该证据后，迅速验证了二次方程与求根公式的一致性，忽略了 $\tau^{(1)}$ 骨架末尾的错误符号，直接输出终极答案 \boxed{3+\sqrt{7}}，仅耗费数十个自回归解码 token 即完成求解。

损失函数 / 训练策略¶

本文提出的 Diffusion Stitching 框架为完全免训练（Training-Free）的测试时即插即用管线。 - 无需额外参数微调：骨干探索模型（LLaDA / Dream / LLaDA 2.0）、评估模型（Qwen-Math-PRM / AceCoderRM / Skywork-PRM）及求解器（Qwen2.5-Math / Qwen2.5-Coder）均直接采用预训练权重，避免了复杂的强化学习对齐或两阶段联合微调的开销与不稳定性。 - 并行调度与分布式同步：将 $N$ 个独立的扩散轨迹采样均匀分发给 $G$ 张 GPU 设备（每卡执行 $\lceil N/G \rceil$ 次推演）。每张 GPU 独立运行扩散去噪并直接在本地执行单次 PRM 打分，最后仅通过一次轻量的 all_gather 通信汇总打分与步骤文本，构建全局候选池，在工程上具有线性的分布式算力扩展比。

实验关键数据¶

主实验¶

论文在严格的零样本（Zero-Shot）设置下，跨越数学推理（GSM8K、Minerva Math、MATH500）与代码生成（HumanEval、HumanEval+、MBPP、MBPP+）等 6 个权威基准进行了详尽评测。默认设置使用 8B 扩散骨干，采样 4 条独立轨迹，生成长度在 256–512 之间。

模型架构	参数规模	HumanEval (pass@1)	HumanEval+ (pass@1)	MBPP (pass@1)	MBPP+ (pass@1)	GSM8K (Acc)	Minerva Math (Acc)	六项平均 (Avg)
Qwen3（自回归基线）	4B	57.32	50.61	67.00	80.69	77.48	47.10	63.37
Qwen3（自回归基线）	8B	64.63	56.71	69.40	83.07	81.80	52.94	68.09
LLaDA（扩散基线）	8B	32.32	27.44	40.80	51.85	70.96	27.30	41.78
↳ 置信度采样（$\gamma=0.9$）	8B	40.24	37.80	41.89	55.82	78.77	36.54	48.51
Dream（扩散基线）	7B	54.88	49.39	56.80	74.60	77.18	39.60	58.74
↳ 置信度采样（$\gamma=0.9$）	7B	59.76	57.32	59.14	71.96	81.96	42.78	62.15
Block Diffusion（混合架构）	4B	56.10	51.22	54.60	69.84	82.87	47.02	60.27
TiDAR (Trust AR)	8B	55.49	52.44	65.40	79.63	79.83	50.58	63.90
TiDAR (Trust Diff)	8B	57.93	55.49	65.40	80.95	80.44	51.64	65.31
Ours（Diffusion Stitching）	8B	70.37	64.02	74.61	81.75	91.51	53.20	72.38

消融实验¶

1. 聚合策略与求解器角色消融（GSM8K 与 MATH500）¶

对比不同轨迹聚合方案（无求解器、全轨迹拼接、单一最优轨迹、纯阈值缝合、带锚点的完整缝合）：

聚合与求解配置	是否使用 PRM	是否使用求解器	GSM8K (Acc %)	MATH500 (Acc %)	说明
LLaDA 单轨迹基线	否	否	78.8	37.6	单一扩散轨迹贪婪生成
多数投票（Majority Vote, N=4）	否	否	85.1	42.0	传统自洽性答案投票
全量轨迹拼接（All CoT）	否	是	86.6	46.0	将 4 条 CoT 全部无差别拼接送入求解器
最优轨迹条件化（Best CoT）	是	是	90.1	49.2	仅选几何平均分最高的一条完整轨迹
纯置信度步骤缝合（Above Confidence）	是	是	90.4	52.0	仅保留所有得分超过 $\delta$ 的步骤，无锚点
完整模型（Best CoT + Above Confidence）	是	是	91.5	54.2	强锚点兜底 + 跨轨迹高分步骤时序缝合

2. 自回归求解器 vs 扩散模型重算消融（Table 3）¶

在相同缝合证据下，评估采用自回归解码与复用扩散去噪作为终极求解器的效果与开销：

缝合轨迹数	终极求解器类型	GSM8K (Acc %)	GSM8K 前向步数	MATH500 (Acc %)	MATH500 前向步数
4 traces	自回归求解器（AR）	91.5	96.9	54.2	135.5
4 traces	扩散求解器（dLLM, $\gamma=0.7$）	90.1	106.3	50.2	132.6
4 traces	扩散求解器（dLLM, $\gamma=0.9$）	90.3	141.6	52.0	186.5
8 traces	自回归求解器（AR）	91.4	110.8	55.0	141.1
8 traces	扩散求解器（dLLM, $\gamma=0.7$）	91.4	115.8	52.6	138.6
8 traces	扩散求解器（dLLM, $\gamma=0.9$）	92.5	151.8	54.4	198.2

3. 缝合设计稳健性消融（Table 9）¶

分析轨迹打分算子、锚点选择、步骤排序方式以及分句切分策略对系统鲁棒性的影响：

维度配置项	具体设定	GSM8K (Acc %)	MATH500 (Acc %)	MBPP (pass@1 %)	结论与分析
轨迹评分算子	几何平均（Geometric Mean，本文）	91.5	54.2	74.6	与算术平均表现极相近，评分算子本身并不敏感
轨迹评分算子	算术平均（Arithmetic Mean）	91.6	54.4	73.6	差异在 0.2% 以内
锚点选择策略	强锚点（Strong Anchor，本文）	91.5	54.2	74.6	保留最佳完整轨迹提供了至关重要的全局连贯性
锚点选择策略	最弱锚点（Weakest Anchor）	89.7	48.8	73.1	在难题（MATH500）上暴跌 5.4%
锚点选择策略	随机锚点（Random Anchor）	90.2	50.4	72.6	掉点明显，验证了骨干质量的关键性
锚点选择策略	无锚点（No Anchor）	90.4	52.0	73.6	证明仅有零散步骤会损害上下文连贯性
步骤排序与切分	规范时序排序（Canonical Order，本文）	91.5	54.2	74.6	依推理进度时序排布能引导求解器顺畅推导
步骤排序与切分	置信度降序排序（Confidence Order）	90.8	51.4	73.2	打乱时序使得推导因果倒置，MATH500 下降 2.8%
步骤排序与切分	随机乱序（Random Step Order）	90.2	53.0	73.0	显著破坏推理上下文因果链
步骤排序与切分	替代切分粒度（Alternative Segmentation）	83.6	39.8	71.5	粒度失配导致逻辑断裂，MATH500 暴跌 14.4%

关键发现¶

步骤级重组的乘积效应在难题上尤为显著：在 GSM8K 上，完整缝合相比单一最优轨迹（Best CoT）提升了 1.4%，但在难度更高、推导链更长、容错率更低的 MATH500 上，缝合直接带来了 5.0% 的绝对精度飞跃（49.2% $\rightarrow$ 54.2%）。这充分证明了在复杂推理中，不同失败轨迹往往共享了互补的局部正确子推导，步骤级缝合成功打捞了这些沉没计算成本。
解耦架构完胜端到端混合系统：解耦的 Diffusion Stitching 在平均表现上以 72.38% 显著超越了此前最先进的端到端混合架构 TiDAR（Trust Diff 65.31%），不仅在数学与代码上双双大幅领跑，而且在同等精度下减少了高达 9.85 倍的串行前向传播次数。
扩散探索容忍低置信度，大幅压降推理延迟：扩散采样阶段无需追求高置信度精细去噪。消融显示，将扩散采样阈值从 $\gamma=0.8$ 降低至 $\gamma=0.6$，平均前向推演步数由 158.5 步减少至 119.8 步（节约 24.4% 算力），而 MATH500 准确率仅微幅波动（54.4% $\rightarrow$ 52.0%）。下游 PRM 与自回归求解器展现出强大的去噪、容错与重构韧性。
失败归因指向探索多样性而非 PRM 误判：针对 50 例 MATH500 预测失败的人工干预实验发现，仅有 1 例（2.0%）归因于 PRM 的假阴性漏判，多达 40 例（80.0%）即便补充单一缺失推导步也无法修复。这表明核心瓶颈在于扩散模型是否能在上游产生具有突破性的正确子推导，而非下游评估模型的微小扰动。

亮点与洞察¶

将扩散探索与因果求解的专业化分工发挥到极致：掩码扩散模型擅长全局并行探索但缺乏严格单向因果闭环能力，自回归模型长于逻辑归纳但串行扩展极耗显存与延迟。本文通过 PRM 建立两者的连接桥梁，既享受了扩散模型的低延迟多样性搜索，又获得了自回归模型的确定性严密逻辑。
置信度标注前缀引导的软约束条件化机制：不强求在离线阶段对步骤进行完美的去重与图结构对齐，而是将置信度标量显式转化为自然语言标记（如 [c=0.93]），巧妙利用了大语言模型自身的上下文注意力偏置，使其在生成答案时自适应倾向于采信高分证据并调和分歧。
测试时扩展的 Pareto 前沿推进：相比于在自回归模型上进行昂贵的 Best-of-$N$ 或树搜索，该方案将探索成本卸载到高吞吐的扩散采样上，在 4 块 H200 硬件下实现了 1.8 倍的端到端端时延缩减，为高吞吐低延迟的 Agent 推理提供了极具落地前景的工程范式。

局限与展望¶

瞬时峰值显存与多卡并行依赖：尽管独立并行的扩散采样显著压缩了端到端物理等待时间，但多路轨迹的同时展开与 PRM/自回归求解器的多模型加载显著拉高了显存开销。在 GSM8K 与 MATH500 上，Qwen-PRM 配置的峰值显存达到了 42.8 GB，限制了其在边缘设备或单卡低显存环境中的部署。
上游扩散探索的有效覆盖率瓶颈：实验证明绝大部分失败归因于扩散轨迹未生成关键推演。当前使用的通用 dLLM 缺乏针对特定数学/代码领域深层符号逻辑的先验引导，若初始采样严重坍缩至相同的错误模式，下游的置信度缝合将“巧妇难为无米之炊”。
未来的扩展空间：未来可探索在扩散去噪过程中直接融入轻量级引导（Guidance），在扩散早期进行有向剪枝；同时可引入 KV 缓存加速技术（如 Fast-dLLM）进一步挖掘扩散采样的极限吞吐潜力。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将扩散模型的非因果多样性并行采样与 PRM 步骤级评分及自回归重算相结合，打破了传统测试时搜索对自回归串行展开与粗粒度轨迹聚合的固有依赖。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 6 大数学与代码基准，包含详细的 Pareto 前沿分析、模块化消融、模型通用性验证及深度失败归因，实验数据详实客观。
写作质量: ⭐⭐⭐⭐⭐ 论文动机清晰明确，方法设计逻辑层层推进，图表直观且附录算法细节与超参完备。
价值: ⭐⭐⭐⭐⭐ 为解决高难度推理任务中“测试时计算爆炸”与“低延迟落地需求”之间的核心冲突开辟了极富启发性的全新技术路径。