PRISM: Efficient Test-Time Scaling via Hierarchical Search and Self-Verification for Discrete Diffusion Language Models

1 · 动机

1.1 历史脉络:dLLM 撞上和 AR 一样的 accuracy 墙

2025 年是 dLLM 从玩具走向 reasoning baseline 的一年——LLaDA-8B、Dream-7B、LLaDA-2.0、Seed Diffusion、Mercury、Gemini Diffusion 都把"masked diffusion + bidirectional 上下文"做到了和同尺寸 AR 模型可比的精度。但 dLLM 一旦真的拿去做数学题、写代码,就会很快撞到 AR 模型熟悉的 accuracy ceiling:N=1 single trajectory 在 GSM8K 上 LLaDA-8B 也就 67.58%,MATH500 上只有 26.40%——继续训也榨不出多少。

AR 世界里这个问题的标准答案是 test-time scaling (TTS):CoT、self-consistency (Wang 2022)、best-of-N (Brown 2024)、o1-style RL + 长 thinking、tree-of-thoughts / MCTS (Yao 2023)。本质都是花推理时算力换正确率。但这些算法的设计前提是 left-to-right、prefix-conditional 的解码——比如 PRM 要给"良构 prefix"打分,MCTS 节点是逐 token 增量的。

dLLM 的 generation 长得完全不一样:

• 不是 prefix → token,而是 [MASK]ᴸ → 全序列并行去噪 T 步;
• 中间状态是 partially-masked sequence,不是 well-formed prefix;
• 每步同时预测所有位置的 logits,具备天然的 z0-prediction head——任何中间步都能 argmax 出一个 candidate full answer。

所以"把 best-of-N 直接搬过来"是 trivial 的(跑 N 条独立 trajectory + majority vote 就行),问题是 每条都要 T 步去噪——LLaDA-8B 默认 256 token / 32 steps per 32-token block,N=16 直接 16× 算力。论文 Figure 1 给了硬数据:GSM8K 上 best-of-16 要 4096 NFE 才到 87.50%,PRISM 用 1048 NFE 就能到 85.30%——用 1/4 算力换 -2.2 pp。这就是 PRISM 想解决的"dLLM-flavor 的 TTS"。

1.2 别的方案为什么不够 (对比表)

1.3 为什么这事不平凡

三个 dLLM 独有的复杂性,直接照搬 AR TTS 就会踩坑:

1. "早期决定全局,晚期决定细节"的非对称性。论文附录 A 测了 Dream-7B 在四个 benchmark 上 token-averaged 预测熵 H(t):熵在 很早期急剧下降,但 mid-trajectory 仍有显著 bumps 和 inter-trajectory variance,直到末期才共同收敛。意味着:整体框架(logic skeleton)在 mid-denoising 已经分化,但局部 token 还没定;在熵已经"分化但未塌缩"的窗口里剪枝最有信息量。
2. Verifier 必须在 partially-masked 状态上工作。AR PRM 训练数据是 prefix,给 [MASK] 包裹的中间态打分会失准。PRISM 的 trick 是不直接评 partially-masked z_t,而是先 argmax 出 ẑ₀(z₀-prediction head 一直输出全序列),再让同一个 dLLM 用 Yes/No prompt 判断这个 hypothetical 答案——把 verify 问题转成"dLLM 擅长的 text understanding"任务。
3. 多 trajectory 之间的"branching"语义不像 AR 那样清晰。AR 树里一个节点 = 一个 prefix,分裂就是 sample 下一个 token。dLLM 一条 trajectory 是整条 256-token 的 partially-denoised 序列,怎么"在它附近"产生 diverse children?如果完全重 sample 就退化成 best-of-N;如果完全 share 就没多样性。PRISM 的答案是 partial remasking:保留高置信度的 logic skeleton,只重 mask 低置信度位置——既复用算力又保多样性。

这三条加在一起决定了 PRISM 不是"把 MCTS 抄过来",而是要重新设计 search 拓扑、verifier 形式、扩展操作三件事,都对齐 dLLM 的 denoising 动力学。

