今天小編分享的科學經驗：模型知識蒸餾新SOTA！告别傳統散度蒸餾，歡迎閱讀。
用大模型 " 蒸餾 " 小模型，有新招了！
甚至能在不同類型和架構的 LLMs（大語言模型）上達到新 SOTA。
這就是來自中科大、騰訊優圖實驗室提出的一種基于 Sinkhorn 距離的知識蒸餾方法，能把大的、復雜的教師模型的知識 " 蒸餾 " 到小的、簡單的學生模型中，從而讓小模型也能像大模型一樣工作。
之所以提出新方法，主要是現有的知識蒸餾（KD）方法都有各自的局限性：
當兩個模型的輸出差異較大時，它們就不太管用了。
KL 散度：會導致學生模型的輸出變得過于平滑，失去了區分性；
RKL 散度：會讓學生的輸出變得太簡單，不能很好地模仿教師模型；
JS 散度：會讓學生模型低估稀有事件的概率；
而基于 Sinkhorn 距離的新方法能更準确地衡量和縮小教師模型和學生模型之間的差異，從而提高了學生模型的性能。
此外，研究還提出了一種基于批量的重構方法，從而在高維空間中捕捉跨樣本分布的幾何復雜性。
最終，通過在兩個流行的自然語言處理測試集（GLUE 和 SuperGLUE）上測試，新方法在編碼器、編碼器 - 解碼器以及解碼器等不同架構的所有類型 LLMs 上均優于當前的最先進方法。
研究背景
知識蒸餾的提出是為了通過對齊教師模型的軟目标（例如輸出 logits 和中間層表示）來将教師模型内在固有的知識傳遞給學生模型。
給定訓練集中的一個樣本 x_i 及其真實标籤 ∈ ℝ，來自教師模型和學生模型的輸出 logits ∈ ℝ和 ∈ ℝ可以由以下式子得到：
其中為 softmax 函數， τ 是溫度參數 , d 是輸出 logits 的維度。基于 logit 的知識蒸餾的目标是 σΤ 最小化測量散度 J（，）以實現知識傳遞。
研究動機
現有研究已經嘗試使用 Kullback-Leibler（KL）散度、反 Kullback-Leibler（RKL）散度和 Jensen-Shannon（JS）散度。
所有這些度量都可以被視為f- 散度度量的變體，而 f- 散度度量在量化缺乏實質性交集的任何兩個分布時都存在明顯局限性。
此外，每種度量都有其自身的缺陷：
KL 蒸餾會導致模式平均，使學生學習到一個過于平滑的分布，涵蓋了教師的整個支撐集；
RKL 會引起模式塌陷，學生僅關注教師分布中高概率的顯著區網域，而忽視了其餘部分；
JS 蒸餾會產生模式低估，由于懲罰不足，學生會低估稀有事件的概率。
為了解決傳統散度度量的問題，研究做出了以下貢獻：
提出了一種知識蒸餾方法 SinKD，采用 Sinkhorn 距離作為散度度量。它不僅解決了 KL、RKL 和 JS 散度在極端場景下的局限性，而且避免了計算 Wasserstein 距離的負擔。
深入探讨了 Sinkhorn 距離的性質，并将 SinKD 重新 reformulated 為 batch-wise OT，擴展了它在 NLP 任務中的适用性。
通過大量的可比性、有效性和泛化性實驗證明了 SinKD 相較于目前最先進的方法的優越性。并為實際應用提供了使用 SinKD 進行蒸餾的實用指導方針。
傳統散度度量的缺陷
首先，KL 散度是不對稱的，表現為 JKL（，）≠ JKL（，），這一性質違反了距離度量的對稱性特性，從而引入了一些不一致性。
其次，由于使用 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
— 完 —
投稿請發郵件到：
标題注明【投稿】，告訴我們：
你是誰，從哪來，投稿内容
附上論文 / 項目主頁鏈接，以及聯系方式哦
我們會（盡量）及時回復你
點這裡關注我，記得标星哦～
一鍵三連「分享」、「點贊」和「在看」
科技前沿進展日日相見 ~
>
