知識蒸餾是一種將更大的教師模型的知識轉移到更小的學生模型的方法，理想情況下可生成緊湊、易於部署的學生，且準確度與教師相當。知識蒸餾在預訓練設置中越來越受歡迎，但用於在監督式微調（Supervised Fine-Tuning，SFT）期間執行知識蒸餾的資源越來越少。
NVIDIA NeMo-Aligner 開源了一個在 SFT 期間使用知識蒸餾的實現，相較於標準 SFT，該實現的資料效率更高，準確性也更高 (Table 1)。
在表 1 中，SFT 是使用數學/代碼資料集執行的。使用知識蒸餾微調的模型版本在所有數學和代碼相關基準測試中均優於基準，即使僅執行 70%的訓練步驟也是如此。
NeMo-Aligner 中的知識蒸餾
在 SFT 期間，有許多方法可以從大型模型傳輸知識。最常見的方法是使用教師模型生成合成資料，我們稱之為 KD-SDG。然後，使用合成生成的資料微調學生模型。
還有一種開創性的方法，即訓練學生以匹配教師的輸出 logits。此方法在 Distilling the Knowledge in a Neural Network 中引入。我們將其稱為 KD-logit。
此方法利用跨類（稱為 暗知識 ）的知識，生成資訊更豐富的梯度信號。有關更多資訊，請參閱神經網路中的 Dark Knowledge。
在本文和 NeMo-Aligner 中，我們將重點介紹在 SFT 期間應用 KD-logit。
NeMo-Aligner 的離線 KD-logit 工作流包含以下關鍵步驟：
教師模型對訓練資料進行預測的預處理步驟。教師模型的 logits 添加到訓練資料中。
這是一個訓練步驟，其中對學生進行了訓練，使其 logits 與教師的 logits 相匹配。
只需緩存一次教師的 logits。與在訓練時動態計算教師邏輯相比，此方法具有以下優勢：
節省 內存： 您不必同時在 GPU 上加載教師和學生模型。
加快訓練速度： 您不必等待老師在訓練期間做出預測。
但是，將所有教師的 logits 保存到磁盤可能需要大量內存。為節省內存，我們僅將教師的最高 K logits 保存到磁盤，其中 K 是從業者選擇的超參數。
K 的值越大，學生可以從教師那裡學習的細粒度資訊越多，但內存壓力就越大。在實踐中，通常選擇 K 值在 100 左右，這比典型的詞彙量小。
將教師 logits 添加到資料集後，學生被訓練以匹配教師的 top- K logits。具體來說，知識蒸餾損失函數等於 K 學生和教師 logits 之間的前向 KL 差異：
L^{kd} (p^S, p^T) = sum_{k=1}^K p_k^T(log p_k^T - log p_k^S)
此損失函數與 Vanilla SFT 交叉熵損失函數結合使用，以生成最終訓練目標，其中 lambda 控制 SFT 損失項相對於 KD 損失項的強度：
L(p^S, p^T, y) = L^{kd} (p^S, p^T) + lambda L^{sft}(p^S, y)
結果
表 1 顯示，與 Vanilla SFT 相比，使用知識蒸餾目標微調模型可獲得更高的準確性和所需的訓練令牌。我們使用 基礎 Nemotron-4 15B 學生模型 和 微調的 Nemotron-4 340B 教師模型 進行實驗。
用於 SFT 的資料集是使用以下論文中描述的技術生成的組合：
數學資料集：OpenMathInstruct-2： 利用海量開源指令資料加速數學 AI (使用 Nemotron-4 340B，而非 Llama-3.1-405B-Instruct)
代碼資料集：Genetic Instruction：Scaling up Synthetic Generation of Coding Instructions for Large Language Models
資料集的數學和代碼部分均使用合成資料生成。這些實驗設置了 K=100 和 lambda=0.1。
在相同數量的訓練步驟中，使用聯合知識蒸餾和 SFT 目標微調的模型在七個評估指標中的六個方面的表現優於 SFT 基準。特別是，我們看到 HumanEval、MBPP 和 MATH 基準測試有了顯著改進，這些基準用於衡量編碼和數學推理技能。在評估各種語言理解任務的 MMLU 上，KD 微調模型的表現至少與零樣本設置中的基準相當，並且在 5 鏡頭設置中優於基準。
KD-finetuned Nemotron-4 僅使用 70% 的訓練令牌，但在相同的六個評估指標上，其性能仍然優於 Vanilla SFT 模型。
結束語
這些結果具有兩個重要含義。首先，我們已證明知識蒸餾（Knowledge Distillation）可用於提高微調模型的準確性。這在資料稀缺的設置中特別有用，因為需要更少的訓練令牌才能實現良好的準確性。
其次，我們已經證明 KD-logit 可以與您的 SDG 資料結合使用，以實現復合優勢。
WeChat
Profile