問題來源

最近讀到一篇模型蒸餾的文章 [1]，其中在設(shè)計軟標(biāo)簽的損失函數(shù)時使用了一種特殊的 softmax： ?

文章中只是簡單的提了一下，其中 T 是 softmax 函數(shù)的溫度超參數(shù)，而沒有做過多解釋。這說明這種用法并非其首創(chuàng)，應(yīng)該是流傳已久。經(jīng)過一番調(diào)研和學(xué)習(xí)，發(fā)現(xiàn)知乎上最高贊的文章《深度學(xué)習(xí)中的 temperature parameter 是什么》[13] 對超參數(shù) T 的講解具有很強的誤導(dǎo)性，所以在此重新寫一篇文章為其正名。 ? 本文的標(biāo)題有兩個雙關(guān)。一個是知識蒸餾的方法用于深度學(xué)習(xí)，同時也需要深入學(xué)習(xí)；另一個則是本文的核心：蒸餾中如何合理運用溫度，讓隱藏的知識更好地?fù)]發(fā)和凝結(jié)。下面我將詳細(xì)講解以上 softmax 公式中溫度系數(shù)的由來以及它起到的作用。 ? ?

蒸餾模型

模型蒸餾或知識蒸餾，最早在 2006 年由 Buciluǎ 在文章 Model Compression?[14]?中提出（很多博主把人名都寫錯了。其后，Hinton 進(jìn)行了歸納和發(fā)展，并在 2015 年發(fā)表了經(jīng)典之作 Distilling the Knowledge in a Neural Network [15]。正是在這篇文章 [2] 中，Hinton 首次提出了 Softmax with Temperature 的方法。 ? 先簡要概括一下模型蒸餾在做什么。出于計算資源的限制或效率的要求，深度學(xué)習(xí)模型在部署推斷時往往需要進(jìn)行壓縮，模型蒸餾是其中一種常見方法。將原始數(shù)據(jù)集上訓(xùn)練的重量級（cumbersome）模型作為教師，讓一個相對更輕量的模型作為學(xué)生。 ? 對于相同的輸入，讓學(xué)生輸出的概率分布盡可能的逼近教師輸出的分布，則大模型的知識就通過這種監(jiān)督訓(xùn)練的方式「蒸餾」到了小模型里。小模型的準(zhǔn)確率往往下降很小，卻能大幅度減少參數(shù)量，從而降低推斷時對 CPU、內(nèi)存、能耗等資源的需求。 ? 對于傳統(tǒng)的監(jiān)督訓(xùn)練，損失函數(shù)可以寫為 KL-散度 ，表示用分布 擬合分布 帶來的誤差。其中 是輸出的真實分布，我們的數(shù)據(jù)集的標(biāo)簽 就從這個分布中采樣而來，對于一個 分類問題， 常常會表示為 one-hot 向量，包含 1? 個 1 和 個 0。對于模型蒸餾，損失函數(shù)可以表示為 ，表示用學(xué)生模型的輸出 來擬合教師模型的輸出 。 ? 我們知道模型在訓(xùn)練收斂后，往往通過 softmax 的輸出不會是完全符合 one-hot 向量那種極端分布的，而是在各個類別上均有概率，推斷時通過 argmax 取得概率最大的類別。Hinton 的文章就指出，教師模型中在這些負(fù)類別（非正確類別）上輸出的概率分布包含了一定的隱藏信息。比如 MNIST 手寫數(shù)字識別，標(biāo)簽為 7 的樣本在輸出時，類別 7 的概率雖然最大，但和類別 1 的概率更加接近，這就說明 1 和 7 很像，這是模型已經(jīng)學(xué)到的隱藏的知識。 ? 我們在使用 softmax 的時候往往會將一個差別不大的輸出變成很極端的分布，用一個三分類模型的輸出舉例： ?