2 · 背景速查

3 · 总览:PRISM 的三个零件

论文起的标题首字母拼出 PRISM = Pruning + Remasking + Integrated Self-verification Method。但实际是三个解耦的机制:

对应 paper Algorithm 1:

• Stage I (Stochastic Exploration, T_p < t ≤ T):N 条 trajectory 同步 denoise,不剪枝。此时熵高、SVF 不可靠,先广撒。
• Stage II (Progressive Thinning, T_r < t ≤ T_p):活跃宽度按 W_t = max(N·d^{-(T_p-t)}, K) 几何衰减,每 i 步触发一次 (SVF 评分 → 取 top-S → 各自 partial-remask 出 b_t 个 children) 操作。
• Stage III (Final Refinement, 1 ≤ t ≤ T_r):剩 K 条独立 denoise 到末端,加 confidence threshold τ=0.95 早提交 + \boxed{} marker 早终止,最后 majority voting 出答案。

4 · SVF — 让 dLLM 自己当 verifier

外部 reward model 是 AR TTS 的标配,但对 dLLM 有两个硬伤:(1) PRM 训练在 prefix 数据上,partially-masked z_t 不在它见过的分布里;(2) 跑一个独立 7-8B verifier 直接把 8B dLLM 的 deployment memory 翻倍。SVF 的 trick 是:不评中间状态,只评 z₀-prediction head 给出的完整 hypothetical 答案。

其中 s_Yes / s_No 分别取 dLLM 在 verification prompt 末位置上,Yes 类 token IDs(如 "Yes", "yes", "YES")的 max logit 和 No 类的 max logit。本质是用 dLLM 自身的 instruction-following 能力做二分类 calibration。

Worked example

问题: If Lisa has 7 boxes of 8 chocolates and gives away 2 boxes, how many chocolates does she have left?

假设 t=80 时(T=128, 已经走了 48 步),有两条 trajectory 的 ẑ₀ 都还带几个 [MASK],argmax 给出:

• traj A 的 ẑ₀ = "7×8=56, give away 2×8=16, left = 56-16 = \boxed{40}"
• traj B 的 ẑ₀ = "7-2=5 boxes left, 5×[M]=\boxed{40}"

两条最终答案都 40 但中间论证完全不同。SVF prompt 把两个 ẑ₀ 各喂回 dLLM 一次:

• traj A: s_Yes=4.1, s_No=0.8 → Φ_SVF = e^4.1/(e^4.1+e^0.8) ≈ 0.964
• traj B: s_Yes=1.9, s_No=2.3 → Φ_SVF = e^1.9/(e^1.9+e^2.3) ≈ 0.401 (B 的论证里 5×[M] 是 garbage,模型自己也判 No)

之后 HTS top-S=2 选中包含 A 的两条作为 seed,B 被淘汰。

SVF 可靠吗?

Table 5 给了诊断:在 GSM8K / MATH-500 上用 Qwen3-235B-A22B-Instruct 做 proxy oracle,SVF 的 F1 在 0.63–0.73 之间(显著优于随机),而且 ECE 从 first SVF (0.20-0.33) 单调降到 last SVF (0.09-0.29)——denoising 越往后,SVF 越自信也越 calibrated。论文把这归因于"hypothetical 答案语义稳定性提升"。

反向论证

如果不用 SVF 而用 random pruning,ablation (Table 4) 显示 K=2/4/8 各掉 3.5/4.3/3.7 pp——SVF 是主要 search 增益来源。如果用 Qwen3-8B 当外部 verifier,GSM8K 能再涨到 87.35(SVF 85.30),但要多 load 一个 8B 模型,总 params 16B vs 8B,在 40GB A100 上直接 OOM。SVF 用同模型换 -2 pp 是值得的成本。

5 · HTS — 三段式 + geometric decay

HTS 的本质是一个 budget allocation 调度:T 个 denoising 步、总 NFE 预算固定,要决定每步保留多少活跃 trajectory。论文画出的关键直觉是 entropy 曲线(附录 A)的形状:

