近年來，思維鏈在大模型訓(xùn)練和推理中愈發(fā)重要。近日，西湖大學(xué) MAPLE 實驗室齊國君教授團隊首次提出擴散式「發(fā)散思維鏈」—— 一種面向擴散語言模型的新型大模型推理范式。該方法將反向擴散過程中的每一步中間結(jié)果都看作大模型的一個「思考」步驟，然后利用基于結(jié)果的強化學(xué)習(xí)去優(yōu)化整個生成軌跡，最大化模型最終答案的正確率。不同于始終單向推理、線性生成的傳統(tǒng)思維鏈（CoT），擴散式「發(fā)散思維鏈」允許模型以任意順序非線性生成，且在生成過程中無需嚴(yán)格遵從語法結(jié)構(gòu)和可讀性要求，能夠鼓勵模型以更加發(fā)散、創(chuàng)造性的方法開展推理。

擴散式「發(fā)散思維鏈」目前已成功應(yīng)用于兩種具有代表性的擴散語言模型中。在連續(xù)時間擴散語言模型中，該方法可以直接優(yōu)化由模型輸出的得分函數(shù)所確定的策略分布；而在離散時間擴散語言模型中，團隊將預(yù)測不同掩碼 Token 的順序當(dāng)作模型決策的一部分，并基于 Plackett-Luce 模型設(shè)計去掩碼策略。據(jù)此，團隊成功訓(xùn)練有序掩碼生成擴散語言模型（Large Language Diffusion with Ordered Unmasking, LLaDOU）。實驗表明，僅用公開數(shù)據(jù)集和 16 張 H800，經(jīng)擴散式「發(fā)散思維鏈」增強后的模型即可在數(shù)學(xué)推理和代碼生成任務(wù)上超越現(xiàn)有擴散語言模型。

擴散式「發(fā)散思維鏈」對基礎(chǔ)大模型的訓(xùn)練與推理給出了重要啟示：傳統(tǒng)的自回歸思維鏈語言模型通過線性預(yù)測下一個 token 生成答案并非唯一的選擇范式。團隊的研究揭示了通過優(yōu)化 token 生成的順序進行非線性語言生成是發(fā)散式思維的重要特點，對于在生成過程中逐步構(gòu)建從早期概念要素的形成、到最終連接成具有完整想法和語法結(jié)構(gòu)的回答起到了關(guān)鍵作用。

相關(guān)研究成果已于 5 月 15 日公開。團隊注意在此后谷歌發(fā)布了 Gemini Diffusion 語言模型，因而期待強化「發(fā)散思維鏈」可以應(yīng)用到更多的擴散語言模型上成為標(biāo)準(zhǔn)訓(xùn)練過程的一部分。

• 論文標(biāo)題：Reinforcing the Diffusion Chain of Lateral Thought with Diffusion Language Models
• arXiv 地址：https://arxiv.org/abs/2505.10446
• GitHub 地址：https://github.com/maple-research-lab/LLaDOU

背景

近期，大型語言模型的推理能力引發(fā)了學(xué)術(shù)界的高度關(guān)注。一般而言，「推理」通常是指模型在生成最終應(yīng)答前所經(jīng)歷的系統(tǒng)性思考過程。當(dāng)前主流的大型語言模型普遍采用分步拆解問題的方法，構(gòu)建一種具有因果順序的線性思維鏈條，形成所謂的「思維鏈」推理范式。

值得注意的是，人類認(rèn)知過程中的思維構(gòu)建機制與此存在本質(zhì)差異。在構(gòu)思階段，人類思維往往呈現(xiàn)非線性的發(fā)散特征，能夠突破既有語言框架的約束，通過非線性、跳躍性的方式自發(fā)生成概念原型、詞匯單元及初始設(shè)想。隨著認(rèn)知加工的持續(xù)深化，這些離散的思維片段經(jīng)歷系統(tǒng)性整合與結(jié)構(gòu)化重組，最終形成邏輯連貫的完整表達體系。認(rèn)知科學(xué)領(lǐng)域?qū)⒋祟愃季S模式定義為「發(fā)散思維」（Lateral Thinking），顯著區(qū)別于傳統(tǒng)思維鏈所采用的線性推理模式。