可以看到原本的分布很接近均勻分布，但經(jīng)過 softmax，不同類別的概率相差很大。這就導(dǎo)致類別間的隱藏的相關(guān)性信息不再那么明顯，有誰知道 0.09 和 0.24 對應(yīng)的類別很像呢？為了解決這個問題，我們就引入了溫度系數(shù)。 ? ?

溫度系數(shù)

我們看看對于隨機(jī)生成的相同的模型輸出，經(jīng)過不同的函數(shù)處理，分布會如何變化： ?

最左邊是我們隨機(jī)生成的分布來模擬模型的輸出：。中間五幅圖是使用 softmax 得到的結(jié)果；其中溫度系數(shù) 時相當(dāng)于原始的 softmax；右側(cè)對比了 argmax 得到的結(jié)果?？梢钥闯觯瑥淖蟮接?，這些輸出結(jié)果逐漸從均勻分布向尖銳分布過渡，其中保留的除正確類別以外的信息越來越少。下圖 [3] 更加直觀地展示了不同的溫度系數(shù) 對輸出分布的影響。 ?

靈感來源：https://www.youtube.com/watch?v=tOItokBZSfU

不同的曲線代表不同類別上的概率輸出，同樣 時代表傳統(tǒng)的 softmax，在 時，分布逐漸極端化，最終等價于 argmax，在 時，分布逐漸趨于均勻分布，10 個類別的概率都趨近于1/10。 ? 這兩幅畫很好的說明了 softmax 的本質(zhì)。相對于 argmax 這種直接取最大的「hardmax」，softmax 采用更溫和的方式，將正確類別的概率一定程度地突顯出來。而引入溫度系數(shù)的本質(zhì)目的，就是讓 softmax 的 soft 程度變成可以調(diào)節(jié)的超參數(shù)。 ?

而至于這個系數(shù)為啥叫 Temperature，其實很有深意。我們知道這個場景最早用于模型蒸餾，一般來說蒸餾需要加熱，而加熱會導(dǎo)致熵增。我們發(fā)現(xiàn)，提高溫度系數(shù)會導(dǎo)致輸出分布的信息熵增大！[4] 而在 Hinton 的這篇論文里，為了充分利用教師模型負(fù)類別的 dark 信息，一般會選用一個較高的溫度系數(shù)，這也是本文標(biāo)題叫做高溫蒸餾的原因。 ? 我們可以輕松地推導(dǎo)出 趨于無窮大時，分布將趨于均勻分布，此時信息熵趨于最大 ?

而當(dāng) 趨于 0 時，正確類別的概率接近 1，softmax 的效果逼近 argmax ?

?

反對意見

在最高贊的那篇文章中提到：

如果我們在訓(xùn)練時將t設(shè)置比較大，那么預(yù)測的概率分布會比較平滑，那么loss會很大

首先，如果原文考慮的問題中數(shù)據(jù)的標(biāo)簽是 one-hot 向量而不是蒸餾這種軟標(biāo)簽， 較大時 loss 確實會較大，因為輸出分布比較均勻，不能很好地凸顯正類別上的概率優(yōu)勢。但在蒸餾時并非如此，Hinton 給出的 Loss 函數(shù)如下圖 [5] 所示，分為兩項： ?

圖源：https://nni.readthedocs.io/en/stable/sharings/kd_example.html

第一項? 是教師模型與學(xué)生模型的輸出之間的交叉熵，第二項 是學(xué)生模型與真實標(biāo)簽之間的交叉熵。傳統(tǒng)訓(xùn)練模型時只有 項，所以 可以看做是引入的正則項。文中指出這個正則項使得學(xué)生模型能夠?qū)W到教師模型中的高度泛化的知識，從而需要更少的真實訓(xùn)練樣本。 ? 文中的實驗只用了 3% 的訓(xùn)練樣本，就達(dá)到了近似教師模型的準(zhǔn)確率。我們可以看到這里的 項中，兩個模型都使用了同樣的、較大的溫度系數(shù) ，對輸出的作用是相同的，未必會使 loss 變大。 ?

