鑒於自回歸大語言模型（LLM）在計算效率、吞吐量上的局限性，擴散語言模型（DLM）正受到越來越多的關注。

然而，業內目前缺少一個可適用於各種 DLM 架構（如全注意力 DLM 和塊注意力 DLM）的統一且有效的強化學習（RL）框架。同時，現有研究也忽視了對齊推理軌跡與訓練目標的重要性。

日前，普林斯頓大學王夢迪教授團隊在一項最新研究中提出了“軌跡感知 RL”框架——TraceRL，其可以應用於全注意力和塊注意力模型，並實現快速優化。

論文鏈接：http://arxiv.org/abs/2509.06949

值得一提的是，使用 TraceRL 訓練的 4B DLM 模型，在多項複雜數學推理任務上的表現超過了 7B 自回歸模型。

通過課程學習，他們還推出了第一個長思維鏈擴散語言模型（long-CoT DLM），在 MATH500 上的相對準確率比 Qwen2.5-7B-Instruct 提高了 18.1%。

圖｜左：不同方法的 RL 訓練動態，其中 TraceRL 達到最佳優化。右：基於 KV 緩存評估的複雜數學推理任務與 LiveCodeBench-V2 基準測試結果對比。

同時，他們也提出了一個基於擴散的價值模型，從而降低方差並提高訓練穩定性。他們還探索了 TraceRL 的更廣泛潛在應用，如增加模型的塊大小和加速推理。

此外，他們還開源了一個完整集成框架——dLLM-RL，用於跨不同架構構建、訓練和部署 DLM。該框架包含了各種後訓練方法和加速 KV 緩存技術的實現，既支持可複現的研究，也支持實際應用。

地址：http://github.com/Gen-Verse/dLLM-RL

亟需解決DLM“不匹配”難題

研究團隊強調，DLM 在後訓練階段所采用的目標，與其在實際推理（生成文本）時所遵循的軌跡之間，存在著顯著的不匹配。標準訓練方法，如完全隨機掩碼，雖然能夠實現並行解碼，但忽略了語言本身固有的、依賴上下文的有序邏輯。這種訓練與推理行為的脫節使得模型優化效率低下。

為了說明這一差異，他們首先通過實驗證明，采用半自回歸式的微調方法，即訓練模型根據先前的上下文生成後續內容，其優化性能顯著優於全隨機掩碼方法，即使在計算負載相同的情況下也是如此。這表明，將訓練目標與推理模式對齊至關重要。

為了進一步驗證對齊的重要性，他們收集了模型自身的“偏好推理軌跡”，即模型在生成內容時實際遵循的步驟順序。實驗結果顯示，使用這些真實的推理軌跡進行微調，所取得的性能超過了其他基線方法，並且計算成本更低或相當。

最後，盡管使用偏好軌跡進行微調效果很好，但收集這些軌跡需要大量額外工作。相比之下，RL 在其“rollouts”（即模型生成樣本）過程中會自然地產生這些推理軌跡。因此，RL 是一種更實用、更有效的訓練後策略，能夠天然地利用這些軌跡來優化模型。

TraceRL：“小擴散語言”>“大自回歸”

在這項工作中，TraceRL 關注由 DLM 生成的中間軌跡且可跨架構應用。

圖｜TraceRL 概述。本示例采用 s=2、L=6 和 B=3 的參數配置。通過聚合每 s 個相鄰步驟，實現了軌跡感知 RL。方框內的數字對應策略推理過程的執行順序。

在數據方麵，研究團隊采用了不同的數據源：

（1）選用 Math 訓練數據集，生成 8000 個硬核任務；

（2）選擇 GSM8K、MATH500、AIME2024 作為測試基準，以評估聚焦於數學與編程領域的推理任務；

（3）在編碼強化學習場景中，采用 PrimeIntellect 平台提供的 6000 道已驗證題目驗證；

（4）針對編程測試，在評估時選用 LiveCodeBench-V2、LiveBench 作為測試基準。

表｜不同數學與編程任務的主要基準測試結果。“靜態”指靜態采樣，“動態”指動態采樣。此處采用閾值 0.9 的動態采樣方法對長 CoT 模型 TraDo-8B-Instruct 進行評估

在模型訓練方麵，包括全注意力模型和塊注意力模型，並在評估過程中同時報告靜態采樣和動態采樣的結果。過程如下：

第一步：由 TraceRL 單獨訓練模型；

