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

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

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

1.模型縮放法則研究現(xiàn)狀：近年來的研究揭示了語言模型性能與模型規(guī)模、訓(xùn)練數(shù)據(jù)量之間的關(guān)系，但在蒸餾場景下的縮放法則尚未有系統(tǒng)研究。

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

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

1、知識蒸餾

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

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

學(xué)生模型的總損失函數(shù)結(jié)合了下一個token預(yù)測損失和知識蒸餾損失：

• λ：混合系數(shù)，控制知識蒸餾損失和下一個token預(yù)測損失的權(quán)重。
• ：Z-loss的權(quán)重，用于提高訓(xùn)練穩(wěn)定性。

2、實驗設(shè)計

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

固定M教師/學(xué)生IsoFLOP配置

固定教師模型規(guī)模和訓(xùn)練數(shù)據(jù)量，改變學(xué)生模型的規(guī)模和蒸餾數(shù)據(jù)量。目的是研究在固定教師模型條件下，學(xué)生模型性能如何隨其規(guī)模和蒸餾數(shù)據(jù)量變化。通過這種方式，可以確定學(xué)生模型在不同規(guī)模和數(shù)據(jù)量下的最優(yōu)性能。

從實驗結(jié)果可以觀察到：

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

IsoFLOP教師/固定M學(xué)生配置

固定學(xué)生模型規(guī)模和蒸餾數(shù)據(jù)量，改變教師模型的規(guī)模和訓(xùn)練數(shù)據(jù)量。目的是研究教師模型性能如何影響學(xué)生模型的最終性能。通過這種方式，可以確定最優(yōu)的教師模型規(guī)模和訓(xùn)練數(shù)據(jù)量，以最大化學(xué)生模型的性能。

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

固定M教師/固定M學(xué)生配置

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

同時變化教師模型和學(xué)生模型的規(guī)模：教師模型和學(xué)生模型的參數(shù)規(guī)模和以及它們的訓(xùn)練數(shù)據(jù)量和同時變化，但保持固定的比例和。

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

2、Distillation Scaling Laws

先介紹一個預(yù)備知識，已有大量理論和實驗證明了交叉熵是可預(yù)測的，在監(jiān)督學(xué)習(xí)場景下，模型的性能（交叉熵?fù)p失L）在模型規(guī)模（參數(shù)數(shù)量 N）和訓(xùn)練數(shù)據(jù)量（token數(shù)量 D）遵循冪律關(guān)系。

可以使用教師的模型參數(shù)量和訓(xùn)練數(shù)據(jù)量來擬合監(jiān)督縮放定律。

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

• ：教師模型的驗證交叉熵?fù)p失。
• ：學(xué)生模型的驗證交叉熵?fù)p失。
• ：如果學(xué)生模型以監(jiān)督學(xué)習(xí)方式訓(xùn)練，其預(yù)期的驗證交叉熵?fù)p失。
• ：學(xué)生模型的非嵌入?yún)?shù)數(shù)量。
• ：學(xué)生模型的蒸餾token數(shù)量。
• 是通過實驗擬合的系數(shù)

簡單介紹下公式的組成：

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

最后兩個公式的關(guān)鍵系數(shù)可以參考下表：

結(jié)果可以看到，監(jiān)督和蒸餾縮放定律在觀察到的數(shù)據(jù)上擬合得非常好，相對預(yù)測誤差小于1%。反映了蒸餾縮放定律的有效性，因此基于教師和學(xué)生的模型規(guī)模和數(shù)據(jù)規(guī)?？梢曰ハ鄬崿F(xiàn)最優(yōu)模型的選擇。

3、Distillation vs Supervised

為了解蒸餾在何時有益，下圖是在固定計算資源下，蒸餾和監(jiān)督學(xué)習(xí)的性能對比。結(jié)果展示，監(jiān)督學(xué)習(xí)在給定足夠的學(xué)生計算或數(shù)據(jù)時總是優(yōu)于蒸餾。對于適中的數(shù)據(jù)預(yù)算，蒸餾是有利的，然而，當(dāng)有大量數(shù)據(jù)可用時，監(jiān)督學(xué)習(xí)優(yōu)于蒸餾。

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

4、教師模型選擇

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

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

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

5、計算最優(yōu)蒸餾

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

從圖8中可以看到：

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

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

6、總結(jié)

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

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