之前我們討論了『模型壓縮與蒸餾！BERT的忒修斯船』，算是一個(gè)開(kāi)篇。本文繼續(xù)討論關(guān)于模型蒸餾（Distilling Knowledge）及關(guān)于BERT模型的知識(shí)蒸餾，分享針對(duì)具體任務(wù)時(shí)可行的簡(jiǎn)潔方案，同時(shí)在新的視角下探討了知識(shí)蒸餾有效的一些原因，并通過(guò)實(shí)驗(yàn)進(jìn)行驗(yàn)證。

模型蒸餾的最重要的一個(gè)特點(diǎn)就是降低資源使用以及加速模型推理速度，而小模型往往性能較低，本文總結(jié)如何通過(guò)蒸餾來(lái)使小模型具有更好的性能。

Distilling the Knowledge in a Neural Network

這篇是2015年Hinton發(fā)表的,也是我看到的最早提出Knowledge Distillation的論文[1]。

在這篇論文中，Hinton指出one-hot 的label只指示了true label 的信息，但是沒(méi)有給出negative label 之間、negative 與 true label之間的相對(duì)關(guān)系，比如：

現(xiàn)在的任務(wù)是給定一個(gè)詞（比如：蘋(píng)果），然后判斷詞對(duì)應(yīng)的類(lèi)別（電視/手機(jī)/水果/汽車(chē)），假如現(xiàn)在我們有兩個(gè)樣本：（蘋(píng)果，[0,0,1,0]）和（小米，[0,1,0,0]）而one-hot 形式的label并不能告訴我們，蘋(píng)果中 label是水果的概率高出label是拖拉機(jī)的概率，稍低于是手機(jī)的概率，而小米中l(wèi)abel是電視的概率稍低于是手機(jī)的概率，但是同時(shí)要高于是汽車(chē)和水果的概率，這些相對(duì)關(guān)系在one-hot 形式的label中是無(wú)法得到的。

而這些信息非常重要，有了這些信息，我們可以更容易的學(xué)習(xí)任務(wù)。于是提出了Teacher-Student模式，即用一個(gè)大的復(fù)雜的模型（也可以是ensemble后的）來(lái)先學(xué)習(xí)，然后得到label的相對(duì)關(guān)系（logits），然后將學(xué)習(xí)到的知識(shí)遷移到一個(gè)小模型（Student）。

Distilling

具體遷移過(guò)程是Student 在進(jìn)行training 時(shí)，除了學(xué)習(xí)ground truth 外，還需要學(xué)習(xí)label 的probability（softmax output），但是不是直接學(xué)習(xí)softmax output，而是學(xué)習(xí)soften labels，所謂soften labels 即經(jīng)過(guò)Temperature 平滑后的 probability，具體形式：

其中T 越大，對(duì)應(yīng)的probability 越平滑，如下圖所示。而平滑probability 可以看作是對(duì)soften label的一種正則化手段。

更直觀的實(shí)驗(yàn)請(qǐng)查閱Knowledge Distillation From Scratch[2]

Distill BERT

看到的第一篇針對(duì) BERT 模型做蒸餾的是Distilling Task-Specific Knowledge from BERT into Simple Neural Networks[3]。

在這篇論文中，作者延續(xù)Hinton 的思路在BERT 上做實(shí)驗(yàn)，首先用BERT-12 做Teacher，然后用一個(gè)單層Bi-LSTM 做Student，loss 上除了ground truth 外，也選擇了使用teacher 的logits，包括Temperature 平滑后的soften labels 的CrossEntropy和 logits 之間的MSE，最后實(shí)驗(yàn)驗(yàn)證MSE效果優(yōu)于CE。

此外，由于是從頭開(kāi)始訓(xùn)練Student，所以只用任務(wù)相關(guān)數(shù)據(jù)會(huì)嚴(yán)重樣本不足，所以作者提出了三種NLP的任務(wù)無(wú)關(guān)的data augment策略：

