近年來，大型語言模型能力的提升焦點，正從「訓練時擴展」轉向「推理時擴展」。從 Best-of-N、Self-Consistency 到更複雜的搜尋與驗證框架，Test-Time Scaling 已成為提升大型模型複雜推理能力的重要範式。
然而，一個長期被忽略的問題是：這些方法大多預設模型是以自迴歸的方式生成。
對於離散擴散語言模型（Discrete Diffusion Language Models, dLLMs）而言，情況完全不同。dLLM 並非由左至右逐個 token 生成，而是從被遮罩的序列出發，透過多步去噪，逐步恢復出完整的答案。這種並行、非自迴歸的生成方式，天生具備全域雙向上下文，也更適合進行規劃與自我修正；但與此同時，傳統為自迴歸模型設計的樹搜尋、過程獎勵模型和 Best-of-N 推理，並不能直接且高效率地套用。
針對此問題，論文提出了 PRISM：Pruning, Remasking, and Integrated Self-verification Method，一個專為離散擴散語言模型設計的高效率 Test-Time Scaling 框架。其核心目標非常明確：不是簡單地讓模型「多跑幾次」，而是在去噪過程中，辨識出更有潛力的軌跡，動態裁剪、局部分支，並利用模型自身完成輕量級驗證，從而在較低的推理預算下，達到甚至超越 Best-of-N 的效果。
論文標題：Prism: Efficient Test-Time Scaling via Hierarchical Search and Self-Verification for Discrete Diffusion Language Models
arXiv 位址：https://arxiv.org/abs/2602.01842
傳統 Best-of-N 代價高昂，PRISM 重新設計 dLLM 的推理搜尋
對於 dLLM 來說，樸素的 Best-of-N 代價非常直接：如果取樣 N 條軌跡、每條軌跡需要 T 步去噪，那麼總函數呼叫次數就是 O(NT)。這意味著所有候選答案都必須完整跑完，即便其中許多軌跡在中途已經明顯品質不佳，也依然會消耗掉完整預算。
PRISM 的關鍵思路，是把推理過程拆解成三個階段：早期隨機探索、中期漸進裁剪和後期精修。
在高噪聲階段，模型輸出仍然不穩定，因此 PRISM 維持較寬的候選集合以保留多樣性；在早中期的去噪時間窗口，當答案的「邏輯骨架」開始形成時，PRISM 使用自我驗證訊號來裁剪低品質軌跡，並將運算資源重新分配給更有前景的候選；最終，只保留較少數量的軌跡繼續完成精修。論文中將此過程稱為 Hierarchical Trajectory Search（HTS）。
這種設計使得 PRISM 的實際複雜度接近 O(N + KT)，其中 K 是最終保留的較小候選寬度。相比傳統 Best-of-N 的 O(NT)，這相當於把「所有路線都跑到底」改成了「先廣泛探索，再集中火力」。
不是重新開始，而是在低信心 token 上局部分支
PRISM 的第二個關鍵組件是 Local Branching via Partial Remasking。直觀來說，模型在中期去噪時，已經形成了一部分高信心 token，這些 token 往往對應答案的穩定結構或邏輯骨架；與此同時，低信心 token 則可能對應不確定的推理細節、實作方式或局部表達。
PRISM 不會粗暴地丟棄整條軌跡重新取樣，而是保留高信心的部分，只對低信心位置進行重新遮罩，然後從這些局部變化中生成新的分支。這樣做的好處是，它既保留了已有的高品質結構，又能繼續探索不同的細節實作，避免過早收斂到單一路徑。論文圖 2 對這個過程做了直觀展示：在漸進裁剪階段，PRISM 會圍繞高分軌跡進行局部分支，並透過部分重遮罩生成新的候選。
不再依賴外部驗證器：模型自己給自己打分
許多 Test-Time Scaling 方法依賴額外的 reward model 或驗證器來判斷候選答案的品質。但這會帶來顯著的系統開銷：部署時需要額外載入一個模型，記憶體、延遲和工程複雜度都會上升。
PRISM 提出了 Self-Verified Feedback（SVF）：直接複用同一個 dLLM 作為二元分類驗證器。具體來說，模型先根據中間去噪狀態生成一個完整的候選答案，然後構造一個 Yes/No 驗證提示，讓模型判斷該答案是否可能正確。PRISM 將 Yes 和 No 的 logits 轉換為一個二元正規化分數，用於軌跡排序、裁剪和最終選擇。
此設計的價值在於，它把驗證從「需要額外模型」變成了「同一模型的一次輕量級自我檢查」。論文進一步指出，SVF 的呼叫次數相較於去噪 NFE 非常少，實驗中通常低於總 NFE 的 10%，因此能在較低的額外開銷下，提供有效的搜尋訊號。
實驗：在數學推理和程式碼生成上實現顯著性價比提升
論文在四個基準上評估 PRISM：數學推理任務 GSM8K、MATH-500，以及程式碼生成任務 HumanEval、MBPP。實驗涵蓋三個離散擴散語言模型：LLaDA-8B-Instruct、Dream-7B-Instruct 和 LLaDA-2.0-mini。
在 LLaDA-8B-Instruct 上，PRISM（K=8）將 GSM8K 從 67.58% 提升到 85.30%，將 MATH-500 從 26.40% 提升到 42.80%；在程式碼任務上，HumanEval 提升了 24.39 個百分點，MBPP 提升了 16.40 個百分點。更重要的是，這些提升並不是透過線性增加 Best-of-N 運算量得來的：例如在 GSM8K 上，PRISM 用 1048 NFE 達到 85.30%，而 Best-of-16 需要 4096 NFE 才達到 87.50%，實現了超過 4 倍的去噪運算節省。
在論文圖 1 中，PRISM 相較 Best-of-N，在多個任務上展現出更優異的效能—運算曲線：在可比的準確率下，GSM8K、MATH500、HumanEval、MBPP 分別呈現 2.9 倍、6.5 倍、1.8 倍、1.7 倍的速度優勢。
論文還與其他推理期擴展方法進行了比較。在 TruthfulQA 上，PRISM 的 ROUGE-1/2/L 達到 31.8/35.5/31.9，推理時間為 1048.0 秒；相比之下，LLaDA-ReMDM 為 29.5/31.8/29.5，推理時間為 1354.8 秒。這表明 PRISM 不僅能提升任務效能，也能維持更好的推理效率。
在外部驗證器的對比中，SVF 在 GSM8K 上達到 85.30%，只需載入原本的 8B 模型；雖然 Qwen3-8B 驗證器可達到 87.35%，但需要額外載入模型，總參數量達到 16B。論文認為，SVF 的優勢並不在於絕對取代所有外部驗證器，而在於它提供了一條更輕量、更易於部署的 dLLM 推理擴展路徑。
意義：為非自迴歸語言模型打開推理期擴展路線
PRISM 的核心貢獻，並不是簡單地提出一個新的搜尋啟發式方法，而是重新定義了 dLLM 上的 Test-Time Scaling 應該如何發生。
對於自迴歸模型，推理搜尋通常圍繞「前綴」展開；而對於離散擴散模型，中間狀態是部分遮罩的全域序列，傳統的前綴式過程獎勵和樹搜尋並不天然適用。PRISM 將搜尋、裁剪、局部分支和自我驗證，都重新放回 dLLM 的去噪動力學中：在結構形成階段集中分配預算，在低信心區域探索替代表達，在無需額外模型的情況下完成驗證。
這意味著，dLLM 不再只是「並行生成更快速」的替代範式，也可能成為一種適合推理、規劃和自我修正的新型語言模型架構。隨著 LLaDA、Dream、Mercury、Gemini Diffusion 等模型推動離散擴散語言模型走向更大規模，PRISM 展示了一條重要方向：讓非自迴歸模型也能像當前主流 LLM 一樣，透過推理期運算，持續獲得能力提升。
從這個角度來看，PRISM 不只是一個更節省算力的 Best-of-N 替代方案，而是離散擴散語言模型邁向高效推理系統的一塊關鍵拼圖。
作者簡介
本文由 Jinbin Bai 等研究者完成。作者團隊長期關注 discrete diffusion 與 masked generative modeling 等新一代生成範式，研究方向涵蓋高解析度文生圖、統一多模態生成、離散擴散模型的偏好對齊與推理最佳化，以及可互動世界模型等。
此前，團隊曾提出 Meissonic [1]，探索 masked generative transformer 在高解析度文字到圖像生成中的潛力；隨後進一步提出 Muddit [2]，將離散擴散建模從圖像生成推進到更統一的多模態生成框架。此次入選 ICML 2026 的 PRISM，則將此一研究脈絡進一步延伸到推理階段，關注如何透過層次化搜尋、自我驗證回饋和局部 remasking，讓離散擴散模型在無需額外驗證器的情況下，實現高效 Test-Time Scaling。
[1] Meissonic: Revitalizing Masked Generative Transformers for Efficient High-Resolution Text-to-Image Synthesis, ICLR 2025, https://arxiv.org/abs/2410.08261
[2] Muddit: Liberating Generation Beyond Text-to-Image with a Unified Discrete Diffusion Model, ICLR 2026, https://arxiv.org/abs/2505.23606
© THE END
轉載請聯繫本公眾號獲得授權
投稿或尋求報導：liyazhou@jiqizhixin.com