TinyBERT 是華為不久前提出的一種蒸餾 BERT 的方法，本文梳理了 TinyBERT 的模型結(jié)構(gòu)，探索了其在不同業(yè)務(wù)上的表現(xiàn)，證明了 TinyBERT 對(duì)復(fù)雜的語義匹配任務(wù)來說是一種行之有效的壓縮手段。
作者：chenchenliu&winsechang，騰訊 PCG 內(nèi)容挖掘工程師
來源：騰訊技術(shù)工程微信號(hào)

一、簡(jiǎn)介

在 NLP 領(lǐng)域，BERT 的強(qiáng)大毫無疑問，但由于模型過于龐大，單個(gè)樣本計(jì)算一次的開銷動(dòng)輒上百毫秒，很難應(yīng)用到實(shí)際生產(chǎn)中。TinyBERT 是華為、華科聯(lián)合提出的一種為基于 transformer 的模型專門設(shè)計(jì)的知識(shí)蒸餾方法，模型大小不到 BERT 的 1/7，但速度提高了 9 倍，而且性能沒有出現(xiàn)明顯下降。目前，該論文已經(jīng)提交機(jī)器學(xué)習(xí)頂會(huì) ICLR 2020。本文復(fù)現(xiàn)了 TinyBERT 的結(jié)果，證明了 Tiny BERT 在速度提高的同時(shí)，對(duì)復(fù)雜的語義匹配任務(wù)，性能沒有顯著下降。

目前主流的幾種蒸餾方法大概分成利用 transformer 結(jié)構(gòu)蒸餾、利用其它簡(jiǎn)單的結(jié)構(gòu)比如 BiLSTM 等蒸餾。由于 BiLSTM 等結(jié)構(gòu)簡(jiǎn)單，且一般是用 BERT 最后一層的輸出結(jié)果進(jìn)行蒸餾，不能學(xué)到 transformer 中間層的信息，對(duì)于復(fù)雜的語義匹配任務(wù)，效果有點(diǎn)不盡人意。

基于 transformer 結(jié)構(gòu)的蒸餾方法目前比較出名的有微軟的 BERT-PKD (Patient Knowledge Distillation for BERT)，huggingface 的 DistilBERT，以及本篇文章講的 TinyBERT。他們的基本思路都是減少 transformer encoding 的層數(shù)和 hidden size 大小，實(shí)現(xiàn)細(xì)節(jié)上各有不同，主要差異體現(xiàn)在 loss 的設(shè)計(jì)上。

二、模型實(shí)現(xiàn)細(xì)節(jié)

整個(gè) TinyBERT 的 loss 設(shè)計(jì)分為三部分：

1. Embedding-layer Distillation

其中:

分別代表 student 網(wǎng)絡(luò)的 embedding 和 teacher 網(wǎng)絡(luò)的 embedding. 其中 l 代表 sequence length, d0 代表 student embedding 維度， d 代表 teacher embedding 維度。由于 student 網(wǎng)絡(luò)的 embedding 層通常較 teacher 會(huì)變小以獲得更小的模型和加速，所以 We 是一個(gè) d 0×d 維的可訓(xùn)練的線性變換矩陣，把 student 的 embedding 投影到 teacher embedding 所在的空間。最后再算 MSE，得到 embedding loss.

2. Transformer-layer Distillation

TinyBERT 的 transformer 蒸餾采用隔 k 層蒸餾的方式。舉個(gè)例子，teacher BERT 一共有 12 層，若是設(shè)置 student BERT 為 4 層，就是每隔 3 層計(jì)算一個(gè) transformer loss. 映射函數(shù)為 g(m) = 3 * m, m 為 student encoder 層數(shù)。具體對(duì)應(yīng)為 student 第 1 層 transformer 對(duì)應(yīng) teacher 第 3 層，第 2 層對(duì)應(yīng)第 6 層，第 3 層對(duì)應(yīng)第 9 層，第 4 層對(duì)應(yīng)第 12 層。每一層的 transformer loss 又分為兩部分組成，attention based distillation 和 hidden states based distillation.

2.1 Attention based loss

其中，

h 代表 attention 的頭數(shù)，l 代表輸入長(zhǎng)度，

代表 student 網(wǎng)絡(luò)第 i 個(gè) attention 頭的 attention score 矩陣，

代表 teacher 網(wǎng)絡(luò)第 i 個(gè) attention 頭的 attention score 矩陣。這個(gè) loss 是受到斯坦福和 Facebook 聯(lián)合發(fā)表的論文，What Does BERT Look At? An Analysis of BERT’s Attention 的啟發(fā)。這篇論文研究了 attention 權(quán)重到底學(xué)到了什么，實(shí)驗(yàn)發(fā)現(xiàn)與語義還有語法相關(guān)的詞比如第一個(gè)動(dòng)詞賓語，第一個(gè)介詞賓語，以及[CLS], [SEP], 逗號(hào)等 token，有很高的注意力權(quán)重。為了確保這部分信息能被 student 網(wǎng)絡(luò)學(xué)到，TinyBERT 在 loss 設(shè)計(jì)中加上了 student 和 teacher 的 attention matrix 的 MSE。這樣語言知識(shí)可以很好的從 teacher BERT 轉(zhuǎn)移到 student BERT.

2.2 hidden states based distillation

其中，

分別是 student transformer 和 teacher transformer 的隱層輸出。和 embedding loss 同理，

投影到 Ht 所在的空間。

3. Prediction-Layer Distillation