第二步：聯合訓練長 CoT 模型；

第三步：TraceRL 與其他 RL 方法進行對比實驗；

第四步：TraceRL 在全注意力模型與編碼任務中的驗證；

第五步：進行區塊大小擴展實驗。

基於實驗結果，他們證明了 TraceRL 的有效性與強性能，完整結果如下：

首先，他們基於 SDAR 基礎模型，利用 TraceRL 開發了 TraDo-4B-Instruct 和 TraDo-8B-Instruct 兩種模型。在數學、編程和 5 個推理數據集的評估中，這些模型不僅與強擴散語言模型、自回歸模型形成有效對比，其生成能力也展現出顯著優勢。

圖｜TraceRL 在數學任務中對 4B 和 8B 模型的訓練曲線。紅色曲線表示動態采樣精度，其采樣速度更快；藍色曲線表示靜態采樣精度，可獲得更高的精度。4B 模型使用價值模型進行訓練，而 8B 模型直接使用 Jpolicy 進行訓練。

TraDo-4B-Instruct 在推理任務上展現出 SOTA 水平，證明了 TraceRL 的有效性。無論是動態采樣（更快）還是靜態采樣（更準確），模型的性能都得到了顯著提升。值得注意的是，在所有數學任務中，TraDo-4B-Instruct 的表現甚至超越了 Qwen2.5-7B Instruct 等強自回歸基線模型。

盡管他們在 RL 訓練過程中采用了動態采樣策略，但動態與靜態準確率均呈現穩步提升趨勢，且有趨勢表明模型仍有擴展潛力。這種 RL 訓練顯著提升了模型的數學推理能力：

在 MATH500 測試中，TraDo-4B Instruct 模型靜態準確率提升 5.4%，動態準確率提升 4.2%，優化後超越 Qwen2.5-7B-Instruct；而 TraDo-8B-Instruct 模型則分別實現靜態準確率 4.2% 和動態準確率 4.8% 的提升。

圖｜基於塊擴散模型和數學 RL 任務的 RL 方法消融實驗。紅色與黃色曲線分別對應啟用和禁用價值模型的 TraceRL 訓練結果，藍色曲線采用類似半自回歸訓練法的隨機掩碼目標進行塊內訓練，綠色曲線則通過在塊內添加互補性掩碼實現額外訓練效果。

他們進一步將 TraceRL 與現有的 RL 方法進行了對比研究，重點關注塊擴散模型。盡管當前 RL 方法主要針對全注意力模型開發，但他們將其直接適配到塊結構中。對於隨機掩碼方法，他們限製每個塊內的采樣操作，使其類似於半自回歸方法。對於耦合強化學習，他們在每個訓練模塊中引入了互補目標函數，從而獲得更穩定且高效的訓練效果。基於數學任務的實驗結果顯示，無論采用何種優化策略，TraceRL均展現出最優性能。

圖｜全注意力模型 Dream-7B-Coder-Instruct 的 RL 訓練消融實驗，重點聚焦編碼任務；使用與不使用價值模型的對比表明，引入價值模型能有效減少訓練過程中的波動。所有實驗均在 4B 模型的數學任務上進行。

此外，為驗證 TraceRL 的廣泛適用性，他們還在編碼 RL 任務中對全注意力模型進行了實驗。基於 Dream-7B-Coder-Instruct，在冷啟動階段使用蒸餾數據進行微調後開展 RL 訓練。為加速訓練進程，將收縮參數設置為 s=8。實驗表明，TraceRL收斂速度更快且性能更優。

未來可期

基於以上實驗結果，研究團隊驗證了 TraceRL 在不同 RL 任務中的有效性。同時，他們還展示了 TraceRL 在加速推理和擴展塊大小方麵的優勢，這為未來的研究提供了有前景的方向。

特別地，將擴散模型的加速推理能力與其潛在的強推理能力相結合，代表了一個令人期待的研究方向。盡管當前的長 CoT LLM 在複雜任務上表現出色，但其推理時間太長。此類整合有望在大規模環境中高效地執行複雜推理任務，從而開辟新的應用可能。

研究團隊還表示，他們提出的擴散價值模型能夠整合過程獎勵，提供比單一可驗證獎勵更強的監督信號。未來，他們將進一步探索基於過程獎勵的 TraceRL 優化。

整理：小瑜

如需轉載或投稿，請直接在公眾號內留言