1. 首頁 > 生活百科 > >

用更一致的軌跡、更少的解碼步數「馴服」掩碼擴散語言模型

ar算法數學ass

用更一致的軌跡、更少的解碼步數「馴服」掩碼擴散語言模型

文章圖片

用更一致的軌跡、更少的解碼步數「馴服」掩碼擴散語言模型

文章圖片

用更一致的軌跡、更少的解碼步數「馴服」掩碼擴散語言模型

文章圖片

用更一致的軌跡、更少的解碼步數「馴服」掩碼擴散語言模型

文章圖片




擴散大語言模型得到了突飛猛進的發展 ， 早在 25 年 2 月 Inception Labs 推出 Mercury—— 第一個商業級擴散大型語言模型 ， 同期人民大學發布第一個開源 8B 擴散大語言模型 LLaDA ， 5 月份 Gemini Diffusion 也接踵而至 。 種種跡象表明 ， 擴散大語言模型很可能是下一代大語言模型基礎范式的有力競爭者 。 但是針對于擴散大語言模型的解碼策略和強化學習算法仍然是欠探索的 。

近期 ， 復旦大學、上海人工智能實驗室、上海交通大學聯合研究團隊發布最新論文《Taming Masked Diffusion Language Models via Consistency Trajectory Reinforcement Learning with Fewer Decoding Step》 。

他們提出了一套對于掩碼擴散大語言模型（Masked Diffusion Large Language Model ， MDLM）的高效解碼策略 + 強化學習訓練組合 ， 顯著提升了掩碼擴散大語言模型的推理性能與效率 ， 為擴散大語言模型的發展開辟了新路徑 。



代碼倉庫：https://github.com/yjyddq/EOSER-ASS-RL 論文鏈接：https://arxiv.org/pdf/2509.23924
解決了什么問題？

掩碼擴散大語言模型（MDLM）如 LLaDA 展現出與自回歸模型相媲美的能力 ， 并具備并行解碼、靈活生成順序、潛在少步推理等優勢 。 然而 ， 完全擴散式（Full Diffusion-Style）解碼策略并未被廣泛使用 ， 取而代之的是分塊解碼（Block-wise） 。 因為目前的完全擴散式解碼存在一大痛點 —— 性能大幅度遜色于分塊解碼 。



但令人疑惑的是 ， 掩碼擴散大語言模型在預訓練和指令微調階段并未針對分塊解碼這種方式進行適配微調 ， 所以這一現象背后的原因仍不為人所知 。 該團隊基于這點發現刨根問底 ， 最終定位到 MDLM 的全擴散式解碼的三個特點：

解碼過程中的 token 置信度變化趨勢：由平緩到陡升； 解碼過程中EOStoken 的置信度一直顯著大于其他非EOStoken； 以上兩點導致在早期解碼時都會有很大概率解碼出EOStoken ， 模型像是陷入了EOS的陷阱無法自拔 。 而分塊解碼由于塊的限制存在 ， 則不會深陷其中 。


此外 ， 在將為 AR LLMs 定制的強化學習算法遷移到 MDLM 時 ， 可能會遇到 rollout 軌跡和優化軌跡不一致的問題 ， 因為 AR 模型具有因果性掩碼 ， 獲得完整軌跡后計算得到的 token 概率與 rollout 時保持一致 。 而 MDLM 采用的是雙向注意力 ， 獲得完整軌跡再計算得到的 token 概率與 rollout 時不一致 。 而現有的方法是（1）使用 prompt masking 近似的一步優化 ， 或者（2）從 fully masked response 進行一步去噪優化 。 但是這二者都面臨 rollout 軌跡和實際優化軌跡不一致的問題 ， 可能會引起較大的優化誤差 。

基于此 ， 該團隊將以上問題匯總為三個關鍵問題：

