Why Deterministic PRM Guidance Underperforms in Discrete Diffusion Reasoning
Yan Zhan, Shaobo Liu, Zhijun Gao
NeurIPS 2026
cs.AI
2026-09-28
同等算力下,dLLM 的 PRM 引导在 GSM8K 输给采样加 ORM 重排近 10 个点,根因是早期剪枝毁池子、终选又错人。
离散扩散语言模型(dLLM)不走自回归逐 token 生成的路子,它维护一条带 mask 的完整序列,逐步去噪出答案。每一步的中间状态(论文叫 snapshot)已经露出算式和答案片段,看上去天然适合 PRM(给推理中间步骤打分的模型)引导:在答案定稿前识别好苗子,把测试时算力花在刀刃上。PRM 在自回归数学推理上确实成过事,但那套方法假设证据按前缀顺序到达;dLLM 的 snapshot 是散落的可见位置,后面的 token 露出来时前面的可能还蒙着。Dream、LLaDA 一线的工作都在改生成器和采样器,没人对过账:同等算力下,PRM 引导到底打不打得过最朴素的独立采样加 ORM 重排。这篇的答案是打不过,差得不少。
核心是一套同单位记账协议,所有方法都按「对 7B 模型做一次前向」计费。vanilla 采样记 128 次前向(去噪 128 步);ORM Rerank@N 记 128N 次去噪加 N 次打分;PRM Guided 每 64 步分叉出 K 个候选并行去噪、用 PRM 打分、只留最高的一个,K=8 时 1040 次前向对 ORM 的 1032 次,PRM 一侧还多拿 0.8%。换算 wall-clock 同样成立:一次 ORM 打分约 0.3 秒,一条完整去噪轨迹 20.7 秒。答案抽取全程用严格正则,避开 lm-eval 式取末位数字带来的 11 个点虚高。
PRM 是冻结的 Dream-7B backbone 加 LoRA 加奖励头,用最终答案对错的二值标签训练,标签挂到轨迹上所有中间状态(cross-mask 训练)。对照组有只训终态的 ORM、只训终态的 PRM、同预算的 SMC 采样器,以及 LLaDA-8B 复验。诊断拆两段:候选池被剪枝毁掉多少(看 Oracle 上限),以及终选本身选得多差(同一个池子换打分器重排)。
Dream-v0-Instruct-7B,GSM8K 全量 1319 题:
| 方法 | N/K=8 | N/K=32 |
| Vanilla(单样本) | 43.14% | – |
| Majority@N | 60.05% | 67.63% |
| PRM Guided | 65.18% | 70.02% |
| ORM Rerank | 75.13% | 82.71% |
| Oracle 上限 | 81.05% | 91.13% |
差 9.95 和 12.69 个点,且 ORM Rerank@8 赢下所有测过的 PRM Guided 预算,包括花四倍算力的 K=32。MATH 上 30.65% 对 20.80%,MBPP 上 63.04% 对 50.88%。wall-clock 对账:ORM Rerank@6 约 126 秒拿 72.40%,PRM Guided K=8 用 156 秒只有 65.18%。
两层失灵各有证据。其一,信号随 mask 比例衰减:PRM 的 ROC-AUC 从接近解码完成的 0.77 掉到接近全 mask 的 0.54,论文给出互信息上界,说明任何打分器都躲不开;改用 8 次新 rollout 的重标标签再训,AUC 只升 0.011。top-1 剪枝在最早打分点一刀下去,46% 的情况把日后能到达正确答案的血统全剪没,池子的 Oracle 上限从 81.05% 掉到 67.30%,每题唯一答案数从 4.31 缩到 1.75。其二,终选错人:SMC 在同预算下把 Oracle 修复到 77.89%,用 PRM 选却停在 65.48%;同一个池子换 final-state PRM 重排,75.40%,追平 ORM。MBPP 把两层分开:PRM 给写完的程序重排有 65.47%,与 ORM 打平,一引导去噪就掉到 50.88%,损失全部来自引导本身。读出上,mean pooling 的因果 PRM 终态 ROC-AUC 只有 0.61,换 last-token pooling 升到 0.73;LLaDA 上双向 PRM 引导 31.64%,因果版 22.25%;48% 的题出现错轨被自信排在对轨前面的倒挂。
对做 dLLM 测试时扩展的人,这是最该补的一类对照:打分器调用算进预算后,不少「引导有提升」的说法要重新过账。操作建议也直接:晚引导、多留候选(终态留 top-4 把剪灭风险从 20% 压到 6%,每 64 步分支比每 16 步好约 7 个点),终选交给训过终态的 verifier。池子伤害加终选误差的拆法不限于扩散模型,自回归的 PRM beam search 同样套得上。快照语料和评测工具已开源。
作者列了八条威胁,要紧的几条:证据只覆盖数学和代码,创意生成没测;核心结论绑定确定性 top-1 这个最弱的引导形态,energy-tilted 采样、自适应分支、随机 beam search 都没测,对「PRM 引导」的否定有边界;ORM 优势依赖任务内训练,GSM8K 训的 ORM 搬到 MATH500 只剩 7% 上下,连 Majority@32 的 17.20% 都不如;LLaDA 上只复验了双向对因果的优势,没跑 ORM 对照。