王夢迪團隊推出TraceRL：邁向擴散語言模型「RL大一統(tǒng)」

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

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

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

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

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

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

圖｜左：不同方法的 RL 訓練動態(tài)，其中 TraceRL 達到最佳優(yōu)化。右：基于 KV 緩存評估的復雜數(shù)學推理任務與 LiveCodeBench-V2 基準測試結果對比。

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

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

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

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

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

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

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

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

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

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

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

在數(shù)據(jù)方面，研究團隊采用了不同的數(shù)據(jù)源：

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

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

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

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

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

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

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

• 第二步：聯(lián)合訓練長 CoT 模型；

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

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

• 第五步：進行區(qū)塊大小擴展實驗。

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

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

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

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

盡管他們在 RL 訓練過程中采用了動態(tài)采樣策略，但動態(tài)與靜態(tài)準確率均呈現(xiàn)穩(wěn)步提升趨勢，且有趨勢表明模型仍有擴展?jié)摿Α＿@種 RL 訓練顯著提升了模型的數(shù)學推理能力：

在 MATH500 測試中，TraDo-4B Instruct 模型靜態(tài)準確率提升 5.4%，動態(tài)準確率提升 4.2%，優(yōu)化后超越 Qwen2.5-7B-Instruct；而 TraDo-8B-Instruct 模型則分別實現(xiàn)靜態(tài)準確率 4.2% 和動態(tài)準確率 4.8% 的提升。

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

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

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

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

未來可期

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

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

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

整理：小瑜

如需轉載或投稿，請直接在公眾號內(nèi)留言