几何衰减的算术

取默认 N=16, K=8, d=1.8, T=32 (LLaDA 每块), W=[0.1, 0.6] → T_p=⌈0.6·32⌉=20, T_r=⌈0.1·32⌉=4。Stage II 跨 16 步,但每 i=3 步才触发剪枝。活跃宽度 W_t = max(16·1.8^{-(20-t)}, 8)。代入:

这里 d=1.8 其实把 N→K 的收缩压在了前两次 prune 内。论文实验也确认 d 太小(慢收缩) / 太大(激进)都差,1.8 是 sweet spot。

反向论证: 不做三段式会怎样

• 不做 Stage I 直接剪枝 → 在 t≈T 时熵未分化,SVF 给出几乎随机的 Yes 分数,等同 random prune(论文 ablation 显示掉 3+ pp)。
• 不做 Stage III 直接在 Stage II 末选答案 → 部分位置还是 [MASK],z₀-pred 的尾部 token 不稳定。
• 整段 uniform 剪枝(无 window)→ 算力浪费在已经塌缩的末端,论文 Tab.6 显示 W=[0.0, 0.6] 或 W=[0.3, 0.6] 都比 [0.1, 0.6] 差。

6 · Local Branching via Partial Remasking

剪枝以后剩下 S 个 survivor seeds,每个 seed 需要扩出 b_t = ⌈W_{t-1}/S⌉ 个 children 才能填满下一步的 W_{t-1} 池。最朴素的做法是把 seed 整个 mask 掉重 sample——但那就丢掉了 seed 已经形成的 logic skeleton。PRISM 的 partial remasking 是这样的:

为什么不直接重 sample

对比三个选项的成本/多样性:

Table 4 ablation: 去掉 partial remask 在 K=2/4/8 各掉 1.13/1.89/1.21 pp——比 SVF 的贡献小,但和 SVF 互补。

7 · 复杂度分析:为什么是 O(N + KT)

三项分别对应:Stage I 的 warm-up(N 条 × (T-T_p) 步),Stage II 的几何衰减总和,Stage III 的 K 条精修。因为

• warm-up 短:T - T_p ≈ (1 - w_max)·T = 0.4·T,而且 Σ|P_t| 在几何衰减下被 K·(几步) 主导;
• Stage II 总和 ≈ K · (T_p - T_r) 当 d 够大;

所以 C_HTS ≈ N·0.4T + K·0.5T + K·0.1T ≈ O(N + KT)(原文化简)。对比 best-of-N 的 O(NT):当 N=16, K=8, T=32 时,best-of-N = 512,HTS ≈ 16 + 256 = 272,理论 ~1.9× speedup,实测 GSM8K 上 4096 → 1048 即 ~3.9×(还吃到 Stage III 的早终止红利)。

8 · Worked Example — 一个 GSM8K 问题端到端

设置:LLaDA-8B, block size M=32, T=32, N=16, S=4, K=8, d=1.8, W=[0.1, 0.6], i=3, τ=0.95。一条 256-token answer 共 8 个 block,这里只追 block 1。

Prompt: Lisa has 7 boxes of 8 chocolates and gives away 2 boxes. How many chocolates left?

关键观察:(1) Stage I 的 192 NFE 占了一半算力——这是必要 warm-up,SVF 在 t>20 时基本不可靠;(2) Stage II 的 SVF 调用只有约 24 次,远小于 256 NFE 的 denoise 成本;(3) Stage III 早终止让最后 8 条没各跑完 4 步,实际 24 NFE 而非 32。

9 · 实验关键数字

9.1 主表 (LLaDA-8B Instruct)

读法:

