Large Language Models to Diffusion Finetuning

目錄

1. 概述

2. 研究背景

3. 方法

3.1 用于 LM 微調的高斯擴散

3.2 架構

4. 主要實驗結果

5. 結論

---

1. 概述

本文提出了一種新的微調方法——LM to Diffusion (L2D)，旨在賦予預訓練的大語言模型（Large Language Models, LLMs）測試時可擴展計算能力。

研究發現，通過在推理過程中增加擴散步數，模型的準確性可以單調增加，從而提高在數學、編程和推理等下游任務中的性能。

此外，該方法可以結合強大的引導技術，優化計算資源的動態分配，同時保持語言模型的單步生成能力。

本文的主要貢獻包括：

• 提出 L2D 微調方法，結合語言模型的自回歸能力和擴散模型的多步推理能力，利用語言模型的預測來更新擴散狀態。
• 證明 L2D 可擴展計算能力，通過增加計算資源提高推理質量。
• 提升 LMs 在數學、編程和推理任務上的表現，同時保留其單步生成能力。
• 實現高效微調：L2D 不修改模型原始參數，僅需修改少量參數，即可在多個任務上取得顯著提升。??

2. 研究背景

當前自回歸（autoregressive）大語言模型在推理能力上存在一定的局限性，主要體現在：

• 無法根據任務的復雜度調整計算資源。
• 生成結果的質量受限于單步預測，而無法進行多步優化。

擴散模型（Diffusion Models）在視覺領域的成功啟發了研究人員嘗試將其應用于語言建模，以引入多步推理的能力。然而，由于離散文本數據的特殊性，現有的擴散語言模型在性能上落后于自回歸模型。

3. 方法

3.1 用于 LM 微調的高斯擴散

L2D 通過將 LMs 視為單步擴散模型，并在其基礎上擴展多步推理能力，從而增強 LLM 的推理能力，而不會影響其原始的單步生成能力。

語言建模在有限詞匯表 V 上定義的目標分布 p_1 上運行。給定一個由標簽 y 索引的 token x_1，該 token 與來自目標數據分布 p_1 的先前 token c 的上下文一起采樣，我們的擴散損失公式如下：

使用公式 3 訓練擴散模型可以解釋為標準的下一個 token 預測，其中為模型提供了一個額外的擴散 token x_t，包含有關目標 y 的一定程度的知識，范圍從無信息（t = 0）到完美信息（t = 1）。

• 當 t = 0 時，LM 本質上是用與 L2D 相同的預測目標進行訓練的，此時 x_0?與目標 y 完全不相關。
• 遵循算法 1 的推理涉及從模型的 logit 中迭代采樣越來越準確的下一個 token ?x，直到采樣預算 T。
• 傳統的 LM 推理可以再次被視為此過程的特例（T = 1），其中僅使用模型的第一個樣本來預測 y。

這些設計選擇的目的是 L2D 旨在通過微調方法擴展預先訓練的 LM，而不是從頭開始學習新模型。

雖然從一開始就完全采用擴散訓練可能看起來更為普遍，但我們認為這可能會失去傳統自回歸建模所固有的一些訓練可擴展性和強大的歸納偏差。正是這些偏差使傳統自回歸建模在語言領域得到廣泛確立。

3.2 架構

L2D 的實現設計為預訓練 transformer 的模塊化擴展，以有效利用擴散的多步擴展能力，同時保留其原始的單步生成能力。

L2D 在其架構中引入了一條 并行擴散路徑，其中傳播擴散 token x_t 的隱藏表示，僅在最后一層影響凍結的主 LM 路徑。

• L2D 使用 transformer 架構以及與主路徑 f_{θ_l} 相同數量的塊來實現擴散路徑 f_{θ_d} 。為了充分利用預訓練 LM 的知識，擴散路徑中的所有層也都使用來自 θ_l 的權重進行初始化。
• 擴散路徑中的 transformer 由一系列殘差 MLP 和交叉注意模塊組成。雖然 MLP 模塊遵循與 f_{θ_l} 中相應模塊相同的結構，但交叉注意模塊專門參數化 query 和輸出線性層。具體而言，在交叉注意期間，目標 token y_k 的擴散 token x^k_t 會關注從 f_{θ_l} 中相應的自注意模塊計算出的所有先前的 key 和 value。
• 最終，只在所有塊之后，即 LM 的線性頭之前，將 f_θ 中處理的信息整合回主路徑。具體來說， 將這兩條路徑與元素加權和 f_{θ_l} + w_d·f_{θ_d} 合并，其中擴散 token x^k_t 的重新縮放潛變量被添加到前一個 token x^{k?1} 的潛變量中。

4. 主要實驗結果

研究者在多個基準測試（benchmark）上對 L2D 進行了實驗，包括：數學任務（GSM8K、MATH），編程任務（HumanEval、MBPP），知識推理任務（MMLU、MMLU-Pro）

增加擴散步數（計算量）能顯著提高推理能力，但性能提升在一定步數后趨于穩定。?

L2D 在所有測試集上均提升了語言模型的性能，特別是在數學和編程任務上，表現尤為明顯。

L2D 兼容現有的微調方法（如 LoRA 和全參數微調），且比傳統微調方法更高效。

引導技術（Classifier-Free Guidance） 使 L2D 在特定任務上的表現進一步提升，允許用戶在生成過程中指定任務需求。

5. 結論

L2D 通過結合擴散模型的推理增強能力，使大語言模型能夠更高效地利用計算資源，提高復雜任務的推理能力。這項研究為未來的大語言模型微調提供了一條新的方向，通過引入擴散計算框架，使模型能夠智能調整計算資源，提高推理質量，在多個領域展現出了良好的潛力。

【一個比較關心且重要的問題，新添加的擴散路徑導致的額外計算開銷是多少？原文并未提及】

論文地址：https://arxiv.org/abs/2501.15781

進 Q 學術交流群：922230617