![]()
近年來,大模型能力提升的焦點正在從「訓練時擴展」轉向「推理時擴展」。從 Best-of-N、Self-Consistency 到更復雜的搜索與驗證框架,Test-Time Scaling 已經成為提升大模型復雜推理能力的重要范式。
然而,一個長期被忽視的問題是:這些方法大多默認模型是自回歸生成的。
對于離散擴散語言模型(Discrete Diffusion Language Models, dLLMs)而言,情況完全不同。dLLM 并不是從左到右逐 token 生成,而是從被 mask 的序列出發,通過多步去噪逐漸恢復完整答案。這種并行、非自回歸的生成方式天然具備全局雙向上下文,也更適合規劃與自我修正;但與此同時,傳統面向自回歸模型設計的樹搜索、過程獎勵模型和 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
- 代碼地址:https://github.com/viiika/Prism
![]()
傳統 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 不會粗暴地丟棄整條軌跡重新采樣,而是保留高置信部分,只對低置信位置進行重新 mask,然后從這些局部變化中生成新的分支。這樣做的好處是,它既保留了已有的高質量結構,又能繼續探索不同的細節實現,避免過早收斂到單一路徑。論文圖 2 對這一過程做了直觀展示:在漸進裁剪階段,PRISM 會圍繞高分軌跡進行局部分支,并通過部分重 mask 生成新的候選。
![]()
不再依賴外部 verifier:模型自己給自己打分
很多 Test-Time Scaling 方法依賴額外的 reward model 或 verifier 來判斷候選答案質量。但這會帶來顯著系統開銷:部署時需要額外加載一個模型,顯存、延遲和工程復雜度都會上升。
PRISM 提出了 Self-Verified Feedback(SVF):直接復用同一個 dLLM 作為二分類驗證器。具體來說,模型先根據中間去噪狀態生成一個完整候選答案,然后構造一個 Yes/No 驗證 prompt,讓模型判斷該答案是否可能正確。PRISM 將 Yes 和 No 的 logits 轉換為一個二元歸一化分數,用于軌跡排序、裁剪和最終選擇。
這一設計的價值在于,它把 verification 從「額外模型」變成了「同一模型的一次輕量自檢」。論文進一步指出,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 不僅能提升任務性能,也能維持更好的推理效率。
![]()
在外部 verifier 對比中,SVF 在 GSM8K 上達到 85.30%,只需加載原本的 8B 模型;雖然 Qwen3-8B verifier 可達到 87.35%,但需要額外加載模型,總參數量達到 16B。論文認為,SVF 的優勢并不在于絕對替代所有外部驗證器,而在于它提供了一條更輕量、更易部署的 dLLM 推理擴展路徑。
![]()
意義:為非自回歸語言模型打開推理期擴展路線
PRISM 的核心貢獻并不是簡單提出一個新的搜索啟發式,而是重新定義了 dLLM 上的 Test-Time Scaling 應該如何發生。
對于自回歸模型,推理搜索通常圍繞「前綴」展開;而對于離散擴散模型,中間狀態是部分 mask 的全局序列,傳統前綴式過程獎勵和樹搜索并不天然適用。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,讓離散擴散模型在無需額外 verifier 的情況下實現高效 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
特別聲明:以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布,本平臺僅提供信息存儲服務。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.