mask：隨機(jī)mask一部分token作為新樣本，讓teacher去生成對(duì)應(yīng)logits ;

根據(jù)POS標(biāo)簽去替換，得到 ”What do pigs eat?" -> " How do pigs ear?"

n-gram采樣：隨機(jī)選取n-gram，n取[1-5]，丟棄其余部分。

在Distilling the Knowledge in a Neural Network[4]中曾指出 logits 之間的CrossEntropy是可以看作是MSE 的近似版本，不過(guò)這里作者的結(jié)論是MSE 更好。

此外，由于Hinton 實(shí)驗(yàn)時(shí)是巨大數(shù)據(jù)量，所以不存在樣本不足的情況，而普通實(shí)驗(yàn)時(shí)都會(huì)遇到遷移時(shí)訓(xùn)練樣本不足，需要做數(shù)據(jù)增強(qiáng)的問(wèn)題。

TinyBERT

TinyBERT 出自TinyBERT: Distilling BERT for Natural Language Understanding[5]。

由于Transformer 結(jié)構(gòu)在NLP 任務(wù)中的強(qiáng)大能力，作者選擇用與BERT 同結(jié)構(gòu)的方式做Student。此外，為了提高KD后模型性能，做了更細(xì)致的工作：

Student選擇一個(gè)更窄更淺的transformer;

將KD也分為兩個(gè)階段：pre-train 和 fine-tuning，并且在兩個(gè)階段上都進(jìn)行KD;

使用了更多的loss：Embedding之間的MSE，Attention Matrix中的logits之間的MSE，Hidden state之間的MSE以及最后的分類(lèi)層的CE;

為了提高下游任務(wù)fine-tuning后的性能，使用了近義詞替換的策略進(jìn)行數(shù)據(jù)增強(qiáng)。

優(yōu)點(diǎn)

6層transformer基本達(dá)到了bert-12的性能，并且hidden size更小，實(shí)際是比bert-6更小的;

因?yàn)橛衟re-train KD，所以可以拿來(lái)當(dāng)bert 一樣直接在下游fine-tuning。

缺點(diǎn)

由于hidden size的不同，所以為了進(jìn)行MSE，需要用一個(gè)參數(shù)矩陣W 來(lái)調(diào)節(jié)，這個(gè)參數(shù)只在訓(xùn)練時(shí)使用，訓(xùn)練完后丟棄，這個(gè)矩陣沒(méi)有任何約束，覺(jué)得不優(yōu)雅;

其次，student model的每一層都需要去學(xué)習(xí)teacher model的對(duì)應(yīng)的block的輸出，如何對(duì)不同的層如何設(shè)計(jì)更好的權(quán)重也是一個(gè)費(fèi)力的事；

雖然student的結(jié)構(gòu)也是transformer，但是由于hidden size 不同，沒(méi)法使用teacher的預(yù)訓(xùn)練結(jié)果，但是我覺(jué)得這里其實(shí)可以用降維的方式用teacher的預(yù)訓(xùn)練結(jié)果，可能不需要pretraining的階段了也說(shuō)不定。

DistilBERT

DistilBERT 出自DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter[6]。

論文中作者通過(guò)調(diào)查發(fā)現(xiàn)BERT 中的hidden size 對(duì)計(jì)算效率的改變比hidden layer nums 的影響小，說(shuō)白了就是「讓模型變矮比讓模型變瘦效率更高」，所以作者使用了一個(gè)更矮的BERT來(lái)做Student 來(lái)遷移BERT 中的知識(shí)。

由于DistilBERT 是一個(gè)與BERT 同結(jié)構(gòu)只是層數(shù)更小，所以DistilBERT 可以用BERT 的預(yù)訓(xùn)練的權(quán)重進(jìn)行初始化。此外，DistilBERT 是一個(gè)與任務(wù)無(wú)關(guān)的模型，即與BERT 一樣，可以對(duì)很多下游任務(wù)進(jìn)行fine-tuning。

