模型知識(shí)蒸餾新SOTA！告別傳統(tǒng)散度蒸餾｜騰訊優(yōu)圖&中科大出品

作者：量子位 2024-12-04 09:15:00

研究引入了SinKD以解決現(xiàn)有蒸餾方法的局限性。此外，作者們提出了基于批次的重構(gòu)方法，以捕捉高維空間中樣本分布的幾何復(fù)雜性。

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

甚至能在不同類型和架構(gòu)的LLMs（大語(yǔ)言模型）上達(dá)到新SOTA。

這就是來(lái)自中科大、騰訊優(yōu)圖實(shí)驗(yàn)室提出的一種基于Sinkhorn距離的知識(shí)蒸餾方法，能把大的、復(fù)雜的教師模型的知識(shí)“蒸餾”到小的、簡(jiǎn)單的學(xué)生模型中，從而讓小模型也能像大模型一樣工作。

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

當(dāng)兩個(gè)模型的輸出差異較大時(shí)，它們就不太管用了。

KL散度：會(huì)導(dǎo)致學(xué)生模型的輸出變得過(guò)于平滑，失去了區(qū)分性；
RKL散度：會(huì)讓學(xué)生的輸出變得太簡(jiǎn)單，不能很好地模仿教師模型；
JS散度：會(huì)讓學(xué)生模型低估稀有事件的概率；

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

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

最終，通過(guò)在兩個(gè)流行的自然語(yǔ)言處理測(cè)試集（GLUE和SuperGLUE）上測(cè)試，新方法在編碼器、編碼器-解碼器以及解碼器等不同架構(gòu)的所有類型LLMs上均優(yōu)于當(dāng)前的最先進(jìn)方法。

研究背景

知識(shí)蒸餾的提出是為了通過(guò)對(duì)齊教師模型的軟目標(biāo)（例如輸出logits和中間層表示）來(lái)將教師模型內(nèi)在固有的知識(shí)傳遞給學(xué)生模型。

給定訓(xùn)練集中的一個(gè)樣本x_i及其真實(shí)標(biāo)簽??_?? ∈ ?^??，來(lái)自教師模型??_??和學(xué)生模型??_??的輸出logits ??_?? ∈ ?^??和??_?? ∈ ?^??可以由以下式子得到：

其中為softmax函數(shù)， τ是溫度參數(shù), d是輸出logits的維度。基于logit的知識(shí)蒸餾的目標(biāo)是σ_Τ最小化測(cè)量散度J（??_??，??_??）以實(shí)現(xiàn)知識(shí)傳遞。

研究動(dòng)機(jī)

現(xiàn)有研究已經(jīng)嘗試使用Kullback-Leibler（KL）散度、反Kullback-Leibler（RKL）散度和Jensen-Shannon（JS）散度。

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

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

KL蒸餾會(huì)導(dǎo)致模式平均，使學(xué)生學(xué)習(xí)到一個(gè)過(guò)于平滑的分布，涵蓋了教師的整個(gè)支撐集；
RKL會(huì)引起模式塌陷，學(xué)生僅關(guān)注教師分布中高概率的顯著區(qū)域，而忽視了其余部分；
JS蒸餾會(huì)產(chǎn)生模式低估，由于懲罰不足，學(xué)生會(huì)低估稀有事件的概率。

為了解決傳統(tǒng)散度度量的問(wèn)題，研究做出了以下貢獻(xiàn)：

提出了一種知識(shí)蒸餾方法SinKD，采用Sinkhorn距離作為散度度量。它不僅解決了KL、RKL和JS散度在極端場(chǎng)景下的局限性，而且避免了計(jì)算Wasserstein距離的負(fù)擔(dān)。
深入探討了Sinkhorn距離的性質(zhì)，并將SinKD重新reformulated為batch-wise OT，擴(kuò)展了它在NLP任務(wù)中的適用性。
通過(guò)大量的可比性、有效性和泛化性實(shí)驗(yàn)證明了SinKD相較于目前最先進(jìn)的方法的優(yōu)越性。并為實(shí)際應(yīng)用提供了使用SinKD進(jìn)行蒸餾的實(shí)用指導(dǎo)方針。

傳統(tǒng)散度度量的缺陷

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

其次，由于使用KL損失進(jìn)行優(yōu)化，學(xué)生模型試圖對(duì)教師模型的多模態(tài)分布進(jìn)行平均化，從而導(dǎo)致對(duì)這些模式的擬合不足。這被稱為“模式平均問(wèn)題”（mode-averaging problem）。

因此，學(xué)生模型無(wú)法捕獲數(shù)據(jù)中的所有關(guān)鍵模式，最終影響模型性能。