……那么loss會很大，這樣可以避免我們陷入局部最優(yōu)解。

為什么 loss 大就可以避免陷入局部最優(yōu)呢？我猜作者想表達(dá)的是 loss 很大，從而隨機(jī)梯度下降的時候梯度很大，步長就會很大，從而更容易跳出局部最優(yōu)。該文章的評論區(qū)也有同樣的聲音，但可惜這并不正確。我們還以硬標(biāo)簽 監(jiān)督訓(xùn)練為例，使用交叉熵?fù)p失函數(shù)，設(shè) softmax 的輸出為 ，我們可以推導(dǎo) loss 對于模型輸出 的梯度： ?

交叉熵的梯度 ?

softmax 的梯度 ?

當(dāng) 時 ?

當(dāng) 時 ?

代入鏈?zhǔn)椒▌t，最終的梯度為（推導(dǎo)參考了 [6][7]） ?

顯然標(biāo)簽 與 softmax 的輸出 之差不總能增長 倍，大家可以自己舉一些反例，會發(fā)現(xiàn)大多數(shù)情況下，梯度都不是增大的。那么對于 Hinton 這篇論文，由于 loss 的數(shù)量級沒有變化，所以梯度實際是減小的，所以文章中特意強調(diào)了要將系數(shù) 設(shè)置大一些來補償，比如設(shè)置為 ，在這里給出的 Pytorch 實現(xiàn) [16] 中也是這么做的。 ? 文章中給出了一個高溫情況下的等價，在 時，利用等價無窮小或者是泰勒展開得到： ?

可以清晰的看出這里是 的關(guān)系。 ?

隨著訓(xùn)練的進(jìn)行，我們將 t 變小，也可以稱作降溫，類似于模擬退火算法，這也是為什么要把 t 稱作溫度參數(shù)的原因。變小模型才能收斂。

我不知道將這里的溫度系數(shù)類比模擬退火算法的溫度系數(shù)有什么依據(jù)（Quora 上有個類似的 [8]），但它們真的是不怎么像。同樣也未必是溫度系數(shù)變小模型才能收斂，需要分情況：如果是模型蒸餾， 項始終都使用較大的溫度；如果是使用真實標(biāo)簽訓(xùn)練，確實選取較小的溫度系數(shù)，更利于模型收斂。 ? 可以這樣理解，溫度系數(shù)較大時，模型需要訓(xùn)練得到一個很陡峭的輸出，經(jīng)過 softmax 之后才能獲得一個相對陡峭的結(jié)果；溫度系數(shù)較小時，模型輸出稍微有點起伏，softmax 就很敏感地把分布變得尖銳，認(rèn)為模型學(xué)到了知識。 ? 所以，使用一個固定的小于 1 的溫度系數(shù)是合理的，這也是那篇文章里提到的推薦系統(tǒng)所做的，它沒有降溫過程，直接設(shè)置了?T=0.05?。如果大家在哪篇文章中看到了降溫過程，還請在評論區(qū)指正。 ? ?

其他場景

? 這里我們天馬行空地設(shè)想一個場景：在一些序列生成任務(wù)中，比如 seq2seq 的機(jī)器翻譯模型，或者是驗證碼識別的 CTC 算法 [9] 中，輸出的每一個時間步都會有一個分布。最終的序列會使用 BeamSearch [10] 或者 Viterbi [11] 等算法搜索 Top-K 概率的序列。 ? 這類方法介于逐時間步 argmax 的完全貪心策略和全局動態(tài)規(guī)劃的優(yōu)化策略之間。雖然 BeamSearch 中我們不需要提前 softmax，但假如我們做了帶溫度系數(shù)的 softmax，就可以控制輸出分布的尖銳程度。對于這類逐步計算累積概率的算法，在每個時間步的概率分布較為均勻時就容易輸出不同的結(jié)果。所以在這類問題下，高溫可能導(dǎo)致輸出序列的多樣性。 ? 對于這類場景，我沒有進(jìn)行嚴(yán)格證明也沒有很深的經(jīng)驗，只是一個猜想。這里有類似的說法 [12]，但都不能作為參考依據(jù)。大家感興趣的話可以將 softmax with temperature 引入 BeamSearch 看看會不會對輸出的豐富性造成影響。假如算法只依賴每個時間步的概率大小關(guān)系，那輸出就是確定的，說明我們猜想失敗?；蛘哂邢嚓P(guān)經(jīng)驗的同學(xué)也可以在評論區(qū)給出參考文獻(xiàn)。 ? ?

