近期，擴散語言模型備受矚目，提供了一種不同于自回歸模型的文本生成解決方案。為使模型能夠在生成過程中持續修正與優化中間結果，西湖大學 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 額外預測一個置信度分數。整個模型采用了一種雙流協同的模型結構：

• TPS（Token Prediction Stream）：負責對掩碼位置給出候選 token 分布 ，類似常規的擴散語言模型；
• UPS（Unmasking Policy Stream）：為序列每一個位置輸出一個置信度分數，表示模型在這一步輸出時，該位置上結果的確定度。分數高即說明模型認為，這一步的結果有更大的概率是正確的，無需再被掩碼。與此同時，得分較低的位置就應當仍然保持掩碼狀態，或是被再掩碼，直到模型能依賴更多上下文做出更準確的預測。

基于這一模型結構，RemeDi 按如下方式逐步執行去噪推理步驟：以上一步的結果 作為輸入，UPS 模塊首先會為序列中每一個位置預測 ，決定哪些位置不再需要被掩碼。然后，對于那些不需要掩碼的位置，如果輸入本身就已經不是掩碼 token，我們會直接保留輸入 token 值；否則，我們會基于 TPS 輸出的 采樣該位置的輸出 token。與 “生成即固定” 的傳統掩碼擴散生成范式不同，RemeDi 在每一步都會依賴輸出的置信度決定需要 / 不需要掩碼的部分。因此，模型有可能對已經生成的 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。

我們從干凈文本 引入兩類噪聲構造訓練樣本 ：首先，隨機采樣一個擴散時間 ，并設定對應的隨機掩碼比率 以及不正確 token 的比率 。我們以比例 隨機掩碼一部分 token；接著，在剩余未被掩碼的位置中，以比例 采樣一個子集，并把其中的每個 token 隨機替換為一個其他 token，用以模擬反向擴散過程中可能出現的不正確 token。

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

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

基于上述約束，我們選擇噪聲調度為 ，以及 ，其中 r 為常數。在實驗中我們設定 r=0.1，此時不難驗證在 區間上，上述不等式總是成立。

在實際訓練過程中，除了常規的 token 預測損失外，我們還需要在所有 token 位置上使用二元交叉熵（BCE）目標函數監督模型預測的 。我們按以下規則構造對應的訓練標簽 y：

• 掩碼 tokens ，即 。此類 token 標簽為 y=1，表示該 token 應保持不被掩碼；
• 可見但錯誤的 tokens 即 。此類 token 標簽為 y=0，表示該 token 應被掩碼；
• 可見且正確的 tokens，即 。對這一類 token，我們會賦予軟標簽 ，即模型預測出對應真值 的概率。該概率越高，說明預測出真值的可能性越大，因此該 token 更不應該被掩碼。

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

2.Remask RL（強化學習階段）

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

在具備「再掩碼」能力的 RemeDi 模型中，這一擴散式「發散思維鏈」同樣也包含了 N 個去噪步驟。對于時刻的第 n 步，我們將從 生成 的去噪過程拆解為兩部分策略：

1）去掩碼策略：UPS 為每個 token 位置生成一個置信度分數 ，表示模型多大程度上確信該位置上的 token 是正確的（若已去掩碼）或可預測的（若仍為掩碼）。在推理時，我們根據該置信度對所有 token 排序，并優先為置信度高的位置去掩碼。在 RL 訓練中，我們基于 Plackett–Luce 模型構造解掩碼策略：根據 無放回地順序采樣該步驟的去掩碼位置集合  。這一去掩碼位置集合的采樣概率為：

2）Token 預測策略：對于包含在去掩碼位置集合 中的每一個位置，如果 ，則模型會依據 采樣預測 token 值；否則，該位置 token 值保持輸入不變。這一步中，給定和 采樣的概率為：

綜合上述兩個策略，在一個去噪步驟中，基于上一步結果 采樣 的最終概率可建模為：

該策略可用于基于結果的強化學習，鼓勵所有能夠得到正確答案的完整軌跡。

實驗結果

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

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

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

總結

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