1. 研究背景與問題定義

基于擴散的大型語言模型(DLLMs)作為自回歸模型(ARMs)的重要擴展，正在成為生成式AI領域的重要創新方向。與傳統ARMs按預定義順序順序生成標記的方式不同，DLLMs提供了并發標記生成、更高輸出多樣性、增強全局一致性以及更好的生成文本可控性等優勢。近期的突破性模型如LLaDA、Mercury和Gemini Diffusion都凸顯了DLLMs的潛力。

然而，當前的掩碼擴散模型(MDM)存在一個關鍵限制：無法有效捕獲并發預測的標記之間的依賴關系。這導致在標記間依賴性較強的推理任務中性能下降。例如，在預測"A poker hand that consists of two English words is: a a"的后續兩個詞時，適合的預測應為"high card"、"two pair"、"full house"或"straight flush"。這些詞對之間存在強依賴關系，但MDM在并發預測時會獨立采樣，無法考慮這種依賴性，從而可能產生不合理的組合。

2. 變分掩碼擴散(VMD)模型

為解決上述問題，研究者提出了變分掩碼擴散(Variational Masked Diffusion, VMD)框架，通過引入潛在變量來建模并發預測期間的聯合標記分布。VMD的核心思想是：通過潛在變量模型捕獲任意聯合分布，而不僅僅是可因式分解的分布。

圖片

2.1 基本變分公式

VMD的基本公式為：

pθ(x0i|xt) = ∫pθ(x0i|xt,z)p(z)dz

其中z是全局潛在變量，不依賴于標記位置i。這使得模型能夠在標記之間建立聯合分布。條件于潛在變量z，標記可以獨立采樣，但邊緣化潛在變量后，能夠從正確的聯合分布中獲得樣本。

訓練目標函數(NELBO)為：

-log pθ(x0) ≤ ∫(0→1) (1/t) ??qt|0(xt|x0)[??q?[∑i:xti=[MASK] -log pθ(x0i|xt,z)] + DKL(q?(·|x0,xt)||p(·))]dt

其中p(·)是潛在空間的先驗分布，q?(·|x0,xt)是由可訓練參數φ參數化的近似后驗分布。

2.2 塊擴散公式

為了解決大規模標記序列中的依賴建模問題，VMD進一步引入了塊擴散公式。將標記x分為B個連續塊，每個塊長度為r，總序列長度L=Br。每個塊使用一個潛在變量zb。數據對數似然可以因式分解為：

log pθ(x0|xt) = ∑(b=1→B) log pθ(x0b|xtb,x0<b) = ∑(b=1→B) log ∫pθ(x0b|xtb,x0<b,z≤b)p(z≤b)dz≤b

這種方法結合了自回歸和擴散模型的優勢，在塊內使用變分擴散，塊間使用自回歸處理。通過調整塊大小B(或等效的r)，可以在自回歸變分模型和變分掩碼擴散模型之間進行插值。

2.3 重新掩碼策略

VMD考慮了兩種重新掩碼策略：

1. 基于概率的置信度：cprob.i = pθ(x0i=v1|xt,z)，其中v1是詞匯表中最可能的值
2. 基于邊際的置信度：cmarg.i = |pθ(x0i=v1|xt,z) - pθ(x0i=v2|xt,z)|，其中v1和v2是兩個最可能的值

這些策略通過潛在變量提供了更全局的上下文，相比傳統MDM的局部置信度指標更為有效。

3. 實驗結果與分析

3.1 控制合成數據實驗

3.1.1 兩標記序列實驗

研究者首先在兩標記序列上進行了實驗，包括確定性依賴、非均勻分布和可變依賴強度三種設置：

• 確定性依賴：數據為{(k,k+1mod10)}k=0→9，第二個標記完全由第一個決定。標準MDM在并發解碼時退化為隨機猜測(約10%準確率)，而VMD能可靠地生成正確對。
• 非均勻分布：分布為P((k,k+1mod10))=(k+1)/55，k∈{0,...,9}。VMD更準確地從目標分布中采樣。
• 可變依賴強度：通過參數p控制依賴強度，VMD在整個依賴強度范圍內都能準確建模真實數據分布。

3.1.2 四標記序列實驗

在四標記序列實驗中，研究者構建了兩個數據集：

• D1：包含10個唯一序列{(k,k+1,k+2,k+3mod10)}k=0→9，具有強標記依賴性
• D2：{(k,k+1,l,l+1mod10)}k,l=0→9，第一個長度為2的塊與第二個塊獨立

結果顯示，在并行解碼所有四個標記時(B=1, NFE=1)，標準MDM表現類似隨機猜測(D1為0.1%，D2為1.1%)，而VMD達到了顯著更高的準確率(D1為81.5%，D2為64.4%)。

