機器之心報道
編輯:陳陳、杜偉
- 大語言模型的推理能力,不再是 AR(自回歸)的專屬。擴散模型現(xiàn)在也能「動腦子」,新框架 d1 讓它們學會了解數(shù)學、懂邏輯、會思考。
當前,強化學習(RL)方法在最近模型的推理任務(wù)上取得了顯著的改進,比如 DeepSeek-R1、Kimi K1.5,顯示了將 RL 直接用于基礎(chǔ)模型可以取得媲美 OpenAI o1 的性能。
不過,基于 RL 的后訓(xùn)練進展主要受限于自回歸的大語言模型(LLM),它們通過從左到右的序列推理來運行。
與此同時,離散擴散大語言模型(dLLM)成為有潛力的語言建模的非自回歸替代。不像以因果方式逐 token 生成文本的自回歸模型那樣,dLLM 通過迭代去噪過程生成文本,在多步驟操作中優(yōu)化序列的同時并通過雙向注意力利用過去和未來的上下文。其中,LLaDA 等開放的掩碼 dLLM 實現(xiàn)了媲美同尺寸自回歸模型的性能,而 Mercury 等閉源 dLLM 進一步展現(xiàn)了出色的推理延遲。
然而,頂級的開源 dLLM 并沒有使用 RL 后訓(xùn)練,使得這一有潛力的研究方向還有很大的挖掘空間。這一范式轉(zhuǎn)變引出了重要的問題:RL 后訓(xùn)練如何在非自回歸上下文中高效地實現(xiàn)?
RL 算法適應(yīng)掩碼 dLLM 面臨一些獨特的挑戰(zhàn),原因在于自回歸模型采用的已有方法(如 PPO、GRPO)通過計算生成序列的對數(shù)概率來估計和優(yōu)化策略分布,導(dǎo)致無法直接應(yīng)用于 dLLM。雖然這種計算在自回歸模型中通過序列因式分解很容易實現(xiàn),但 dLLM 由于它們的迭代、非序列生成過程而缺乏這種自然分解。
為了解決這些問題,來自 UCLA 和 Meta AI 的研究者提出了一個兩階段后訓(xùn)練框架 d1,從而可以在掩碼 dLLM 中進行推理。在第一階段,模型在高質(zhì)量推理軌跡中進行監(jiān)督微調(diào);在第二即 RL 階段,研究者引入了用于掩碼 dLLM 的新穎策略梯度方法 diffu-GRPO,它利用提出的高效一步(one-step)對數(shù)概率估計在 GRPO 的基礎(chǔ)上創(chuàng)建。
研究者表示,他們的估計器利用了隨機提示詞掩碼,作為策略優(yōu)化的一種正則化,使得可以擴展 per batch 的梯度更新數(shù)量并減少 RL 訓(xùn)練所需的在線生成數(shù)量。這將極大地降低計算時間。
- 論文標題:d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learning
- 論文地址:https://arxiv.org/pdf/2504.12216
- 項目主頁:https://dllm-reasoning.github.io/
- GitHub 地址:https://github.com/dllm-reasoning/d1
在實驗部分,研究者使用 LLaDA-8B-Instruct 作為基礎(chǔ)模型實例化 d1。他們將 d1-LLaDA 的性能與基礎(chǔ) LLaDA 模型以及僅使用 SFT 和僅使用 diffu-GRPO 訓(xùn)練的 LLaDA 模型進行比較。結(jié)果表明,d1 在四個數(shù)學和邏輯推理基準測試中始終優(yōu)于基礎(chǔ)模型,如下圖 1 所示。d1-LLaDA 同樣優(yōu)于僅使用 SFT 方法和僅使用 diffu-GRPO 方法的模型。
方法概覽
d1 是一個兩階段框架,通過依次結(jié)合監(jiān)督微調(diào)(SFT)和在線強化學習(RL)來增強預(yù)訓(xùn)練掩碼 dLLMs 的推理性能。
其中,在線強化學習(特別是 GRPO 算法)已被證明能有效提升離線訓(xùn)練語言模型的性能。然而,GRPO 的學習策略并不能直接泛化到 dLLMs。
GRPO 的目標函數(shù)(如公式 3 所示)需要同時計算當前策略 π_θ 和舊策略 π_θold 在以下兩個層面的(對數(shù))似然比:
- token 層面(用于優(yōu)勢權(quán)重計算);
- 序列層面(用于反向 KL 散度項)。
核心問題在于:研究者需要高效計算 dLLMs 生成內(nèi)容的逐 token 對數(shù)概率和序列對數(shù)概率。
自回歸(AR)模型,如 Transformer,直接對每個 token 的對數(shù)概率進行建模,并且可以通過鏈式法則使用一次前向傳遞輕松計算出序列級別的對數(shù)概率
同樣,KL 項可以分解為。
與 AR 模型不同,dLLMs 不遵循序列對數(shù)概率的順序分解。同時,每個 token 的對數(shù)概率計算成本也很高,因為解碼過程中需要多次調(diào)用掩碼預(yù)測器 f_θ。基于此,該研究提出了一個高效的對數(shù)概率估計器。
對于序列對數(shù)概率,該研究使用均場近似方法,將其分解為獨立的每個 token 對數(shù)概率的乘積。
對于每個 token 的對數(shù)概率,該研究引入了一種估計方法,該方法僅調(diào)用一次 f_θ。
基于新引入的對數(shù)概率估計器,該研究將 GRPO 擴展到掩碼 dLLMs,推導(dǎo)出 diffu-GRPO 的損失函數(shù)。
算法如下圖所示。
實驗結(jié)果
表 1 報告了基線模型 LLaDA-8B-Instruct 與采用不同后訓(xùn)練優(yōu)化方案的模型,在四項任務(wù)上的零樣本性能對比。
圖 3 繪制了有效 token 的平均數(shù)量:
基于實驗,該研究得出以下主要發(fā)現(xiàn):
diffu-GRPO 在所有 12 種設(shè)置中都一致優(yōu)于基礎(chǔ)的 LLaDA 和 SFT(監(jiān)督式微調(diào))。diffu-GRPO 和 SFT 都相較于 LLaDA-8B-Instruct 基線有所提升,但 diffu-GRPO 顯示出更持續(xù)且幅度更大的增益。具體來說,diffu-GRPO 在所有 12 種設(shè)置中都優(yōu)于 LLaDA-8B-Instruct 和 SFT,而 SFT 僅在其中的 7 種設(shè)置中優(yōu)于 LLaDA-8B-Instruct,這表明diffu-GRPO 相比于單獨的 SFT 實現(xiàn)了更強的整體性能提升。
LLaDA+diffu-GRPO 在所有設(shè)置中都優(yōu)于基礎(chǔ)的 LLaDA-8B-Instruct 模型,而 d1-LLaDA 在每種情況下都超過了 LLaDA+SFT。這表明,無論初始化是來自預(yù)訓(xùn)練模型還是經(jīng)過 SFT 調(diào)整的檢查點,diffu-GRPO 都能提供可靠的性能提升。
d1 訓(xùn)練方案實現(xiàn)了最顯著的性能提升。通過先進行監(jiān)督微調(diào)(SFT)、再結(jié)合 diffu-GRPO 訓(xùn)練所形成的 d1-LLaDA 模型,產(chǎn)生了超越單一方法的疊加增益。這種組合式方法在 12 個實驗設(shè)置中有 11 項優(yōu)于純 diffu-GRPO 方案,表明兩個訓(xùn)練階段存在協(xié)同效應(yīng)。
定性結(jié)果表明,在 SFT 和 d1-LLaDA 生成中出現(xiàn)了頓悟時刻。盡管與 LLaDA-8B-Instruct 相比,生成序列長度為 128 和 256 的性能隨著 SFT、diffu-GRPO 和 d1 有所提高,但從質(zhì)的方面看,在生成的推理軌跡中并未觀察到顯著差異。然而當序列長度達到 512 時,該研究開始觀察到 SFT 和 d1-LLaDA 模型展現(xiàn)出兩種關(guān)鍵能力:自我修正機制和回溯行為。
特別聲明:以上內(nèi)容(如有圖片或視頻亦包括在內(nèi))為自媒體平臺“網(wǎng)易號”用戶上傳并發(fā)布,本平臺僅提供信息存儲服務(wù)。
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.