謔!大語言擴散模型來了,何必只預測下一個token | 人大高瓴&螞蟻
明敏 發自 凹非寺
量子位 | 公眾號 QbitAI
用擴散模型替代自回歸,大模型的逆詛咒有解了!
人大高瓴人工智能研究院、螞蟻共同提出LLaDA(a Large Language Diffusion with mAsking)。
LLaDA-8B在上下文學習方面與LLaMA3-8B能力相當,而且在反轉詩歌任務中超越GPT-4o。
在大語言模型領域,反轉詩歌是一個特殊任務,它用來評估模型在處理語言模型的雙向依賴關係和邏輯推理能力。
比如讓大模型寫出「一行白鷺上青天」的上一句。
通常情況,自回歸模型(如GPT)根據下文推斷上文的表現上總是不夠好。這是因為自回歸模型的原理就是利用序列中前面的元素來預測當前元素,即預測下一個token。

而LLaDA是基於擴散模型的雙向模型,天然能夠更好捕捉文本的雙向依賴關係。
作者在摘要中表示:LLaDA挑戰了LLMs關鍵能力與自回歸模型之間的固有聯繫。
這些研究也引發了不少討論。
有人提出:
我們正在重構掩碼語言模型建模?

RAG和嵌入式相似性搜索上,這種範式也可能表現更好?

值得一提的是,LLaDA僅用13萬H800GPU時訓練了2.3萬億token語料,然後對450萬對token進行SFT。
正向掩碼+反向預測
論文核心提出了一個問題:自回歸是否是實現LLMs智能的唯一路徑?
畢竟自回歸範式的LLMs目前仍存在諸多弊端,比如逐個生成token的機制導致計算成果很高,從左到右建模限制了逆推理任務中的性能。
這都限制了了LLMs處理更長、更複雜任務的能力。
為此,他們提出了LLaDA。通過正向掩碼和反向預測機制,讓模型更好捕捉文本的雙向依賴關係。

研究採用標準的數據準備、預訓練、監督微調(SFT)和評估流程,將LLaDA擴展到80億參數。
在2.3萬億token上從零開始預訓練,使用13萬H800 GPU時,隨後在450萬對數據上進行監督微調。
在語言理解、數學、代碼和中文等多樣化任務中,表現如下:
強大可擴展性:LLaDA 能夠有效擴展到10²³ FLOPs計算資源上,在六個任務(例如MMLU和GSM8K)上,與在相同數據上訓練的自建自回歸基線模型結果相當。

上下文學習:值得注意的是,LLaDA-8B 在幾乎所有 15 個標準的零樣本/少樣本學習任務上都超越了 LLaMA2-7B,並且與 LLaMA3-8B表現相當。
指令遵循:LLaDA在SFT後顯著增強了指令遵循能力,這在多輪對話等案例研究中得到了展示。
反轉推理:LLaDA有效地打破了反轉詛咒,在正向和反轉任務上表現一致。特別是在反轉詩歌完成任務中,LLaDA 的表現優於 GPT-4o。

LLaDA使用Transformer架構作為掩碼預測器。與自回歸模型不同,LLaDA的transformer不使用因果掩碼(Causal Mask),因此它可以同時看到輸入序列中的所有token。
模型參數量與傳統大語言模型(如GPT)相當,但架構細節(如多頭注意力的設置)略有不同,以適應掩碼預測任務。
其正向掩碼過程如下:
LLaDA採用隨機掩碼機制,對一個輸入序列x0,模型會隨機選擇一定比例的標記進行掩碼(masking),生成部分掩碼的序列xt。
每個token被掩碼的概率為t,其中t是從[0,1]中均勻采樣的。這與傳統的固定掩碼比例(如BERT中的15%)不同,LLaDA的隨機掩碼機制在大規模數據上表現出更好的性能。
模型的目標是學習一個掩碼預測器,能夠根據部分掩碼的序列xt預測出被掩碼的token。訓練時,模型只對被掩碼的token計算損失。

其中1[·]是指示函數,表示只對被掩碼的token計算損失。
在SFT階段,LLaDA使用監督數據(如對話對、指令-響應對)進一步優化模型,使其在特定任務上表現更好。
對於每個任務,模型會根據任務數據的特點進行微調。例如在對話生成式任務中,模型會學習如何根據給定對話歷史生成合適響應。
在SFT階段,模型會根據任務數據的特點選擇性地掩碼響應部分token,這使得模型能夠更好地學習任務相關的模式。
推理部分,在生成任務中,LLaDA通過反向采樣過程生成文本。從一個完全掩碼的序列開始,逐步預測出被掩碼的token,直到生成完整的文本。
采樣過程中,LLaDA採用多種策略(如隨機重掩碼、低置信度重掩碼、半自回歸重掩碼)來平衡生成效率和質量。
在條件概率評估任務中,LLaDA會根據給定的提示(prompt)和部分掩碼的響應(response)來評估模型的條件概率。這使得LLaDA能夠在各種基準任務上進行性能評估。

預訓練LLM在不同基準上的表現如下。

後訓練後在不同benchmark上的表現如下。其中LLaDA只進行了SFT,其他模型有進行額外的強化學習對齊。

在反轉詩歌任務中,LLaDA超越了GPT-4o。

在多輪對話任務中LLaDA的表現如下,較深顏色表示采樣後期階段預測的token,較淺顏色表示在采樣早期預測的token。

網民:期待能被真正用起來
研究團隊同時放出了一些LLaDA的實際表現。
可以解決普通的數學推理問題。

編程問題也OK。

有國外網民表示:這肯定會推動中國AI研究更加關注小模型。不過也不代表他們放棄scaling。

同時也有人評價說,這或許可以開啟一些混合模型的可能。

以及有人提及Meta也有過類似的工作,將transformer和diffusion相結合。

當然也有人關心,此前也提出了不少超越Transformer的架構,但是它們都還沒有被學術界/工業界真正採納。
讓我們期待後續吧。

本項研究由人大高瓴人工智能學院與螞蟻集團共同帶來。通訊作者為李崇軒,他目前為人大高瓴人工智能學院準聘副教授,目前focuses的方向為深度生成模型,瞭解現有模型的能力和局限,設計有效且可擴展的下一代模型。
論文地址:
https://arxiv.org/abs/2502.09992
項目主頁:
https://ml-gsai.github.io/LLaDA-demo/