初次閱讀這篇文章，感到非常震撼。在DeepSeek將知識蒸餾方法帶入大眾視野后，Apple與牛津大學的研究人員迅速提出了蒸餾縮放定律，并已于2月28日完成了所有實驗及一篇長達67頁的論文上傳至arXiv。這種效率和深度，無疑展示了大公司的研究實力。這種高效的研究速度和成果展示，再次彰顯了大廠在科研領域的強大實力。

作為算力有限的獨立研究者，正好借此深入研究知識蒸餾技術，如何利用蒸餾方法最大化地利用有限的計算資源。通過學習這些前沿研究，或許能夠找到一些優化的思路和方法，從而在有限的條件下實現更好的模型性能。

簡單看下研究動機，總結為以下幾點：

1.模型縮放法則研究現狀：近年來的研究揭示了語言模型性能與模型規模、訓練數據量之間的關系，但在蒸餾場景下的縮放法則尚未有系統研究。

2.模型推理成本問題：隨著語言模型規模增大，推理成本顯著增加，研究如何在不損失性能的前提下降低推理成本成為重要課題。

3. 蒸餾的效率與性能：蒸餾理論上可降低推理成本，然而，學界對蒸餾方法缺乏共識，尤其是在如何合理分配計算資源以產生最強大的模型方面，仍存在很大的不確定性。

1、知識蒸餾

知識蒸餾的核心思想是將大型教師模型的知識遷移到小型學生模型中。具體來說，教師模型的預測分布被用作學生模型的目標分布，學生模型通過最小化與教師模型的Kullback-Leibler散度（KL Divergence）來學習：

其中，和分別是教師模型和學生模型的輸出logits，是蒸餾溫度，用于控制教師模型輸出的概率分布的熵。

學生模型的總損失函數結合了下一個token預測損失和知識蒸餾損失：

• λ：混合系數，控制知識蒸餾損失和下一個token預測損失的權重。
• ：Z-loss的權重，用于提高訓練穩定性。

2、實驗設計

為了擬合蒸餾縮放法則，設計了以下三種實驗方案來生成縮放定律的擬合數據：

固定M教師/學生IsoFLOP配置

固定教師模型規模和訓練數據量，改變學生模型的規模和蒸餾數據量。目的是研究在固定教師模型條件下，學生模型性能如何隨其規模和蒸餾數據量變化。通過這種方式，可以確定學生模型在不同規模和數據量下的最優性能。

從實驗結果可以觀察到：

• 在算力較大的情況下，學生模型參數規模越大，其損失函數越小，而且教師模型規模越大，這種趨勢越明顯。
• 在學生和教師模型規模確定的情況下，算力越多，學生模型的確表現越好，這也很符合常理。
• 在算力較小的情況下，模型性能會隨著模型規模先提高，后減弱，這里很容易理解，較大的模型在較小的算力下沒有得到充分訓練。
• 在特殊情況下，學生模型能夠優于教師模型，表現出泛化能力，個人猜測可能是教師模型訓練欠擬合。

IsoFLOP教師/固定M學生配置

固定學生模型規模和蒸餾數據量，改變教師模型的規模和訓練數據量。目的是研究教師模型性能如何影響學生模型的最終性能。通過這種方式，可以確定最優的教師模型規模和訓練數據量，以最大化學生模型的性能。

結果可以看到，并非教師模型參數越大，學生模型的交叉熵越低。表明教師模型的性能需要與學生模型的能力相匹配，才能達到最佳蒸餾效果。

固定M教師/固定M學生配置

在固定比例下，教師模型和學生模型的性能如何隨其規模變化。這種配置雖然不是擬合縮放法則所必需的，但它提供了一個更廣泛的視角，展示了學生模型性能在不同教師模型條件下的變化。

同時變化教師模型和學生模型的規模：教師模型和學生模型的參數規模和以及它們的訓練數據量和同時變化，但保持固定的比例和。

在圖4中，觀察到容量差距現象，教師模型的規模對學生的性能有顯著影響，但存在一個最優的教師模型規模（但并不一定是最大的模型），使得學生模型的性能最低。即提高教師的表現并不總是提高學生的表現，甚至最終會降低學生的表現。

2、Distillation Scaling Laws

先介紹一個預備知識，已有大量理論和實驗證明了交叉熵是可預測的，在監督學習場景下，模型的性能（交叉熵損失L）在模型規模（參數數量 N）和訓練數據量（token數量 D）遵循冪律關系。

可以使用教師的模型參數量和訓練數據量來擬合監督縮放定律。

基于前面實驗結果，研究人員建立了蒸餾擴展定律的數學模型，用以下公式來表示：

• ：教師模型的驗證交叉熵損失。
• ：學生模型的驗證交叉熵損失。
• ：如果學生模型以監督學習方式訓練，其預期的驗證交叉熵損失。
• ：學生模型的非嵌入參數數量。
• ：學生模型的蒸餾token數量。
• 是通過實驗擬合的系數

簡單介紹下公式的組成：