由于DistilBERT 與 BERT 的前幾層一致，所以loss 的選擇上就更多一些，作者選擇了triple loss：MLM loss + embedding cosin loss + soften labels cross entropy loss

優(yōu)點(diǎn)

DistilBERT 做到了與BERT 一樣，完全與任務(wù)無(wú)關(guān)，不需要添加額外的Distillation 階段（添加后結(jié)果會(huì)更好)。

MobileBERT

MobileBERT 出自MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices[7]。

作者同樣采用一個(gè)transformer 作為基本結(jié)構(gòu)，但作者認(rèn)為深度很重要，寬度較小對(duì)模型損壞較小，所以整體架構(gòu)是保持模型深度不變，通過(guò)一個(gè)矩陣來(lái)改變feature size，即bottleneck，再通過(guò)在block的前后插入兩個(gè)bottleneck，來(lái)scale feature size。

由于MobileBERT太窄太深，所以不好訓(xùn)練，作者提出新的方式，通過(guò)一個(gè)同深但是更寬的同架構(gòu)的模型來(lái)訓(xùn)練 作為teacher，然后用MobileBERT遷移。

loss 設(shè)計(jì)上主要包括三部分：feature map之間的MSE，Attention logits之間的KL，以及pre-training MLM + pre-training-NSP + pre-training-KD

訓(xùn)練策略上，有三種方式：

將KD作為附加預(yù)訓(xùn)練的附加任務(wù)，即一起訓(xùn)練；

分層訓(xùn)練，每次訓(xùn)練一層，同時(shí)凍結(jié)之前的層；

分開(kāi)訓(xùn)練，首先訓(xùn)練遷移，然后單獨(dú)進(jìn)行pre-training。

此外，為了提高推理速度，將gelu 替換為更快的 relu ，LayerNormalization 替換為 更簡(jiǎn)單的NoNorm，也做了量化的實(shí)驗(yàn)。

優(yōu)點(diǎn)

首先mobileBERT容量更小，推理更快，與任務(wù)無(wú)關(guān)，可以當(dāng)bert來(lái)直接在下游fine-tuning，而之前的KD大多數(shù)時(shí)候需要與任務(wù)綁定并使用數(shù)據(jù)增強(qiáng)，才能達(dá)到不錯(cuò)的性能；

論文實(shí)驗(yàn)非常詳實(shí)，包括如何選擇inter-block size, intra-block size, 不同訓(xùn)練策略如何影響等;

訓(xùn)練策略上，除了之前的一起訓(xùn)練完，實(shí)驗(yàn)了兩種新的訓(xùn)練方式，而最終的一層一層的訓(xùn)練與skip connection 有異曲同工的作用：每層都學(xué)一小部分內(nèi)容，從而降低學(xué)習(xí)的難度；

替換了gelu 和 LayerNormalization,進(jìn)一步提速。

缺點(diǎn)

要訓(xùn)練一個(gè)IBBERT作為teacher，而這個(gè)模型容量與BERT-Large差不多，增加了訓(xùn)練難度.

總結(jié)

以上論文的遷移過(guò)程其實(shí)可以總結(jié)為兩類(lèi)：

soft label遷移，即主要遷移Teacher 模型最后分類(lèi)層的logits 及相應(yīng)的soft label；

feature遷移，即除了最后分類(lèi)層外，還遷移Teacher 模型中的output/attention/embedding等特征。

Student 的選擇上，除了自定義外，還可以選擇跟Teacher 同結(jié)構(gòu)，而為了降低參數(shù)量，可以選擇將模型變矮/變窄/減小hidden size 等方式。

而為了蒸餾后的模型能更加的general，適應(yīng)更多的task，就需要遷移更多的信息，設(shè)計(jì)上也越復(fù)雜。

想法

實(shí)際工作上，大多數(shù)時(shí)候我們都是需要一個(gè)task 來(lái)做模型，而以上論文中告訴我們，遷移的信息越多，Student 的性能越好。

