知識蒸餾(KD)是一種將大模型(教師)的知識遷移到小模型(學生)的技術,學生通過模仿教師預測分布,充分利用軟標簽信息,通常優于傳統監督微調,已在圖像分類、文本生成等任務及最新工作(如 DeepSeek-R1、Qwen-3)中得到驗證。其核心在于分布匹配方式的選擇,主流方法多用前向 KL 散度(FKLD)或反向 KL 散度(RKLD),但 FKLD 易導致輸出過度平滑,難以聚焦關鍵類別,RKLD 則使學生過度自信、降低多樣性。實驗證明,兩者在多任務中表現有限,且目前缺乏系統方法揭示其深層問題,阻礙了更通用高效 KD 框架的發展。因此,一個自然的問題產生了:
究竟是什么潛在因素導致了 FKLD 和 RKLD 的次優表現?
為了解答這個問題,我們通過追蹤對數質量比(LogR),分析不同散度在訓練中如何影響學生分布的概率分配。進一步分析表明在溫和假設下,LogR 與損失函數對 logits 的梯度成正比,這使我們將問題轉化為分析:不同散度算法如何影響 ∣LogR∣ 下降。
- 標題:ABKD: Pursuing a Proper Allocation of the Probability Mass in Knowledge Distillation via α-β-Divergence
- 論文:https://arxiv.org/abs/2505.04560
- 代碼:https://github.com/ghwang-s/abkd
在此框架下,我們識別出兩種關鍵的「模式集中效應」:難度集中與置信集中。
- 「難度集中」指的是更關注于在教師分布 p 與學生分布 q 之間誤差較大的模式。
- 「置信集中」指的是更關注于學生分布 q 本身高度自信的模式。
進一步研究發現,FKLD 和 RKLD 的局限性正源于對這兩種集中效應的極端利用。
- FKLD 集中效應較弱,對所有類別誤差一視同仁,導致學生難以聚焦目標類別,易出現錯誤預測(見圖 1d)。
- RKLD 集中效應過強,兼顧高誤差和高置信度類別,易陷入「平凡解」,即學生只關注目標類別,忽略教師分布的其他知識(見圖 1e)。
揭示這些局限性后,我們提出一個關鍵問題:我們能否找到一種通用且理論上有依據的方法,來平衡「難度集中」與「置信集中」效應?
方法
為此,我們引入了 α-β 散度。
其廣義統一了 FKLD、RKLD 及 Hellinger 距離等多種散度。
正如下一節理論表明,α-β 散度可通過調節 α 和 β 在 FKLD 與 RKLD 間平滑插值,實現對難度集中和置信集中效應的靈活權衡。這一機制帶來更合理的概率分配,據此我們提出了通用蒸餾框架 ABKD,形式為最小化:
基于 α-β 散度平衡難度集中和置信集中
ABKD 提供了一個統一空間權衡難度集中與置信集中。為解釋這一點,回到對數質量比(LogR)。下列命題解釋了超參數 α 和 β 如何影響 ∣LogR∣ 的減小。
- 當 β→1 時,退化為 RKLD 的效應。
- 當 β→0 時,退化為 FKLD 的效應。
較大的 β 值會增強置信集中效應,使匹配性能更加專注于模型最有信心的類別(見圖 1c)。
在 (b) 和 (b1) 中,使用以下方式衡量難度集中效應:
- 當 α→1 時,退化為 FKLD 的效應。
- 當 α→0 時,退化為 RKLD 的效應。
較小的 α 會放大難度集中效應,在困難類別上進行更強的誤差懲罰,從而實現更激進的匹配(見圖 1b)。
通過調節 α 和 β,ABKD 實現對兩種效應的精細平衡,避免了傳統方法(FKLD/RKLD)的極端情況(見圖 1g)。
與 α-散度的比較
α-散度是 ABKD 框架的一個特例(當 β=1?α 時成立)。
雖然它也能實現 FKLD(α→1)和 RKLD(α→0)之間的轉換,但根據命題 4.2,為了減小 α,必須相應地增大 β,以滿足 α+β=1 的限制條件。這種不必要的限制阻礙了其性能的進一步提升,具體如圖 1(a) 和圖 1(f) 所示。
與 WSD 的比較
另一種方案是對 FKLD 和 RKLD 取加權和(WSD):
不幸的是,這種組合度量會過度強調 p 和 q 中概率較小的模式。當 q(k)≈0,p(k)>0 或 p(k)≈0,q(k)>0 時,有:
因此,算法必須關注極端情況以最小化目標函數,導致概率分配不合理。此外這種情況下梯度范數也會過度增長,導致參數更新中可能含有噪聲,破壞優化穩定性。
最后一種方法是使用 Jensen-Shannon 散度:
其中 m=(p+q)/2。它的缺點是當分布 p 和 q 相距較遠時(這是訓練初期的常見情況),會出現梯度消失,阻礙模型收斂。
實驗
自然語言任務
我們在五個指令跟隨基準上評估方法,使用 ROUGE-L 指標,對比了 SFT、FKLD、GKD 和 DISTILLM 等方法。
實驗結果表明僅修改蒸餾目標,ABKD 在不同數據集上均優于或匹配 FKLD、SFT。對比基于 SGO 的方法(如 GKD、DISTILLM)效果提升明顯,ABKD 在不公平對比下依然表現優越。
- 損失函數影響
α-β 散度優于 α 或 β 散度及 WSD。在五個數據集上相較基線有穩定提升。訓練過程中的優勢也體現在圖 2。
視覺任務
在 12 個常用圖像識別數據集上進行實驗,對比方法比較了多種先進蒸餾方法:KD、DKD、LSD 和 TTM 等。
結果表明 ABKD 能根據不同教師-學生組合自動選擇合適的蒸餾目標。此外 ABKD 可作為簡單的插件工具,優化現有方法的損失函數,帶來額外性能提升。
敏感性分析
α 對難度集中的影響:圖 6(a)(b) 展示了不同 α 下的訓練表現。在 CIFAR-100 中,較小 α 提升有限;在 Dolly 等大輸出空間任務中,適當減小 α 有助于避免陷入局部最優。
β 對置信集中的影響:圖 6(c)(d) 顯示了 β 對輸出分布的 Shannon 熵和 Self-BLEU 的影響。較小 β 使模型更關注低置信度類別,提升輸出分布的平滑性和多樣性(熵更高,Self-BLEU 更低)。
結語
本文指出,知識蒸餾的核心在于平衡「難度集中」和「置信集中」兩種效應,而傳統 FKLD 和 RKLD 只覆蓋兩個極端,導致概率分配不合理。為此,我們提出基于 α-β 散度的 ABKD 框架,統一并推廣了現有方法,實現兩種效應的靈活權衡。理論與大量實驗均驗證了 ABKD 的有效性。
特別聲明:以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布,本平臺僅提供信息存儲服務。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.