• GSM8K K=8 用 1048 NFE 达到 85.30,vs best-of-16 的 4096 NFE 87.50——1/4 算力换 -2.2 pp。
• MATH500 K=8 用 1304 NFE 拿到 42.80,反超 best-of-16 的 38.00,4.80 pp 提升 + 3.1× 算力节省。
• HumanEval 是 PRISM 最弱的场景:K=8 (2480 NFE) 79.27 vs best-of-16 (8192) 82.32,相差 3.05 pp 但 3.3× 加速。代码生成的 inter-trajectory variance 更大(附录 entropy 图也证实),SVF 在 code 上 calibration 略差。

9.2 SVF 可靠性 (Table 5)

关键:ECE 在 denoising 推进中单调下降(.330 → .091),意味着 SVF 越往后越 calibrated。这是 PRISM 把 SVF 调用集中在 Stage II 中后段而非全程的理论依据。F1 不算特别高,但 PRISM 是 ranking(只要相对序对就够),不依赖绝对 calibration。

9.3 Ablation (Table 4, GSM8K LLaDA-8B)

SVF 是主贡献(-3 到 -4 pp 影响),partial remask 是次贡献(-1 到 -2 pp)。两者互补且都为正。

9.4 vs 外部 verifier (Table 3)

用更强的外部 verifier 能再涨 2 pp,代价是 2× deployment memory + 单 A100 (40GB) 装不下。SVF 是"自包含"的工程选择。

10 · 与同类工作对比

简而言之:PRISM 是 dLLM 一侧首个把 "search 拓扑 + verifier 形式" 都对齐 denoising entropy 动力学的工作。它的差异化不在 search 本身(MCTS / SMC 这些都已经存在),而在 把这些 search 工具调度到对的 timestep window(W=[0.1, 0.6])。

11 · 局限 / 个人 take / 待验证

• SVF 在 code 上的可靠性。HumanEval 上 K=8 也只追到 best-of-16 的 -3.05 pp(其他 benchmark 反超或追平)。Dream-7B / LLaDA-2.0 在代码上 entropy 曲线"高 plateau + 大 variance",SVF 的 Yes/No prompt 在代码任务上可能没那么 calibrated。Table 5 没给 HumanEval / MBPP 的 F1。
• 没有探究 SVF prompt 的 sensitivity。论文只给"Is this answer correct? Yes/No",但 prompt 设计对 small LLM 的二分类影响极大,不同 dLLM 可能需要不同 prompt。Reproducibility 略担忧。
• 超参 W=[0.1, 0.6] 是 LLaDA-8B 上调出来的。不同 dLLM 的 entropy 曲线形态不同(论文附录就显示 GSM8K 和 MATH500 在 Dream 上 shape 不同),W 是否能用 entropy curve 自动选?
• SVF + partial remask 都需要 dLLM 有较好的 instruction-following 和 z₀-prediction 稳定性。对于较弱的 dLLM(比如 base 模型 + lightweight 微调)是否仍有效未验证。
• 与 RL 训练的关系不清。论文是 pure training-free 视角,但 SVF score 自然可以做 RL reward(像 self-rewarding LM),可以想象 PRISM + RL on dLLM 的组合。
• Majority voting 的 boundary。K=8 时如果 8 个答案都不一致,fallback 行为论文未说明。对于 open-ended 任务这是个空白。

待验证的问题

1. SVF 的 F1 在 code 上是否系统性低于 math?如果是,是否因为 code 的二元判断更难(语法对但语义错)?
2. 把 PRISM 应用到 long-context reasoning(参考 LongLLaDA 的扩展)时,256 token answer 之外的场景是否仍工作?
3. 如果把 Stage I 的 N 调到更大(N=32, 64),性能是否 saturate?Entropy 曲线告诉我们,Stage I 的 N 越大,Stage II 第一次剪枝的 SVF 信号越有信息量,但 NFE 成本也线性增长——sweet spot 在哪?
4. PRISM 能否和 Fast-dLLM(KV cache + 并行 decoding)叠加?这是 NFE 之外的 wall-clock 优化方向,理论上正交。
5. partial remask 的"低置信位"是基于 token entropy,如果换成 mutual information / DLM 自己的 noise schedule,效果会不会更好?

12 · Memory points