3.2 數獨數據實驗

數獨是一個9×9網格的邏輯謎題，需要填充空白單元格，使每行、每列和每個3×3子網格都包含1到9的所有數字。這種全局和局部依賴性使其成為評估生成模型捕獲標記依賴能力的良好基準。

實驗結果表明，VMD在各種采樣方法和NFE值下都優于基線模型。特別是在較低NFE值時，VMD的優勢更為明顯，表明它能更高效地生成有效解決方案。

3.3 文本數據實驗

在text8數據集上的實驗(包含維基百科的前1億個字符)，VMD達到了比標準擴散基線(SEDD, MDLM)更低的困惑度，并在兩種塊大小上都比BD3-LM有輕微但一致的改進。特別是在塊大小為4時，VMD達到了2.858的困惑度，這是基于擴散模型中的最佳結果。

圖片

4. 技術細節與實現

4.1 模型架構

VMD的編碼器和解碼器都采用了類似于BD3-LM的架構，但進行了修改以包含潛在信息。為確保公平比較，解碼器主干與基線保持相同，而編碼器的層數進行了調整，使其架構鏡像解碼器但參數數量減半。

在數獨實驗中，模型使用了DiT架構，編碼器和解碼器分別有4層和6層Transformer層。每個塊分配了一個128維的潛在向量，通過共享的自適應層歸一化模塊注入到解碼器的每個DiT塊中。

4.2 訓練細節

• 合成數據：批量大小10,000，訓練2,000步，使用Adam優化器，固定學習率1e-3
• 數獨：批量大小1,024，學習率1e-3，余弦學習率調度器，300個周期
• 文本：批量大小512，使用AdamW優化器，學習率3e-4，恒定學習率調度，2,500步預熱

4.3 推理過程

推理過程遵循塊擴散方法，應用KV緩存以提高采樣效率。每個生成的塊x0b存儲在緩存中，這意味著訓練期間使用的前綴上下文x0<b對應于累積的鍵和值(K1:b-1,V1:b-1)。解碼器對塊b的預測因此僅依賴于當前噪聲輸入xtb、潛在變量z≤b和來自早期生成塊的緩存上下文。

5. 與相關工作的比較

VMD與其他離散擴散模型和非自回歸模型相比具有以下優勢：

1. 相比標準掩碼擴散模型(如LLaDA、LLaDA-V)，VMD能更好地捕獲并發采樣標記之間的依賴關系
2. 相比BD3-LM，VMD在保持塊自回歸結構的同時，進一步提高了并行采樣的效率和準確性
3. 相比其他非自回歸模型(如BERT)，VMD通過潛在變量提供了更豐富的文本表示

6. 未來展望

VMD為掩碼擴散模型開辟了新的研究方向，未來可能的發展包括：

1. 多層次潛在變量：引入層次化潛在變量結構，以更好地捕獲不同抽象級別的依賴關系。例如，可以設計一個包含句子級、段落級和文檔級潛在變量的模型，每個級別負責不同范圍的依賴建模。
2. 條件生成增強：將VMD擴展到更復雜的條件生成任務，如文本到圖像、跨模態翻譯等。潛在變量可以作為不同模態之間的橋梁，促進更一致的跨模態生成。
3. 可解釋性研究：探索潛在空間的語義意義，開發可視化和分析工具，幫助理解模型如何表示和利用標記間依賴關系。這可以通過潛在空間的降維分析和特定語言現象的干預實驗來實現。
4. 計算效率優化：開發更高效的推理算法，減少VMD的計算開銷。可能的方向包括自適應采樣策略、稀疏注意力機制和模型量化技術。
5. 應用擴展：將VMD應用于更多需要強依賴建模的任務，如程序合成、數學推理和科學發現。特別是在需要保持全局一致性的復雜推理任務中，VMD的優勢可能更為明顯。

這些方向不僅具有理論創新性，也有實際應用價值，可以進一步提升生成模型在復雜依賴建模方面的能力，推動生成AI向更智能、更一致的方向發展。

7. 結論

變分掩碼擴散(VMD)模型通過引入潛在變量，成功解決了標準掩碼擴散在并發標記預測中無法有效捕獲依賴關系的問題。在合成數據、數獨謎題和文本數據上的實驗都證明了VMD的有效性，特別是在標記間依賴關系重要的場景中。VMD不僅提高了生成質量，也增強了對依賴關系的感知，凸顯了將變分推理集成到掩碼擴散中的價值。

參考資源

• 論文鏈接：https://arxiv.org/abs/2510.23606
• 代碼實現：https://riccizz.github.io/VMD