為模擬這一思考過程，西湖大學(xué) MAPLE 實驗室齊國君教授團隊首次提出擴散式 「發(fā)散思維鏈」這一概念。如圖所示，模型的思考過程從一段不包含任何信息的掩碼序列開始，在思考過程中，模型會逐步生成推理所需要的關(guān)鍵信息，將掩碼轉(zhuǎn)換為具有實際語義內(nèi)涵的文字內(nèi)容，如數(shù)字和計算過程。最終，在整個擴散去噪流程結(jié)束后，模型將生成具有連貫語義內(nèi)涵且包含正確答案的文字回復(fù)。通過僅基于結(jié)果的強化學(xué)習(xí)訓(xùn)練，團隊鼓勵模型探索多樣化的、創(chuàng)造性的和非線性的思維路徑，最終得出正確的答案。

擴散式「發(fā)散思維鏈」

為了建模真實文本數(shù)據(jù)的分布，擴散語言模型構(gòu)建了一個從 t=0 到 t=T 的離散擴散過程，描述了在 t=0 處的一個未知的真實數(shù)據(jù)分布 p_data，是如何逐步演化為 t=T 時的一個已知的先驗分布 p_prior。而生成一段文本則通過反轉(zhuǎn)這一擴散過程來實現(xiàn)：首先從先驗分布 p_prior 中采樣 x_0，然后在一系列去噪時間 t_0:N 迭代去噪。在每一個去噪步驟 n 中，模型 θ 會估計時間 t_n 對應(yīng)的擴散分布，并從中采樣一個中間結(jié)果 x_n。隨著 n 的增加，擴散時間 t_n 逐漸減小，直到 t_N=0 時得到最終生成結(jié)果 x_N。

在這一過程中，為了得到最終的輸出 x_N，模型天然需要生成一系列中間結(jié)果 x_1:N-1。這一過程與「思維鏈」（Chain-of-Thought, CoT）技術(shù)相似。然而，與 CoT 采用線性因果推理不同，擴散過程中的模型能夠在思考過程中自由地生成任何有助于達到正確答案的中間內(nèi)容，更符合發(fā)散思維的概念 —— 即通過間接、具有探索性的方法解決問題。正因如此，團隊將由去噪過程中所有的中間結(jié)果組成的序列稱為擴散式「發(fā)散思維鏈」（Diffusion Chain of Lateral Thoughts, DCoLT），并通過強化學(xué)習(xí)算法優(yōu)化模型的這些中間擴散「推理」過程。

團隊采用基于最終結(jié)果監(jiān)督的強化學(xué)習(xí)方法：如果一條思維鏈推導(dǎo)出的最終答案 x_N 正確，就會予以激勵。具體而言，團隊生成整個思維鏈 x_1:N 的過程視為一個多步動作序列進行優(yōu)化。在第 n 步時，擴散模型會在所有可能的結(jié)果上定義一個輸出分布，即是模型用于采樣 x_n 的策略分布。獎勵信號 r 可以簡單地通過驗證最終生成結(jié)果的正確性得到。值得注意的是，團隊不會對推理過程的中間步驟設(shè)置任何顯式監(jiān)督，從而鼓勵模型探索多樣化、非線性的推理策略。

在下圖中，團隊以 GRPO 為例詳細(xì)闡述了算法訓(xùn)練框架。類似地，其他強化學(xué)習(xí)算法也可應(yīng)用于所提出的框架中。

連續(xù)時間擴散語言模型：DCoLT 強化的 SEDD

首先團隊考慮以 SEDD 為代表的連續(xù)時間擴散語言模型。這類模型通過如下線性常微分方程描述該演化過程。

其中，表示擴散過程中的瞬時轉(zhuǎn)移率矩陣，不妨首先考慮單個 token 的簡單情形。為了生成樣本，這一擴散過程存在一個對應(yīng)的反向過程，其中包括一個反向轉(zhuǎn)移率矩陣。

