追求概率質量的帕累托最優：廣義α-β散度引導知識蒸餾框架ABKD

知識蒸餾（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. 當 β→1 時，退化為 RKLD 的效應。
2. 當 β→0 時，退化為 FKLD 的效應。

較大的 β 值會增強置信集中效應，使匹配性能更加專注于模型最有信心的類別（見圖 1c）。

在 (b) 和 (b1) 中，使用以下方式衡量難度集中效應：

1. 當 α→1 時，退化為 FKLD 的效應。
2. 當 α→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 的有效性。