而針對(duì)具體task ，我覺(jué)得比較簡(jiǎn)潔有效的一種方式是采用更矮的Teacher 來(lái)作為Student ，這樣可以直接將Teacher 中的前幾層的信息完全遷移過(guò)來(lái)，然后在object 上，加入遷移Teacher 在train data 上的logits ，這樣就可以比較有效的進(jìn)行蒸餾了。

除此之外，讓我們換個(gè)角度看看為什么logits 能增強(qiáng)Student 模型的性能呢？除了遷移的角度外，其實(shí)logits 提供了label 更多的信息（不同類(lèi)別的相對(duì)關(guān)系），而這個(gè)額外信息只要優(yōu)于隨機(jī)分布，就能對(duì)模型提供更多的約束信息，從而增強(qiáng)模型性能，即當(dāng)前的模型可以看作是分別擬合ground truth 和 logits的兩個(gè)模型的ensemble，只不過(guò)是兩個(gè)模型共享參數(shù)。

上面我們提到只要logits 優(yōu)于隨機(jī)，對(duì)Student 模型來(lái)說(shuō)就會(huì)有所提升，那logits 由誰(shuí)產(chǎn)生的其實(shí)并不重要。所以，我們除了可以用Teacher 產(chǎn)生的logits來(lái)增強(qiáng)Student 模型外，我們還可以增強(qiáng)Teacher 模型，或者直接用Student 先學(xué)習(xí)一下，產(chǎn)生logits，再用Student 去遷移上次產(chǎn)生的logits。

想到這里，我不禁的有個(gè)大膽的想法：既然我可以一邊生成logits， 一邊學(xué)習(xí)logits，那我不是可以持續(xù)這個(gè)過(guò)程，直到模型完全擬合train data，生成的logits退化為one-hot，那此時(shí)的模型是不是能得到一個(gè)非常大的提升呢？

實(shí)驗(yàn)

實(shí)驗(yàn)的基本設(shè)置是用12層bert 作為T(mén)eacher model ，用3層bert 作為Student model 。soften labels 采用Temperature 平滑后的結(jié)果，此外，Student model 除了學(xué)習(xí) soften labels 的外，也需要學(xué)習(xí)ground truth。

Teacher-to-Student

Teacher model 在train data 上訓(xùn)練，然后在train data 上生成對(duì)應(yīng)的soften labels，Student model 學(xué)習(xí)ground truth 和 soften labels。

student-to-student

既然soften labels 是一種對(duì)labels 的一種平滑估計(jì)，那我們可以用任何方式去估計(jì)他，所以這里我們就用student 去做一個(gè)估計(jì)：student model 在train data 上進(jìn)行訓(xùn)練，然后在train data 上生成對(duì)應(yīng)的soften labels ，將 student model 利用bert 預(yù)訓(xùn)練結(jié)果重新初始化，然后去學(xué)習(xí)ground truth 和 soften labels.

normal-noise-training

既然是對(duì)labels 的一個(gè)估計(jì)，那假如給一個(gè)隨機(jī)的估計(jì)，只要保證生成的logits 中true label 對(duì)應(yīng)的值最大，就能對(duì)Student 模型進(jìn)行一定程度的提升：直接在train label 上添加一個(gè)normal noise ，然后重新進(jìn)行平滑后歸一，作為soften labels讓student model 去學(xué)習(xí)。

實(shí)驗(yàn)結(jié)果

從結(jié)果中可以看到：

優(yōu)于隨機(jī)的logits 對(duì)Student 模型有一定的提升，估計(jì)越準(zhǔn)確，提升越高；

越大的模型性能越好;

迭代進(jìn)行l(wèi)ogits 的生成與訓(xùn)練不能進(jìn)一步提高模型性能，原因主要是新的logits 分布相比之前的對(duì)模型的提升非常小，此外這個(gè)分布也比較容易擬合，所以無(wú)法進(jìn)一步提升。

責(zé)任編輯：lq