跳到主要內容
AI News HubLIVE
站內改寫1 分鐘閱讀

學習擴散語言模型的去掩碼策略

文章摘要

擴散大型語言模型(dLLM)在多項任務上已達到自迴歸模型的效能,同時推理效率更高。其關鍵設計之一是每個擴散步驟選擇去掩碼令牌的取樣過程。當前啟發式方法需要手動調整且效能隨塊大小增大而下降。本文提出使用強化學習訓練取樣策略,形式化掩碼擴散取樣為馬爾可夫決策過程,並基於單層Transformer設計輕量策略。實驗表明,該策略在半自迴歸生成中與先進啟發式方法效能相當,在全擴散設定中更優。

學習擴散語言模型的去掩碼策略
回報錯誤

更正管道尚未開通,可先複製下方文章資訊留存。

查看更正說明
直接讀正文

擴散語言模型(dLLM)近年來在自然語言處理領域發展迅速,尤其是在生成任務上表現出色。這類模型透過並行解碼多個令牌,顯著提升了推理效率,使其成為自迴歸模型的有力替代方案。然而,dLLM的一個關鍵設計挑戰在於取樣過程:在每個擴散步驟中,模型需要決定哪些令牌被去掩碼(即生成)。傳統的啟發式方法,例如基於置信度閾值的去掩碼策略,雖然比隨機去掩碼在樣本質量和令牌吞吐量上有所提升,但存在明顯缺陷。首先,這些方法需要人工仔細調整引數;其次,研究觀察到,當採用更大的塊大小進行半自迴歸生成時,啟發式方法的效能會顯著下降,這限制了其在實際應用中的可擴充套件性。

為了克服這些侷限性,來自阿姆斯特丹大學、麻省理工學院和蘋果公司的研究人員提出了一種基於強化學習(RL)的創新方法。他們將掩碼擴散取樣過程形式化為一個馬爾可夫決策過程(MDP),其中擴散語言模型本身充當環境,而取樣策略則作為智慧體。具體而言,他們設計了一個輕量級的策略網路,基於單層Transformer架構,該網路將模型輸出的每個令牌的置信度對映為是否去掩碼的決策。透過強化學習訓練,該策略能夠最佳化生成樣本的質量和效率,而無需手動調整。

在實驗中,研究人員將訓練得到的策略與當前最先進的啟發式方法進行了對比。結果表明,在半自迴歸(塊狀)生成設定下,訓練策略的效能與啟發式方法持平;而在全擴散(即每個步驟只解碼單個令牌)設定下,訓練策略則顯著優於啟發式方法。這一發現表明,透過強化學習學習取樣策略不僅消除了手動調整的需求,還能在更廣泛的設定中提升效能。

該研究為擴散語言模型的取樣策略提供了新的視角,有望推動該類模型在更廣泛場景中的應用,例如程式碼生成、文本摘要等任務。相關研究成果已被2026年ICML會議接收,並發表於2026年7月。

展開要點與分析

文章情報

工程師進階

要點

  • 擴散語言模型透過並行解碼提升推理效率,但取樣策略依賴啟發式方法,存在手動調整和效能退化問題。
  • 作者提出用強化學習訓練取樣策略,將掩碼擴散取樣形式化為馬爾可夫決策過程。
  • 輕量策略基於單層Transformer,根據令牌置信度決定去掩碼操作。
  • 實驗證明,訓練策略在半自迴歸塊生成中匹配啟發式方法,在全擴散中表現更優。

要點與分析由自動化流程生成,可能有誤,請結合原始來源核實。