后話

寫完這篇文章才發(fā)現(xiàn)，潘小小【經(jīng)典簡讀】知識蒸餾（Knowledge Distillation）經(jīng)典之作 [17] 一文中已有類似的探討。盡管如此，我相信這篇文章還是可以起到一定的科普作用，讓那些和我一樣對知識蒸餾不太了解的同學(xué)，從溫度系數(shù)這個關(guān)鍵詞入手，能夠快速得到想要的答案。 ? 讀完 Hinton 的文章，有兩個強烈的感受：一是感覺他太牛了，3 句話讓我讀了? 18 遍，全文很少用公式，基本沒有配圖，但把算法講得清清楚楚；二就是，他的寫作中長從句實在太多了，一句話 60 個單詞，讀起來很不友好。如果對這篇文章感興趣，也可以看上面潘小小的那篇解讀。文章最后講到了一種和 MOE 很像的分布式集成學(xué)習(xí)方法，在潘的文章中沒有介紹，由于這不是今天的主題，所以我也沒用筆墨，大家如果對這部分感興趣也可以來找我討論。 ? 說出來很難相信，我其實不是做 AI 方向的，我是做系統(tǒng)的，所以歡迎大家懟我（°ー°〃）。 ? ? ?

參考文獻(xiàn)

[1] Group knowledge transfer: Federated learning of large cnns at the edgehttps://proceedings.neurips.cc/paper/2020/file/a1d4c20b182ad7137ab3606f0e3fc8a4-Paper.pdf

[2]Distilling the Knowledge in a Neural Network?https://arxiv.org/abs/1503.02531

[3] PR-009: Distilling the Knowledge in a Neural Network (Slide: English, Speaking: Korean)?https://www.youtube.com/watch?v=tOItokBZSfU

[4] What is the role of temperature in Softmax?https://stats.stackexchange.com/questions/527080/what-is-the-role-of-temperature-in-softmax#answer-527082

[5] Knowledge Distillation on NNIhttps://nni.readthedocs.io/en/stable/sharings/kd_example.html

[6] softmax, CrossEntropyLoss 與梯度計算公式https://blog.csdn.net/jiongjiongai/article/details/88324000

[7] 關(guān)于Softmax的數(shù)值穩(wěn)定性和梯度反向傳播https://zhuanlan.zhihu.com/p/92714192

[8] What is the temperature parameter in deep learning?https://www.quora.com/What-is-the-temperature-parameter-in-deep-learning

[9] 詳解CTChttps://zhuanlan.zhihu.com/p/42719047

[10] 文本生成解碼之 Beam Searchhttps://zhuanlan.zhihu.com/p/43703136

[11] 如何通俗地講解 viterbi 算法？https://www.zhihu.com/question/20136144/answer/763021768

[12]What is Temperature in LSTM??https://www.quora.com/What-is-Temperature-in-LSTM

[13] https://zhuanlan.zhihu.com/p/132785733?

[14] https://dl.acm.org/doi/abs/10.1145/1150402.1150464?

[15] https://arxiv.org/abs/1503.02531?

[16]?https://nni.readthedocs.io/en/stable/sharings/kd_example.html?

[17] https://zhuanlan.zhihu.com/p/102038521

編輯：黃飛