用大模型“蒸餾”小模型，有新招了！

甚至能在不同類型和架構的LLMs（大語言模型）上達到新SOTA。

這就是來自中科大、騰訊優圖實驗室提出的一種基于Sinkhorn距離的知識蒸餾方法，能把大的、復雜的教師模型的知識“蒸餾”到小的、簡單的學生模型中，從而讓小模型也能像大模型一樣工作。

之所以提出新方法，主要是現有的知識蒸餾（KD）方法都有各自的局限性：

當兩個模型的輸出差異較大時，它們就不太管用了。

• KL散度：會導致學生模型的輸出變得過于平滑，失去了區分性；
• RKL散度：會讓學生的輸出變得太簡單，不能很好地模仿教師模型；
• JS散度：會讓學生模型低估稀有事件的概率；

而基于Sinkhorn距離的新方法能更準確地衡量和縮小教師模型和學生模型之間的差異，從而提高了學生模型的性能。

此外，研究還提出了一種基于批量的重構方法，從而在高維空間中捕捉跨樣本分布的幾何復雜性。

最終，通過在兩個流行的自然語言處理測試集（GLUE和SuperGLUE）上測試，新方法在編碼器、編碼器-解碼器以及解碼器等不同架構的所有類型LLMs上均優于當前的最先進方法。

研究背景

知識蒸餾的提出是為了通過對齊教師模型的軟目標（例如輸出logits和中間層表示）來將教師模型內在固有的知識傳遞給學生模型。

給定訓練集中的一個樣本x_i及其真實標簽??_?? ∈ ?^??，來自教師模型??_??和學生模型??_??的輸出logits ??_?? ∈ ?^??和??_?? ∈ ?^??可以由以下式子得到：

其中為softmax函數， τ是溫度參數, d是輸出logits的維度?；趌ogit的知識蒸餾的目標是σ_Τ最小化測量散度J（??_??，??_??）以實現知識傳遞。

研究動機

現有研究已經嘗試使用Kullback-Leibler（KL）散度、反Kullback-Leibler（RKL）散度和Jensen-Shannon（JS）散度。

所有這些度量都可以被視為f-散度度量的變體，而f-散度度量在量化缺乏實質性交集的任何兩個分布時都存在明顯局限性。

此外，每種度量都有其自身的缺陷：

KL蒸餾會導致模式平均，使學生學習到一個過于平滑的分布，涵蓋了教師的整個支撐集；

RKL會引起模式塌陷，學生僅關注教師分布中高概率的顯著區域，而忽視了其余部分；

JS蒸餾會產生模式低估，由于懲罰不足，學生會低估稀有事件的概率。

為了解決傳統散度度量的問題，研究做出了以下貢獻：

1. 提出了一種知識蒸餾方法SinKD，采用Sinkhorn距離作為散度度量。它不僅解決了KL、RKL和JS散度在極端場景下的局限性，而且避免了計算Wasserstein距離的負擔。
2. 深入探討了Sinkhorn距離的性質，并將SinKD重新reformulated為batch-wise OT，擴展了它在NLP任務中的適用性。
3. 通過大量的可比性、有效性和泛化性實驗證明了SinKD相較于目前最先進的方法的優越性。并為實際應用提供了使用SinKD進行蒸餾的實用指導方針。

傳統散度度量的缺陷

首先，KL散度是不對稱的，表現為J_KL（??_??，??_??）≠ J_KL（??_??，??_??），這一性質違反了距離度量的對稱性特性，從而引入了一些不一致性。

其次，由于使用KL損失進行優化，學生模型試圖對教師模型的多模態分布進行平均化，從而導致對這些模式的擬合不足。這被稱為“模式平均問題”（mode-averaging problem）。

因此，學生模型無法捕獲數據中的所有關鍵模式，最終影響模型性能。

第三，KL散度對應的是一個非平滑函數，這為優化過程帶來了挑戰。

與KL散度一樣，具有內在的不對稱性，從而導致在捕捉分布差異時出現不一致性。

此外，優化的學生模型傾向于僅關注教師分布中概率較高的事件，這被稱為“模式崩塌問題”（mode-collapsing）。

如果教師對某個事件賦予零概率，學生模型也被迫做出相同的預測。

其中m_?? = 1/2（??_??+??_??）受制于非平滑性，JS損失在優化過程中面臨挑戰。

另外，由于JS損失在低概率區域的匹配上懲罰不足，學生模型可能會過度低估稀有事件的概率。

對于分布之間重疊較少甚至完全不重疊的情況退化為常數時，還存在梯度消失的風險。

最優傳輸距離的優勢

Wasserstein距離通過求解兩個分布之間的最優傳輸計劃來量化它們的差異。