通過歐拉法數(shù)值求解，可以計算每一步的轉(zhuǎn)移概率，進而得到用于多步生成的迭代公式。此處團隊將簡化為 x_n 以避免標(biāo)記過于冗余。

在經(jīng)典的離散擴散模型 SEDD 模型中，SEDD 模型通過預(yù)測，來表示各個 token 的轉(zhuǎn)移概率。因此，團隊可以將公式中的替換為模型估計的，從而確定轉(zhuǎn)移概率。

擴展到整個序列時，其轉(zhuǎn)移概率可以看作所有 token 轉(zhuǎn)移概率的累乘，即可通過以下公式計算 DCoLT 生成過程中每一步動作對應(yīng)的采樣概率。

離散時間擴散語言模型：DCoLT 強化的 LLaDA

一些擴散語言模型直接在離散的時間步上執(zhí)行多步生成過程。對于這些模型，需要為每個離散步驟定義其輸出策略分布。在這其中，考慮最為常見的掩碼擴散語言模型。

以 LLaDA 模型為例：生成過程從一個完全掩碼序列開始，逐步去除掩碼直至生成最終文本。在每個生成步驟中，模型接收一個帶有掩碼的序列作為輸入，將其中部分掩碼預(yù)測為有實際含義的文本內(nèi)容。在整個生成過程進行時，掩碼的數(shù)量會逐漸減少，直到模型最終輸出完整的生成序列。

據(jù)此，團隊基于 LLaDA 設(shè)計了一種有序掩碼生成擴散語言模型，LLaDOU。他們將模型在每一步的動作拆解為兩部分：首先，確定本步中需要去除的掩碼集合，記為；其次，為這一部分中的每一個掩碼預(yù)測新的值，以獲得新的序列。

要確定為哪些掩碼 token 執(zhí)行去掩碼操作，我們可以用一個得分函數(shù)對所有掩碼 token 排序。為此，團隊設(shè)計了一個「去掩碼策略模塊」（Unmask Policy Module，UPM），該模塊在當(dāng)前擴散步驟 n 下，為第 i 個掩碼字符預(yù)測一個得分值。基于這些得分，團隊采用 Plackett–Luce 模型定義了一個策略，從中采樣一個由 K 個掩碼字符的列表。

具體而言，團隊首先根據(jù)預(yù)測的得分構(gòu)建一個多項分布，隨后以無放回的方式依次采樣出 K 個掩碼 token，這樣，得分較高的 token 有更大的可能性被首先取出，從而使序列中的掩碼得分值更傾向滿足非遞增排序關(guān)系，即：。

令表示第 n 步之后仍然保持掩碼的 token 集合，即滿足：。那么，采樣得到某個特定的去掩碼列表的概率可由下式計算所得。

具體而言，在第 n 步去噪過程中，UPM 會取 LLaDA 中最后一層的輸出特征作為模塊輸入，為每一個 token i 預(yù)測一個得分。UPM 僅包含一層 transformer，因此對模型計算量影響很小。此外，考慮到當(dāng)前的去噪步數(shù) n 和每個 token 的掩碼狀態(tài)同樣也對去掩碼策略十分重要，團隊將這些信息作為自適應(yīng)歸一化層編碼在 UPM 模塊中。為簡單起見，團隊將經(jīng) DCoLT 訓(xùn)練后，包含 UPM 的整個擴散語言模型記為 LLaDOU（LLaDA with Ordered Unmasking）。整體模型結(jié)構(gòu)如下圖所示：

而去掩碼的 token 集合一經(jīng)確定，模型就會根據(jù)詞匯表上的輸出分布預(yù)測它們相應(yīng)的 token 值，此即第二階段動作。在給定以及的情況下，生成的概率為：

綜上，從到的完整策略由這兩部分乘積共同決定：

