知識(shí)蒸餾（KD）是一種將大模型（教師）的知識(shí)遷移到小模型（學(xué)生）的技術(shù)，學(xué)生通過(guò)模仿教師預(yù)測(cè)分布，充分利用軟標(biāo)簽信息，通常優(yōu)于傳統(tǒng)監(jiān)督微調(diào)，已在圖像分類、文本生成等任務(wù)及最新工作（如 DeepSeek-R1、Qwen-3）中得到驗(yàn)證。其核心在于分布匹配方式的選擇，主流方法多用前向 KL 散度（FKLD）或反向 KL 散度（RKLD），但 FKLD 易導(dǎo)致輸出過(guò)度平滑，難以聚焦關(guān)鍵類別，RKLD 則使學(xué)生過(guò)度自信、降低多樣性。實(shí)驗(yàn)證明，兩者在多任務(wù)中表現(xiàn)有限，且目前缺乏系統(tǒng)方法揭示其深層問(wèn)題，阻礙了更通用高效 KD 框架的發(fā)展。因此，一個(gè)自然的問(wèn)題產(chǎn)生了：

究竟是什么潛在因素導(dǎo)致了 FKLD 和 RKLD 的次優(yōu)表現(xiàn)？

為了解答這個(gè)問(wèn)題，我們通過(guò)追蹤對(duì)數(shù)質(zhì)量比（LogR），分析不同散度在訓(xùn)練中如何影響學(xué)生分布的概率分配。進(jìn)一步分析表明在溫和假設(shè)下，LogR 與損失函數(shù)對(duì) logits 的梯度成正比，這使我們將問(wèn)題轉(zhuǎn)化為分析：不同散度算法如何影響 ∣LogR∣ 下降。

• 標(biāo)題：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

在此框架下，我們識(shí)別出兩種關(guān)鍵的「模式集中效應(yīng)」：難度集中與置信集中。

• 「難度集中」指的是更關(guān)注于在教師分布 p 與學(xué)生分布 q 之間誤差較大的模式。
• 「置信集中」指的是更關(guān)注于學(xué)生分布 q 本身高度自信的模式。

進(jìn)一步研究發(fā)現(xiàn)，F(xiàn)KLD 和 RKLD 的局限性正源于對(duì)這兩種集中效應(yīng)的極端利用。

• FKLD 集中效應(yīng)較弱，對(duì)所有類別誤差一視同仁，導(dǎo)致學(xué)生難以聚焦目標(biāo)類別，易出現(xiàn)錯(cuò)誤預(yù)測(cè)（見(jiàn)圖 1d）。
• RKLD 集中效應(yīng)過(guò)強(qiáng)，兼顧高誤差和高置信度類別，易陷入「平凡解」，即學(xué)生只關(guān)注目標(biāo)類別，忽略教師分布的其他知識(shí)（見(jiàn)圖 1e）。

揭示這些局限性后，我們提出一個(gè)關(guān)鍵問(wèn)題：我們能否找到一種通用且理論上有依據(jù)的方法，來(lái)平衡「難度集中」與「置信集中」效應(yīng)？

方法

為此，我們引入了 α-β 散度。

其廣義統(tǒng)一了 FKLD、RKLD 及 Hellinger 距離等多種散度。

正如下一節(jié)理論表明，α-β 散度可通過(guò)調(diào)節(jié) α 和 β 在 FKLD 與 RKLD 間平滑插值，實(shí)現(xiàn)對(duì)難度集中和置信集中效應(yīng)的靈活權(quán)衡。這一機(jī)制帶來(lái)更合理的概率分配，據(jù)此我們提出了通用蒸餾框架 ABKD，形式為最小化：

基于 α-β 散度平衡難度集中和置信集中

ABKD 提供了一個(gè)統(tǒng)一空間權(quán)衡難度集中與置信集中。為解釋這一點(diǎn)，回到對(duì)數(shù)質(zhì)量比（LogR）。下列命題解釋了超參數(shù) α 和 β 如何影響 ∣LogR∣ 的減小。

1. 當(dāng) β→1 時(shí)，退化為 RKLD 的效應(yīng)。
2. 當(dāng) β→0 時(shí)，退化為 FKLD 的效應(yīng)。

較大的 β 值會(huì)增強(qiáng)置信集中效應(yīng)，使匹配性能更加專注于模型最有信心的類別（見(jiàn)圖 1c）。

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

1. 當(dāng) α→1 時(shí)，退化為 FKLD 的效應(yīng)。
2. 當(dāng) α→0 時(shí)，退化為 RKLD 的效應(yīng)。

較小的 α 會(huì)放大難度集中效應(yīng)，在困難類別上進(jìn)行更強(qiáng)的誤差懲罰，從而實(shí)現(xiàn)更激進(jìn)的匹配（見(jiàn)圖 1b）。