• 教師模型的交叉熵：教師模型的性能通過其交叉熵損失來衡量。較低的表示教師模型具有更好的性能，教師模型的性能直接影響學生模型的最終性能。
• 容量差距（Capacity Gap）：當教師模型的性能遠優于學生模型時，學生模型的性能可能會下降。這被稱為容量差距現象。公式中的和用于捕捉這種現象。具體來說，當遠小于時，會變得很大，導致學生模型的性能下降。
• 學生模型的性能極限：當蒸餾數據量足夠大時，學生模型的性能會趨近于教師模型的性能。公式中的用于描述學生模型在不同規模和數據量下的性能變化。
• 學生模型的參數數量和蒸餾數據量：學生模型的參數數量和蒸餾數據量對學生模型的性能有顯著影響。公式中的α′和分別描述學生模型的參數數量和蒸餾數據量對性能的影響。
• 指數項

最后兩個公式的關鍵系數可以參考下表：

結果可以看到，監督和蒸餾縮放定律在觀察到的數據上擬合得非常好，相對預測誤差小于1%。反映了蒸餾縮放定律的有效性，因此基于教師和學生的模型規模和數據規模可以互相實現最優模型的選擇。

3、Distillation vs Supervised

為了解蒸餾在何時有益，下圖是在固定計算資源下，蒸餾和監督學習的性能對比。結果展示，監督學習在給定足夠的學生計算或數據時總是優于蒸餾。對于適中的數據預算，蒸餾是有利的，然而，當有大量數據可用時，監督學習優于蒸餾。

總結即，在計算資源有限的情況下，蒸餾通常比監督學習更高效。這是因為蒸餾可以通過教師模型的指導，更快地學習到有效的特征表示，從而在較少的計算資源下達到更好的性能。

4、教師模型選擇

• 學習信號強度：不同大小的教師模型可能提供不同的學習信號強度，這通常通過交叉熵損失來衡量。較大的教師模型可能會提供更強的學習信號（較低的交叉熵），從而有助于學生模型更好地學習。
• 成本增加：使用較大的教師模型會帶來更高的成本，這是因為需要計算教師模型的logits。這意味著更大的教師模型不僅自身訓練成本更高，在用于蒸餾時也會消耗更多的計算資源。

圖7展示了在不同蒸餾數據預算下學生模型的交叉熵損失變化情況。結果顯示，最優教師損失（用紅線表示）隨著學生大小的增大而按冪律遞減，直到學生的損失與最優教師損失相匹配為止。這驗證了最佳教師規模幾乎始終與學生規模呈線性比例關系。

圖15可以看到，隨著蒸餾數據量的增加，最優教師模型的交叉熵X逐漸降低，但這種降低的速率會隨著的增加而減緩。因此可以得到：在計算資源有限的情況下，選擇較小的教師模型可以減少推理成本，同時仍然能夠提供有效的學習信號給學生模型。

5、計算最優蒸餾

計算最優蒸餾的目標是，在給定計算預算C的情況下，確定如何生成具有最低交叉熵的所需大小的學生模型。具體來說，需要找到最優的學生訓練數據量、教師模型大小和教師訓練數據量，使得在滿足計算預算約束的情況下，學生的交叉熵最小化。

從圖8中可以看到：

• 監督學習在足夠的計算預算下總是匹配最優蒸餾：監督學習總是在某個總計算預算下匹配最佳情況蒸餾設置。這意味著，當計算預算足夠大時，監督學習可以達到與蒸餾相同的性能。
• 當教師訓練包含在計算中時，學生交叉熵總是高于監督設置：這意味著，如果你的唯一目標是生成具有目標大小的最佳模型，并且你沒有訪問教師的權限，那么你應該選擇監督學習，而不是訓練教師然后進行蒸餾。相反，如果目的是蒸餾出一系列模型，或使用教師作為服務模型，那么蒸餾可能比監督學習在計算上更有益。
• 較小的模型更可能從監督預訓練中受益，而較大的模型更可能從蒸餾中受益：較小的模型在較大的計算預算下更傾向于監督學習，而較大的模型在較大的計算預算下更傾向于蒸餾。

圖9展示了隨著計算預算的變化，最優教師大小和訓練數據量的變化趨勢。學生和教師模型的token按冪律縮放，學生token的增長速度更快。最佳教師模型規模最初增加，直到略大于學生，之后趨于穩定。這是因為使用大型教師模型進行推理成本很高，并且隨著學生token數量的增加，過擬合教師模型更有效。

6、總結

Distillation Scaling Laws（蒸餾縮放定律）為知識蒸餾技術提供了理論基礎，優化了教師與學生模型間的資源分配，使得基于計算預算更科學地預測模型性能成為可能，同時降低了大規模應用蒸餾技術的風險，并揭示了“能力差距”現象的本質，為解決蒸餾過程中的挑戰提供了方向。

知識蒸餾的潛力是在有限計算成本下能夠訓練出性能更優的模型，因此在得到蒸餾縮放定律后，可以指導有限資源下的許多場景。