從以上推導(dǎo)可以看出，某種意義上，LLaDOU 模型和基于 next token 預(yù)測的自回歸（auto-regressive) 語言模型并沒有本質(zhì)區(qū)別。兩者都是在給定了 prompt 和 context 作為前綴后，去預(yù)測后續(xù)的 token。區(qū)別僅在于，自回歸模型要求預(yù)測的是緊鄰的下一個 token；而 LLaDOU 模型允許通過一個 UPM 模塊，從所有可能的后續(xù)位置，選擇一個或多個 token 進行預(yù)測。后者相對于前者更加靈活，可以根據(jù)當(dāng)前生成的結(jié)果，打破語言自左到右的自然順序，在中間步驟，跳躍式地選擇合適的 token 進行生成。當(dāng)然，最終生成的完整結(jié)果，仍然滿足各種語言語法結(jié)構(gòu)的要求。

在同一時期，業(yè)界也推出了一些其他面向 diffusion model 的強化訓(xùn)練方法，如 d1 和 MMaDA。這些方法首先采樣得到生成結(jié)果以及對應(yīng)的獎勵值，然后對生成結(jié)果或問題部分再次進行隨機掩碼處理，以估算每個 token 的生成概率，用于強化訓(xùn)練。這種情況下，實際采樣生成的中間過程和計算概率時的再掩碼過程并不一致，可能導(dǎo)致所強化的再掩碼采樣過程并不是模型真正的采樣過程。不同于這些方法，團隊直接基于采樣過程中每一步所選中的 unmask token 計算概率，據(jù)此進行強化訓(xùn)練，保持訓(xùn)練和采樣過程一致。同時，更重要的是，團隊注意到每步如何選擇要 unmask 的 token 也是擴散語言模型采樣的關(guān)鍵步驟。基于此，本方法將 unmask token 生成的順序也作為強化學(xué)習(xí)所優(yōu)化策略的一部分，進一步提升擴散語言模型采樣的性能。

實驗結(jié)果

團隊基于兩個具有代表性的擴散語言模型 ——SEDD 和 LLaDA 開展實驗進行驗證。

首先，團隊基于 SEDD 模型，在數(shù)獨解題和數(shù)學(xué)推理兩個任務(wù)上與其他方法展開了公平對比。DCoLT 取得了比 CoT 和 DoT 更好的實驗結(jié)果。比如在 GSM8K-Aug 數(shù)據(jù)集上，同樣是使用 SEDD 模型，DCoLT 取得了 57.0% 準(zhǔn)確率，超越了 DoT，即使后者使用的訓(xùn)練數(shù)據(jù)中帶有逐步驟的詳細(xì) CoT 標(biāo)注。

而后，團隊在 LLaDA 8B 權(quán)重的基礎(chǔ)上訓(xùn)練 LLaDOU 模型，充分驗證了這一思考技術(shù)在數(shù)學(xué)推理和代碼生成任務(wù)上的能力。結(jié)果顯示，該技術(shù)顯著提升了模型對復(fù)雜數(shù)學(xué)邏輯問題的推理準(zhǔn)確率，和生成代碼的測試通過率。在相關(guān)的評測基準(zhǔn)上，LLaDOU 超越了其他擴散語言模型，取得了最好的性能。

在下圖中，團隊用不同顏色展示了同一回答中不同 token 的先后生成順序 —— 越淺的顏色代表 token 在更早的步數(shù)生成。可以看出，整個推理過程傾向于首先生成關(guān)鍵數(shù)字和計算符號，然后填充其他相關(guān)的文本內(nèi)容，逐漸滿足語法約束。

在這里，團隊也以視頻形式展示了 LLaDOU 在解決數(shù)學(xué)問題的完整生成過程。

總結(jié)

這篇文章介紹了由西湖大學(xué) MAPLE 實驗室提出的一種全新的大模型推理范式，擴散式「發(fā)散思維鏈」。該框架將反向擴散過程中的中間結(jié)果看作模型的推理過程，并將模型最終輸出結(jié)果的正確性作為獎勵開展強化學(xué)習(xí)訓(xùn)練，大幅提升了大模型的推理能力，在數(shù)學(xué)推理、代碼生成等任務(wù)上取得了超越其他擴散語言模型的性能。擴散式「發(fā)散思維鏈」這一理論打破了大模型推理過程的固有范式，為復(fù)雜推理問題提供了創(chuàng)新性的方法解決方案，值得我們進一步挖掘。