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

• 論文地址：https://arxiv.org/abs/2509.23653
• 代碼與模型地址：https://github.com/maple-research-lab/RemeDi

背景

擴(kuò)散語(yǔ)言模型已成為自回歸語(yǔ)言模型的有力替代方案。這一類方法首先定義了一個(gè)將文本逐步破壞為噪聲的前向過程，然后讓模型學(xué)習(xí)從噪聲中恢復(fù)出干凈文本的逆向過程。在這一類方法中，當(dāng)前最主流的是基于掩碼的擴(kuò)散語(yǔ)言模型。該方案要求模型在訓(xùn)練中學(xué)習(xí)恢復(fù)被掩碼的 token，而已經(jīng)被恢復(fù)的 token 則在之后的生成步驟中保持不變，直到生成結(jié)束。這其中蘊(yùn)含了一則假設(shè)：每一步中預(yù)測(cè)的 token 都必然是正確的，無(wú)需修正，直接可以當(dāng)作最后的生成內(nèi)容。這一假設(shè)顯然過于理想 —— 生成過程中，模型不可避免地會(huì)產(chǎn)生預(yù)測(cè)錯(cuò)誤，而我們應(yīng)當(dāng)賦予模型通過自我反思發(fā)現(xiàn)并修正這些錯(cuò)誤的能力。

為解決這一問題，提出一種面向擴(kuò)散語(yǔ)言模型的自我反思式生成范式 —— 再掩碼（remask），并基于這一范式訓(xùn)練了具有「再掩碼」能力的擴(kuò)散語(yǔ)言模型 RemeDi。如圖所示，RemeDi 具備發(fā)現(xiàn)錯(cuò)誤 token，并通過再掩碼將其修正的能力：模型首先生成了 “l(fā)eft”，但隨后在生成完整句子的語(yǔ)義表示時(shí)，發(fā)現(xiàn) “l(fā)eft for the pies” 這一表述與實(shí)際含義不符，因此，將 “l(fā)eft” 一詞再掩碼，修改為更合適的 “used”?？梢钥闯?，通過再掩碼，模型能利用在后續(xù)步驟中生成的上下文信息，識(shí)別較早步驟中存在的錯(cuò)誤，將其改正，并基于更豐富的上下文信息進(jìn)行更精確的預(yù)測(cè)。

用置信度識(shí)別「再掩碼」目標(biāo)

為了讓 RemeDi 能夠通過再掩碼修改已經(jīng)生成的文本內(nèi)容，一個(gè)核心的挑戰(zhàn)是讓模型能夠找到需要修改的 token，執(zhí)行再掩碼操作。為此，我們對(duì)網(wǎng)絡(luò)結(jié)構(gòu)進(jìn)行了修改，讓其在預(yù)測(cè)序列中每個(gè) token 輸出分布的同時(shí)，能夠?yàn)槊總€(gè) token 額外預(yù)測(cè)一個(gè)置信度分?jǐn)?shù)。整個(gè)模型采用了一種雙流協(xié)同的模型結(jié)構(gòu)：

• TPS（Token Prediction Stream）：負(fù)責(zé)對(duì)掩碼位置給出候選 token 分布 ，類似常規(guī)的擴(kuò)散語(yǔ)言模型；
• UPS（Unmasking Policy Stream）：為序列每一個(gè)位置輸出一個(gè)置信度分?jǐn)?shù)，表示模型在這一步輸出時(shí)，該位置上結(jié)果的確定度。分?jǐn)?shù)高即說明模型認(rèn)為，這一步的結(jié)果有更大的概率是正確的，無(wú)需再被掩碼。與此同時(shí)，得分較低的位置就應(yīng)當(dāng)仍然保持掩碼狀態(tài)，或是被再掩碼，直到模型能依賴更多上下文做出更準(zhǔn)確的預(yù)測(cè)。

基于這一模型結(jié)構(gòu)，RemeDi 按如下方式逐步執(zhí)行去噪推理步驟：以上一步的結(jié)果 作為輸入，UPS 模塊首先會(huì)為序列中每一個(gè)位置預(yù)測(cè) ，決定哪些位置不再需要被掩碼。然后，對(duì)于那些不需要掩碼的位置，如果輸入本身就已經(jīng)不是掩碼 token，我們會(huì)直接保留輸入 token 值；否則，我們會(huì)基于 TPS 輸出的 采樣該位置的輸出 token。與 “生成即固定” 的傳統(tǒng)掩碼擴(kuò)散生成范式不同，RemeDi 在每一步都會(huì)依賴輸出的置信度決定需要 / 不需要掩碼的部分。因此，模型有可能對(duì)已經(jīng)生成的 token 預(yù)測(cè)出較低的置信度，將其「再掩碼」，使其后續(xù)可以依據(jù)更充分的上下文重寫，使推理過程具備 “邊寫邊改” 的能力。

此外，在語(yǔ)言生成任務(wù)中，許多場(chǎng)景下的輸出并非固定長(zhǎng)度。如果模型只能在固定長(zhǎng)度下生成，將導(dǎo)致資源浪費(fèi)或生成結(jié)果被壓縮、截?cái)唷Ｒ虼?，使擴(kuò)散語(yǔ)言模型具備靈活的不定長(zhǎng)生成能力 （variable-length generation）是必要的。在 RemeDi 中，我們采用分塊自回歸生成的方法實(shí)現(xiàn)這一點(diǎn)：模型每次會(huì)通過一個(gè)完整的反向擴(kuò)散過程生成一段長(zhǎng)為 L=32 的序列。完成后，如果該序列中沒有生成結(jié)束符，則將已生成的這一段序列拼接在上下文中，繼續(xù)往后生成下一段長(zhǎng)為 L=32 的序列，如此重復(fù)直到生成結(jié)束符為止。與自回歸模型類似，我們采用分塊因果注意力掩碼機(jī)制，確保在生成時(shí)，每個(gè) token 能看到自己所在的 block 內(nèi)的其他 token，和之前已生成 block 內(nèi)的 token，而無(wú)法看到未來將要生成的 block。

