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

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

背景

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

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

用置信度識別「再掩碼」目標

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

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

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

兩階段訓練，賦予「再掩碼」能力

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

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

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

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

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

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

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

實驗結(jié)果

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

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

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

總結(jié)

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