1. 首頁 > 生活百科 > >

掩碼生成到「再掩碼」訓練：RemeDi讓擴散語言模型學會糾正與反思

廈門大疆運動相機

掩碼生成到「再掩碼」訓練：RemeDi讓擴散語言模型學會糾正與反思

文章圖片

掩碼生成到「再掩碼」訓練：RemeDi讓擴散語言模型學會糾正與反思

文章圖片

掩碼生成到「再掩碼」訓練：RemeDi讓擴散語言模型學會糾正與反思

文章圖片

掩碼生成到「再掩碼」訓練：RemeDi讓擴散語言模型學會糾正與反思

文章圖片

【掩碼生成到「再掩碼」訓練：RemeDi讓擴散語言模型學會糾正與反思】掩碼生成到「再掩碼」訓練：RemeDi讓擴散語言模型學會糾正與反思

文章圖片

掩碼生成到「再掩碼」訓練：RemeDi讓擴散語言模型學會糾正與反思

文章圖片

掩碼生成到「再掩碼」訓練：RemeDi讓擴散語言模型學會糾正與反思

文章圖片

掩碼生成到「再掩碼」訓練：RemeDi讓擴散語言模型學會糾正與反思

文章圖片

掩碼生成到「再掩碼」訓練：RemeDi讓擴散語言模型學會糾正與反思

文章圖片

掩碼生成到「再掩碼」訓練：RemeDi讓擴散語言模型學會糾正與反思

文章圖片

掩碼生成到「再掩碼」訓練：RemeDi讓擴散語言模型學會糾正與反思

文章圖片

掩碼生成到「再掩碼」訓練：RemeDi讓擴散語言模型學會糾正與反思

文章圖片

掩碼生成到「再掩碼」訓練：RemeDi讓擴散語言模型學會糾正與反思

文章圖片



近期 ， 擴散語言模型備受矚目 ， 提供了一種不同于自回歸模型的文本生成解決方案 。 為使模型能夠在生成過程中持續修正與優化中間結果 ， 西湖大學 MAPLE 實驗室齊國君教授團隊成功訓練了具有「再掩碼」能力的擴散語言模型（Remasking-enabled Diffusion Language Model RemeDi 9B） 。 在擴散去噪的多步過程中 ， 通過進行再掩碼 SFT 和 RL 訓練 ， 為每個 token 輸出一個去掩碼置信度 ， RemeDi 能夠從序列中已經生成的內容中識別無法確定的位置進行再掩碼（remask） ， 從而修正錯誤內容并提升文本質量 ， 在各方面都超越了現有的擴散語言模型 。 該模型還具有可變長生成（variable-length generation）能力 ， 打破了現有中大規模擴散語言模型僅支持定長生成的限制 ， 提高了模式能力的靈活性 。

論文地址：https://arxiv.org/abs/2509.23653 代碼與模型地址：https://github.com/maple-research-lab/RemeDi背景
擴散語言模型已成為自回歸語言模型的有力替代方案 。 這一類方法首先定義了一個將文本逐步破壞為噪聲的前向過程 ， 然后讓模型學習從噪聲中恢復出干凈文本的逆向過程 。 在這一類方法中 ， 當前最主流的是基于掩碼的擴散語言模型 。 該方案要求模型在訓練中學習恢復被掩碼的 token ， 而已經被恢復的 token 則在之后的生成步驟中保持不變 ， 直到生成結束 。 這其中蘊含了一則假設：每一步中預測的 token 都必然是正確的 ， 無需修正 ， 直接可以當作最后的生成內容 。 這一假設顯然過于理想 —— 生成過程中 ， 模型不可避免地會產生預測錯誤 ， 而我們應當賦予模型通過自我反思發現并修正這些錯誤的能力 。
為解決這一問題 ， 提出一種面向擴散語言模型的自我反思式生成范式 —— 再掩碼（remask） ， 并基于這一范式訓練了具有「再掩碼」能力的擴散語言模型 RemeDi 。 如圖所示 ， RemeDi 具備發現錯誤 token ， 并通過再掩碼將其修正的能力：模型首先生成了 “left” ， 但隨后在生成完整句子的語義表示時 ， 發現 “left for the pies” 這一表述與實際含義不符 ， 因此 ， 將 “left” 一詞再掩碼 ， 修改為更合適的 “used” 。 可以看出 ， 通過再掩碼 ， 模型能利用在后續步驟中生成的上下文信息 ， 識別較早步驟中存在的錯誤 ， 將其改正 ， 并基于更豐富的上下文信息進行更精確的預測 。