在實(shí)驗(yàn)中，我們基于 LLaDA 的權(quán)重繼續(xù)訓(xùn)練，將其改造成一個(gè)具有不定長(zhǎng)生成能力的分塊擴(kuò)散模型。上面表 4 中的 baseline 模型即展示了不定長(zhǎng)生成模型在經(jīng)過再掩碼訓(xùn)練前的性能。

兩階段訓(xùn)練，賦予「再掩碼」能力

1.Remask SFT（監(jiān)督微調(diào)階段）

傳統(tǒng)的掩碼擴(kuò)散語(yǔ)言模型通常通過在輸入序列上隨機(jī)掩碼進(jìn)行有監(jiān)督微調(diào)（SFT）。與之不同的是，RemeDi 在反向擴(kuò)散過程中還需要能夠找到潛在的不正確 token 并再掩碼。我們?cè)?SFT 過程中將這類不正確 token 視為除掩碼 token 之后的第二類噪聲。因此，在 SFT 階段，我們不僅要訓(xùn)練模型從掩碼 token 恢復(fù)原文本的能力，同時(shí)也需要訓(xùn)練識(shí)別那些需要再掩碼的不正確 token。

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

由于在反向擴(kuò)散過程中，噪聲水平（定義為 mask token 的數(shù)量）應(yīng)當(dāng)單調(diào)遞減。由于在 SFT 設(shè)計(jì)中，長(zhǎng)度為 L 的輸入序列中，所有不正確 token 都必須被重新掩碼，因此需要滿足以下不等式約束：

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

基于上述約束，我們選擇噪聲調(diào)度為 ，以及 ，其中 r 為常數(shù)。在實(shí)驗(yàn)中我們?cè)O(shè)定 r=0.1，此時(shí)不難驗(yàn)證在 區(qū)間上，上述不等式總是成立。

在實(shí)際訓(xùn)練過程中，除了常規(guī)的 token 預(yù)測(cè)損失外，我們還需要在所有 token 位置上使用二元交叉熵（BCE）目標(biāo)函數(shù)監(jiān)督模型預(yù)測(cè)的 。我們按以下規(guī)則構(gòu)造對(duì)應(yīng)的訓(xùn)練標(biāo)簽 y：

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

整個(gè)再掩碼微調(diào)算法流程如下圖：

2.Remask RL（強(qiáng)化學(xué)習(xí)階段）

在完成 Remask SFT 訓(xùn)練后，我們進(jìn)一步通過基于結(jié)果的強(qiáng)化學(xué)習(xí)對(duì)模型進(jìn)行微調(diào)。根據(jù)實(shí)驗(yàn)室先前的研究，反向擴(kuò)散過程中的每一步中間結(jié)果都可以視為大模型的一個(gè)「思考」步驟，而基于結(jié)果的強(qiáng)化學(xué)習(xí)可以優(yōu)化整個(gè)生成軌跡，提升模型生成正確最終答案的概率。這種面向擴(kuò)散語(yǔ)言模型的大模型推理范式稱為擴(kuò)散式「發(fā)散思維鏈」，在機(jī)器之心的往期報(bào)道中已有詳細(xì)闡述。

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

1）去掩碼策略：UPS 為每個(gè) token 位置生成一個(gè)置信度分?jǐn)?shù) ，表示模型多大程度上確信該位置上的 token 是正確的（若已去掩碼）或可預(yù)測(cè)的（若仍為掩碼）。在推理時(shí)，我們根據(jù)該置信度對(duì)所有 token 排序，并優(yōu)先為置信度高的位置去掩碼。在 RL 訓(xùn)練中，我們基于 Plackett–Luce 模型構(gòu)造解掩碼策略：根據(jù) 無(wú)放回地順序采樣該步驟的去掩碼位置集合  。這一去掩碼位置集合的采樣概率為：

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

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

該策略可用于基于結(jié)果的強(qiáng)化學(xué)習(xí)，鼓勵(lì)所有能夠得到正確答案的完整軌跡。

實(shí)驗(yàn)結(jié)果

在同規(guī)模與相近計(jì)算預(yù)算下，RemeDi 在數(shù)學(xué)推理、代碼生成與通用問答三類任務(wù)上均取得穩(wěn)定提升。其中，僅采用 Remask SFT 帶來顯著增益；在此基礎(chǔ)上加入 Remask RL，多數(shù)基準(zhǔn)再獲得進(jìn)一步提升。

我們?cè)诓煌愋偷娜蝿?wù)上對(duì)再掩碼次數(shù)進(jìn)行了統(tǒng)計(jì)，可以看出：對(duì)輸出約束更強(qiáng)的任務(wù)（如代碼生成）會(huì)更頻繁觸發(fā)再掩碼。

而具體的生成示例也表明，通過再掩碼機(jī)制，RemeDi 可以實(shí)現(xiàn)糾錯(cuò)、插入、刪除等多種文本修改手段。

總結(jié)

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