通過(guò)調(diào)節(jié) α 和 β，ABKD 實(shí)現(xiàn)對(duì)兩種效應(yīng)的精細(xì)平衡，避免了傳統(tǒng)方法（FKLD/RKLD）的極端情況（見(jiàn)圖 1g）。

與 α-散度的比較

α-散度是 ABKD 框架的一個(gè)特例（當(dāng) β=1?α 時(shí)成立）。

雖然它也能實(shí)現(xiàn) FKLD（α→1）和 RKLD（α→0）之間的轉(zhuǎn)換，但根據(jù)命題 4.2，為了減小 α，必須相應(yīng)地增大 β，以滿足 α+β=1 的限制條件。這種不必要的限制阻礙了其性能的進(jìn)一步提升，具體如圖 1(a) 和圖 1(f) 所示。

與 WSD 的比較

另一種方案是對(duì) FKLD 和 RKLD 取加權(quán)和（WSD）：

不幸的是，這種組合度量會(huì)過(guò)度強(qiáng)調(diào) p 和 q 中概率較小的模式。當(dāng) q(k)≈0,p(k)>0 或 p(k)≈0,q(k)>0 時(shí)，有：

因此，算法必須關(guān)注極端情況以最小化目標(biāo)函數(shù)，導(dǎo)致概率分配不合理。此外這種情況下梯度范數(shù)也會(huì)過(guò)度增長(zhǎng)，導(dǎo)致參數(shù)更新中可能含有噪聲，破壞優(yōu)化穩(wěn)定性。

最后一種方法是使用 Jensen-Shannon 散度：

其中 m=(p+q)/2。它的缺點(diǎn)是當(dāng)分布 p 和 q 相距較遠(yuǎn)時(shí)（這是訓(xùn)練初期的常見(jiàn)情況），會(huì)出現(xiàn)梯度消失，阻礙模型收斂。

實(shí)驗(yàn)

自然語(yǔ)言任務(wù)

我們?cè)谖鍌€(gè)指令跟隨基準(zhǔn)上評(píng)估方法，使用 ROUGE-L 指標(biāo)，對(duì)比了 SFT、FKLD、GKD 和 DISTILLM 等方法。

實(shí)驗(yàn)結(jié)果表明僅修改蒸餾目標(biāo)，ABKD 在不同數(shù)據(jù)集上均優(yōu)于或匹配 FKLD、SFT。對(duì)比基于 SGO 的方法（如 GKD、DISTILLM）效果提升明顯，ABKD 在不公平對(duì)比下依然表現(xiàn)優(yōu)越。

• 損失函數(shù)影響

α-β 散度優(yōu)于 α 或 β 散度及 WSD。在五個(gè)數(shù)據(jù)集上相較基線有穩(wěn)定提升。訓(xùn)練過(guò)程中的優(yōu)勢(shì)也體現(xiàn)在圖 2。

視覺(jué)任務(wù)

在 12 個(gè)常用圖像識(shí)別數(shù)據(jù)集上進(jìn)行實(shí)驗(yàn)，對(duì)比方法比較了多種先進(jìn)蒸餾方法：KD、DKD、LSD 和 TTM 等。

結(jié)果表明 ABKD 能根據(jù)不同教師-學(xué)生組合自動(dòng)選擇合適的蒸餾目標(biāo)。此外 ABKD 可作為簡(jiǎn)單的插件工具，優(yōu)化現(xiàn)有方法的損失函數(shù)，帶來(lái)額外性能提升。

敏感性分析

α 對(duì)難度集中的影響：圖 6(a)(b) 展示了不同 α 下的訓(xùn)練表現(xiàn)。在 CIFAR-100 中，較小 α 提升有限；在 Dolly 等大輸出空間任務(wù)中，適當(dāng)減小 α 有助于避免陷入局部最優(yōu)。

β 對(duì)置信集中的影響：圖 6(c)(d) 顯示了 β 對(duì)輸出分布的 Shannon 熵和 Self-BLEU 的影響。較小 β 使模型更關(guān)注低置信度類別，提升輸出分布的平滑性和多樣性（熵更高，Self-BLEU 更低）。

結(jié)語(yǔ)

本文指出，知識(shí)蒸餾的核心在于平衡「難度集中」和「置信集中」兩種效應(yīng)，而傳統(tǒng) FKLD 和 RKLD 只覆蓋兩個(gè)極端，導(dǎo)致概率分配不合理。為此，我們提出基于 α-β 散度的 ABKD 框架，統(tǒng)一并推廣了現(xiàn)有方法，實(shí)現(xiàn)兩種效應(yīng)的靈活權(quán)衡。理論與大量實(shí)驗(yàn)均驗(yàn)證了 ABKD 的有效性。