第三，KL散度對(duì)應(yīng)的是一個(gè)非平滑函數(shù)，這為優(yōu)化過(guò)程帶來(lái)了挑戰(zhàn)。

與KL散度一樣，具有內(nèi)在的不對(duì)稱性，從而導(dǎo)致在捕捉分布差異時(shí)出現(xiàn)不一致性。

此外，優(yōu)化的學(xué)生模型傾向于僅關(guān)注教師分布中概率較高的事件，這被稱為“模式崩塌問(wèn)題”（mode-collapsing）。

如果教師對(duì)某個(gè)事件賦予零概率，學(xué)生模型也被迫做出相同的預(yù)測(cè)。

其中m_?? = 1/2（??_??+??_??）受制于非平滑性，JS損失在優(yōu)化過(guò)程中面臨挑戰(zhàn)。

另外，由于JS損失在低概率區(qū)域的匹配上懲罰不足，學(xué)生模型可能會(huì)過(guò)度低估稀有事件的概率。

對(duì)于分布之間重疊較少甚至完全不重疊的情況退化為常數(shù)時(shí)，還存在梯度消失的風(fēng)險(xiǎn)。

最優(yōu)傳輸距離的優(yōu)勢(shì)

Wasserstein距離通過(guò)求解兩個(gè)分布之間的最優(yōu)傳輸計(jì)劃來(lái)量化它們的差異。

直觀地看，它可以被認(rèn)為是將一個(gè)分布（即學(xué)生的logits分布）轉(zhuǎn)換為另一個(gè)分布（即教師的logits分布）所需的最小“代價(jià)”，其中“代價(jià)”可以定義為被移動(dòng)的質(zhì)量與移動(dòng)距離的乘積。

與傳統(tǒng)的散度度量相比，Wasserstein距離作為蒸餾的成本函數(shù)更為合理，因?yàn)樗灰蕾囉趯?duì)被測(cè)量分布的隱式假設(shè)。此外，它幾乎處處可微，從而便于優(yōu)化。

另外，現(xiàn)有的散度度量只能獨(dú)立處理每個(gè)樣本對(duì)，進(jìn)行逐一logit的匹配，對(duì)于一批樣本，這些方法無(wú)法定位來(lái)自同一樣本的教師和學(xué)生的logits對(duì)，從而無(wú)法實(shí)現(xiàn)整體距離的最小化。

由于計(jì)算Sinkhorn距離的過(guò)程可以實(shí)現(xiàn)來(lái)自同一樣本的兩個(gè)輸出之間的精確逐元素匹配，研究提出了“批量化”的SinKD方法（batchified SinKD）。

通過(guò)這種方式，即使通過(guò)低維觀測(cè)，也能夠捕捉復(fù)雜且隱式分布的幾何結(jié)構(gòu)。

方法介紹

這里簡(jiǎn)要介紹SinKD的核心方法，詳細(xì)推導(dǎo)過(guò)程可以參閱原論文。

批量重構(gòu)的Sinkhorn距離

對(duì)于本問(wèn)題，Wasserstein距離的定義如下：

其中，

Wasserstein距離本身在解析計(jì)算上存在困難，其計(jì)算成本對(duì)于蒸餾大型語(yǔ)言模型來(lái)說(shuō)高得難以承受。

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

Sinkhorn距離的定義如下：

逐樣本蒸餾將每個(gè)實(shí)例獨(dú)立處理，但忽略了一個(gè)批次樣本中的整體趨勢(shì)。

研究摒棄了僅在每對(duì)教師-學(xué)生樣本對(duì)上工作的逐樣本知識(shí)蒸餾方法，轉(zhuǎn)而在教師和學(xué)生樣本組上執(zhí)行知識(shí)蒸餾。

一個(gè)包含b個(gè)樣本的批次會(huì)整體參與散度度量。通過(guò)批量重構(gòu)，這種方法有效地增加了“觀測(cè)”空間的維度，特別是在d遠(yuǎn)小于b的情況下表現(xiàn)尤為顯著。

對(duì)于常規(guī)分類任務(wù)的蒸餾，研究使用如下“batchified”代價(jià)函數(shù)：

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

通過(guò)重構(gòu)和化簡(jiǎn)，研究可以使用如下迭代式計(jì)算最優(yōu)傳輸矩陣（具體推導(dǎo)過(guò)程參見(jiàn)論文）：

由此，可以算出最優(yōu)傳輸距離：

SinKD的變體

拓展到回歸任務(wù)：對(duì)于回歸任務(wù)，模型不會(huì)為每個(gè)選項(xiàng)生成概率，而是僅生成一個(gè)標(biāo)量（d=1）。對(duì)于一個(gè)包含b個(gè)樣本的批次，教師模型和學(xué)生模型的輸出分別表示為?? ∈ ?^bx1和?? ∈ ?^bx1。

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