用置信度識別「再掩碼」目標
為了讓 RemeDi 能夠通過再掩碼修改已經生成的文本內容 ， 一個核心的挑戰是讓模型能夠找到需要修改的 token ， 執行再掩碼操作 。 為此 ， 我們對網絡結構進行了修改 ， 讓其在預測序列中每個 token 輸出分布的同時 ， 能夠為每個 token 額外預測一個置信度分數 。 整個模型采用了一種雙流協同的模型結構：



此外 ， 在語言生成任務中 ， 許多場景下的輸出并非固定長度 。 如果模型只能在固定長度下生成 ， 將導致資源浪費或生成結果被壓縮、截斷 。 因此 ， 使擴散語言模型具備靈活的不定長生成能力 （variable-length generation）是必要的 。 在 RemeDi 中 ， 我們采用分塊自回歸生成的方法實現這一點：模型每次會通過一個完整的反向擴散過程生成一段長為 L=32 的序列 。 完成后 ， 如果該序列中沒有生成結束符 ， 則將已生成的這一段序列拼接在上下文中 ， 繼續往后生成下一段長為 L=32 的序列 ， 如此重復直到生成結束符為止 。 與自回歸模型類似 ， 我們采用分塊因果注意力掩碼機制 ， 確保在生成時 ， 每個 token 能看到自己所在的 block 內的其他 token ， 和之前已生成 block 內的 token ， 而無法看到未來將要生成的 block 。

在實驗中 ， 我們基于 LLaDA 的權重繼續訓練 ， 將其改造成一個具有不定長生成能力的分塊擴散模型 。 上面表 4 中的 baseline 模型即展示了不定長生成模型在經過再掩碼訓練前的性能 。
兩階段訓練 ， 賦予「再掩碼」能力
1.Remask SFT（監督微調階段）
傳統的掩碼擴散語言模型通常通過在輸入序列上隨機掩碼進行有監督微調（SFT） 。 與之不同的是 ， RemeDi 在反向擴散過程中還需要能夠找到潛在的不正確 token 并再掩碼 。 我們在 SFT 過程中將這類不正確 token 視為除掩碼 token 之后的第二類噪聲 。 因此 ， 在 SFT 階段 ， 我們不僅要訓練模型從掩碼 token 恢復原文本的能力 ， 同時也需要訓練識別那些需要再掩碼的不正確 token 。

由于在反向擴散過程中 ， 噪聲水平（定義為 mask token 的數量）應當單調遞減 。 由于在 SFT 設計中 ， 長度為 L 的輸入序列中 ， 所有不正確 token 都必須被重新掩碼 ， 因此需要滿足以下不等式約束：

以確保輸出中掩碼位置的數量單調減少 。 若該不等式不成立 ， 則在下一步重新掩碼所有不正確 token 會增加總的掩碼數量 ， 從而違反擴散過程中掩碼比例應逐步減少的基本原則 。



整個再掩碼微調算法流程如下圖：

2.Remask RL（強化學習階段）
在完成 Remask SFT 訓練后 ， 我們進一步通過基于結果的強化學習對模型進行微調 。 根據實驗室先前的研究 ， 反向擴散過程中的每一步中間結果都可以視為大模型的一個「思考」步驟 ， 而基于結果的強化學習可以優化整個生成軌跡 ， 提升模型生成正確最終答案的概率 。 這種面向擴散語言模型的大模型推理范式稱為擴散式「發散思維鏈」 ， 在機器之心的往期報道中已有詳細闡述 。 （與Gemini Diffusion共振！首個擴散式「發散思維鏈」來了）








實驗結果
在同規模與相近計算預算下 ， RemeDi 在數學推理、代碼生成與通用問答三類任務上均取得穩定提升 。 其中 ， 僅采用 Remask SFT 帶來顯著增益；在此基礎上加入 Remask RL ， 多數基準再獲得進一步提升 。


我們在不同類型的任務上對再掩碼次數進行了統計 ， 可以看出：對輸出約束更強的任務（如代碼生成）會更頻繁觸發再掩碼 。

而具體的生成示例也表明 ， 通過再掩碼機制 ， RemeDi 可以實現糾錯、插入、刪除等多種文本修改手段 。


總結
這篇文章介紹了由西湖大學 MAPLE 實驗室推出的 ， 具有再掩碼反思機制的擴散語言模型 ， RemeDi 。 基于額外的置信度預測 ， RemeDi 能夠識別生成過程中的錯誤 ， 并通過「再掩碼」機制重新預測 ， 從而做到生成過程中的自我反思與優化 。 針對「再掩碼」機制設計的有監督訓練與強化學習算法確保了這一機制的有效性 。 實驗結果表明 RemeDi 在數學推理、代碼生成、通用知識問答等多個任務上都取得了超越其他擴散語言模型的性能 。 這些結果說明「再掩碼」能有效提升擴散語言模型的文本生成質量 ， 值得進一步探討 。

    推薦閱讀