其中 t 是 temperature value，暫時(shí)設(shè)為 1.除了模仿中間層的行為外，這一層用來模擬 teacher 網(wǎng)絡(luò)在 predict 層的表現(xiàn)。具體來說，這一層計(jì)算了 teacher 輸出的概率分布和 student 輸出的概率分布的 softmax 交叉熵。這一層的實(shí)現(xiàn)和具體任務(wù)相關(guān)，我們的兩個(gè)實(shí)驗(yàn)分別采取了 BERT 原生的 masked language model loss + next sentence loss 和單任務(wù)的 classification softmax cross-entropy.

另外，值得一提的是 prediction loss 有很多變化。在 TinyBERT 中，這個(gè) loss 是 teacher BERT 預(yù)測(cè)的概率和 student BERT 預(yù)測(cè)概率的 softmax 交叉熵，在 BERT-PKD 模型中，這個(gè) loss 是 teacher BERT 和 student BERT 的交叉熵和 student BERT 和 hard target( one-hot)的交叉熵的加權(quán)平均。我們?cè)跇I(yè)務(wù)中有試過直接用 hard target loss，效果比使用 teacher student softmax 交叉熵下降 5-6 個(gè)點(diǎn)。因?yàn)?softmax 比 one-hot 編碼了更多概率分布的信息。并且實(shí)驗(yàn)中，softmax cross-entropy loss 容易發(fā)生不收斂的情況，把 softmax 交叉熵改成 MSE, 收斂效果變好，但泛化效果變差。這是因?yàn)槭褂?softmax cross-entropy 需要學(xué)到整個(gè)概率分布，更難收斂，因?yàn)閿M合了 teacher BERT 的概率分布，有更強(qiáng)的泛化性。MSE 對(duì)極值敏感，收斂的更快，但泛化效果不如前者。

所以總結(jié)一下，loss 的計(jì)算公式為:

其中，

三、實(shí)驗(yàn)

TinyBERT 論文中提出了兩階段學(xué)習(xí)框架，比較新穎。類似于原生的 BERT 先 pre-train, 根據(jù)具體任務(wù)再 fine-tine, TinyBERT 先在 general domain 數(shù)據(jù)集上用未經(jīng)微調(diào)的 BERT 充當(dāng)教師蒸餾出一個(gè) base 模型，在此基礎(chǔ)上，具體任務(wù)通過數(shù)據(jù)增強(qiáng)，利用微調(diào)后的 BERT 再進(jìn)行重新執(zhí)行蒸餾。

這種兩階段的方法給 TinyBERT 提供了像 BERT 一樣的泛化能力。不過為了快速得到實(shí)驗(yàn)結(jié)果，并且論文中的控制變量實(shí)驗(yàn)顯示 general 的蒸餾對(duì)各項(xiàng)下游任務(wù)的影響較小，我們此次選擇直接用 fine-tune 過的 teacher BERT，蒸餾得到 student BERT.

所以我們蒸餾 TinyBERT 的流程是:

制作任務(wù)相關(guān)數(shù)據(jù)集;

fine-tune teacher BERT;

固定 teacher BERT 參數(shù)，蒸餾得到 TinyBERT.

關(guān)于實(shí)驗(yàn)結(jié)果，先上 TinyBERT 論文中的結(jié)論：

可以看到 TinyBERT 表現(xiàn)優(yōu)異。在 GLUE 上，相較于完整的 BERT，性能下降 3 個(gè)點(diǎn)，但是推理性能卻得到了巨大提升，快了 9 倍多。

我們?cè)谧约旱臉I(yè)務(wù)上，也用 TinyBERT 得到了相似的結(jié)果。

3.1 文章連貫性特征任務(wù)

做這個(gè)特征的目的是為了過濾東拼西湊或者機(jī)器生成前后沒有邏輯的文章。由于語義的復(fù)雜性還有語義的轉(zhuǎn)移，這個(gè)任務(wù)和語義相似度任務(wù)略有不同，文章的上下句之間語義會(huì)有不同。在這個(gè)背景下，實(shí)驗(yàn)過 DSSM, Match-Pyramid 等模型，表現(xiàn)效果較差。

由于 BERT 能學(xué)到豐富的語義，這個(gè)任務(wù)目前采用 BERT 的 next sentence 任務(wù)較為合適。但是一旦文章很長(zhǎng)，原生 BERT 需要算 1 秒甚至更久，這樣的速度是不能接受的。

TinyBERT 在不同實(shí)驗(yàn)參數(shù)下的表現(xiàn)如下：

可以看到 4 層 encoder 的 TinyBERT 在 next sentence 任務(wù)下準(zhǔn)確率較 BERT base 準(zhǔn)確率下降了不到 3 個(gè)點(diǎn)，在 mlm 任務(wù)上下降較多。在 CPU 上，TinyBERT 相較于 base 速度獲得了將近 8 倍的提升。

3.2 問答 FAQ 任務(wù)

業(yè)務(wù)場(chǎng)景：為用戶的 query 匹配最接近的 question，將其 answer 返回，是一個(gè) Query-question 語義匹配任務(wù)。

下面是蒸餾到兩層 encoder 的 TinyBERT 結(jié)果：

可以看到 ACC 損失 3 個(gè)點(diǎn)，AUC 損失 4 個(gè)點(diǎn)，取得了不錯(cuò)的效果。

四、總結(jié)

我們證明了 TinyBERT 作為一種蒸餾方法，能有效的提取 BERT transformer 結(jié)構(gòu)中豐富的語意信息，在不犧牲性能的情況下，速度能獲得 8 到 9 倍的提升。下一步可能會(huì)嘗試蒸餾一個(gè) general 的 TinyBERT base。

更多騰訊AI相關(guān)技術(shù)干貨，請(qǐng)關(guān)注專欄騰訊技術(shù)工程

審核編輯 黃昊宇