直觀地看，它可以被認為是將一個分布（即學生的logits分布）轉換為另一個分布（即教師的logits分布）所需的最小“代價”，其中“代價”可以定義為被移動的質量與移動距離的乘積。

與傳統的散度度量相比，Wasserstein距離作為蒸餾的成本函數更為合理，因為它不依賴于對被測量分布的隱式假設。此外，它幾乎處處可微，從而便于優化。

另外，現有的散度度量只能獨立處理每個樣本對，進行逐一logit的匹配，對于一批樣本，這些方法無法定位來自同一樣本的教師和學生的logits對，從而無法實現整體距離的最小化。

由于計算Sinkhorn距離的過程可以實現來自同一樣本的兩個輸出之間的精確逐元素匹配，研究提出了“批量化”的SinKD方法（batchified SinKD）。

通過這種方式，即使通過低維觀測，也能夠捕捉復雜且隱式分布的幾何結構。

方法介紹

這里簡要介紹SinKD的核心方法，詳細推導過程可以參閱原論文。

批量重構的Sinkhorn距離

對于本問題，Wasserstein距離的定義如下：

其中，

Wasserstein距離本身在解析計算上存在困難，其計算成本對于蒸餾大型語言模型來說高得難以承受。

在這種情況下，研究使用Sinkhorn距離作為一種高效的近似方法。它不僅保留了Wasserstein距離的所有優點，同時也大大緩解了其在在線蒸餾中所面臨的成本問題。

Sinkhorn距離的定義如下：

逐樣本蒸餾將每個實例獨立處理，但忽略了一個批次樣本中的整體趨勢。

研究摒棄了僅在每對教師-學生樣本對上工作的逐樣本知識蒸餾方法，轉而在教師和學生樣本組上執行知識蒸餾。

一個包含b個樣本的批次會整體參與散度度量。通過批量重構，這種方法有效地增加了“觀測”空間的維度，特別是在d遠小于b的情況下表現尤為顯著。

對于常規分類任務的蒸餾，研究使用如下“batchified”代價函數：

并初始化如下候選傳輸矩陣：

通過重構和化簡，研究可以使用如下迭代式計算最優傳輸矩陣（具體推導過程參見論文）：

由此，可以算出最優傳輸距離：

SinKD的變體

拓展到回歸任務：對于回歸任務，模型不會為每個選項生成概率，而是僅生成一個標量（d=1）。對于一個包含b個樣本的批次，教師模型和學生模型的輸出分別表示為?? ∈ ?^bx1和?? ∈ ?^bx1。

為了計算教師和學生之間的批量化Sinkhorn距離，成本矩陣的元素由“批量化”回歸輸出之間的絕對差值確定：

拓展到獨熱標簽微調：SinKD方法也適用于僅有獨熱（one-hot）標簽且無法獲取教師模型logits的模型微調。

在這種情況下，可以將單熱標簽視為“假想”的單熱教師模型的logits。由于單熱logits中以零為主，傳統的散度度量（例如KL散度）在處理這種極端情況下的散度量化時顯得無能為力。

實驗與分析

（1）數值結果。與基線和SOTA方法對比，論文方法在大部分任務上均取得了更好的性能。

（2）消融實驗。得出的結論如下：

• Sinkhorn損失在所有損失中對學生模型的收益最大
• 批量化的SinKD優于逐樣本的SinKD
• SinKD超越了基于f-散度變體的蒸餾方法

（3）生成式大語言模型實驗。SinKD可以推廣到生成式大語言模型，并在基于類GPT架構的模型的蒸餾上取得不俗的成績表現。

但同時研究也觀察到，蒸餾效果的影響會隨著PROMPT模板的變化而改變。

這意味著，同樣的任務設置下，更加合理的PROMPT設計能夠更充分地利用教師模型的固有知識。

（4）可視化結果如下。

為了增強內在評估，研究還進行了以下附加分析：

• 隱藏狀態的表示
• 注意力機制的模式
• 層級性能分析

（5）拓展到獨熱標簽微調。與現有的散度度量方法（例如KL散度）不同，SinKD方法還可以擴展用于使用獨熱標簽 (one-hot label) 微調語言模型。

（6）拓展到計算機視覺領域深度網絡。SinKD在所有測試的配置中均穩定地超越了所有基線方法。

總結

研究引入了SinKD以解決現有蒸餾方法的局限性。此外，作者們提出了基于批次的重構方法，以捕捉高維空間中樣本分布的幾何復雜性。最后，研究在各類任務、數據集和模型架構上進一步驗證SinKD的有效性。

更多細節歡迎查閱原論文。

COLING 2024會議論文：https://arxiv.org/abs/2402.17110
IEEE TNNLS期刊論文：https://hal.science/hal-04803835