解碼陷阱：全擴散解碼容易在早期步驟中過早生成過多的EOStoken ， 導致生成中斷 ， 性能下降 。 等間隔解碼步長：根據 token 置信度的變化曲線 ， 等間隔的解碼步長可能不是最優的 ， 因為解碼前期對解碼置信度并不高 ， 而解碼后期很確定 。 rollout 軌跡和優化軌跡不一致：直接將自回歸模型的強化學習算法遷移至 MDLM 可能會面臨 rollout 軌跡和優化軌跡不一致問題 ， 影響優化效果 。
提出了什么方法？

該團隊提出三大核心貢獻 ， 致力于解決了上述問題：



針對于EOS陷阱！該團隊提出EOS早期拒絕機制在解碼早期主動抑制EOS的置信度 ， 避免生成過早終止 。 后期逐步恢復EOS的置信度 ， 確保句子正常結束 。 顯著提升全擴散式解碼性能 ， 在部分任務上顯著超越半自回歸解碼 。

基于「早期置信度低、后期急劇上升」的觀察 ， 前期謹慎解碼 ， 后期大膽解碼 ， 設計冪次遞增解碼步長調度器 。 將推理步數從 O (L) 降至 O (logL) ， 大幅加速推理過程 。



針對 rollout 軌跡和優化軌跡不一致問題 ， 該團隊提出一致性軌跡分組策略優化（CJ-GRPO） ， 在 rollout 過程中存儲每一步解碼的中間狀態 ， 逐步優化相鄰兩個狀態之間的轉變 ， 從而有效緩解跳步優化（或者說是不一致軌跡優化）帶來的誤差 ， 提升訓練穩定性與效果 。

但是 ， 由于過長的步數會引入較大的存儲開銷 ， 而結合遞增步長調度器后 ， 既能保證軌跡的一致性 ， 同時能大大縮減中間狀態存儲開銷 。 于是該團隊結合EOS早期拒絕機制、遞增步長調度器和 CJ-GRPO 算法 ， 削減訓練時 CJ-GRPO 的中間狀態存儲開銷 ， 同時使得訓練后的模型在少解碼步數推理下甚至能達到 Baseline 方法多解碼步數時的性能 —— 一石三鳥 。

【用更一致的軌跡、更少的解碼步數「馴服」掩碼擴散語言模型】模型在少解碼步數 (logL) 時仍能保持可觀的性能（與 L/2 解碼步數性能可比） ， 真正激發擴散語言模型的推理速度優勢的潛能 。訓練時的解碼時間 / 空間復雜度從 O (L) 降至 O (logL) ， 大幅加速訓練過程 。

實驗結果



在數學推理（GSM8K、MATH500）和規劃任務（Countdown、Sudoku）上 ， 使用 LLaDA-8B-Instruct 模型進行了廣泛的實驗 ， 結果顯示：
一致性軌跡優化在數學和規劃任務上實現了對 baseline 的全面超越 。 CJ-GRPO + Semi-AR 在數學類任務上表現尤為突出 。 CJ-GRPO + EOSER + ASS 在規劃類任務上表現尤為突出 ， 某些 benchmark 的性能提升高達 2–4 倍 。 發現規劃任務適合并行推理 ， 數學問題適合順序推理 。

在都僅使用 log (L) 步數的情況下 ， EOSER + ASS 性能仍優于分塊解碼和全擴散式解碼策略 。 一致性軌跡優化結合遞增步長調度器和 EOSER 解碼策略 ， 在少步數推理場景下保持良好的競爭力 ， 真正實現「又快又好」 。
研究意義與未來展望

本工作探究了擴散語言模型的全擴散式解碼策略 ， 并對其使用更一致的軌跡、更少的解碼步數進行優化 ， 用更少的步數 ， 越快越好地完成更復雜的推理任務 ， 推動擴散語言模型的全擴散式解碼、少步數解碼、強化學習算法的發展 。

并行推理 v.s. 順序推理：裝配了并行解碼的 MDLM 在規劃類任務中表現更佳 ， 而數學類任務則更適合半自回歸 / 分塊解碼 ， 這比較符合人類的直觀感受 。

未來可探索混合推理模式 ， 結合擴散與自回歸優勢 ， 適應多樣化任務需求 。

歡迎引用、關注與合作交流 。

    推薦閱讀