基于知識蒸餾的短文本分類方法

孫 紅，黃甌嚴

（上海理工大學光電信息與計算機工程學院，上海 200093）

0 引言

隨著深度學習在自然語言處理（Natural Language Processing，NLP）任務中的不斷發展與性能指標的不斷提高，NLP 任務的工業落地成為可能。然而，深度學習模型結構復雜，占用較大的存儲空間且計算資源消耗大，因此高性能模型很難直接部署在移動端。為了解決這些問題，需對模型進行壓縮以減小模型在計算時間和空間上的消耗。

模型壓縮的目的是在保證模型預測效果的前提下，盡可能減小模型體積，提升模型推演速度。常用模型壓縮方法有剪枝（Pruning）、權重分解（Weight Factorization）、削減精度（Quantization）、權重共享（Weight Sharing）及知識蒸餾（Knowledge Distillation）。其中，知識蒸餾方法能將模型壓縮至最小規模，使性能效果最佳，近年來備受關注。Nakashole 等［1］采用零次學習（Zero-shot）的方式在翻譯任務上使用知識蒸餾；Wang 等［2］在基于強化學習的對話系統中使用知識蒸餾，提高系統可維護性與拓展性；Siddhartha等［3］使用多種知識蒸餾策略結合的方式應用于問答系統（QA）響應預測；Sun 等［4］利用知識蒸餾壓縮BERT 模型，壓縮得到的模型可以用于移動端；Liu 等［5］采用樣本自適應機制，不依賴標注數據，利用自蒸餾的方式訓練模型；廖勝蘭等［6］利用卷積神經網絡蒸餾模型（Bidirectional Encoder Representation from Transformers，BERT），用于意圖識別分類；Subhabrata 等［7］在多語言命名實體識別任務中采用多階段蒸餾框架，使用階段優化的方式提高性能。

本文在現有研究基礎上，提出一種基于知識蒸餾的短文本分類模型，其中教師模型為BERT 模型，學生模型為雙向長短時記憶網絡模型（Bi-directional Long Short-Term Memory，BiLSTM）。實驗結果表明，經過知識蒸餾的學生模型比單獨訓練的學生模型分類效果更佳，并且與復雜的教師模型相比，本文模型可極大降低預測所需的響應時間，有利于模型在工業場景中有效部署與使用。

1 相關工作

很多NLP 任務場景可以歸結為文本分類任務。依據分類對象，文本分類可分為短文本分類（標題、評論）和長文本分類（文章、文檔）；依據分類體系，文本分類可分為新聞分類、情感分析和意圖識別等；依據分類模式，文本分類可分為二分類問題、多分類問題以及多標簽問題（一個文本屬于多個類別）。

傳統文本分類方法常使用基于規則的特征匹配，或依賴專家系統，往往能做到快速分類，然而與數據集所屬領域高度相關，需要不同領域的專家構建特定規則，耗力費時，且準確率并不高，無法達到工業要求。

隨著統計學習方法的興起及互聯網文本數據集數量爆炸式增長，利用機器學習處理文本分類問題成為主流。機器學習方法一般包含3 個步驟：文本預處理、文本特征提取和分類模型分類。文本預處理首先對文本進行分詞，再建立停用詞詞典，去除副詞、形容詞以及連接詞，有些任務還需要進行詞性標注，對分詞后得到的詞直接判斷詞性；文本特征提取方式可以考慮詞頻，在一段文本中反復出現越多的詞越重要，權重越大，也可以考慮詞的重要性，以TF-IDF（Term Frequency-inverse Document Frequency）作為特征，表征詞重要程度；分類模型通常有邏輯回歸模型（Logistic Regression，LR）、支持向量機（Support Vector Machine，SVM）、隨機森林（RandomForest，RF）等。

機器學習的方法雖然在文本分類上取得了較好效果，但也存在問題，文本特征提取得到的文本表示是高緯度、高稀疏的，表達特征能力很弱，且往往需要人工進行特征工程，成本很高。于是深度學習的方法被應用于文本分類任務，用端到端的方式解決復雜耗時的人工特征工程。深度學習文本分類模型包括訓練速度快的FastText 模型［8］、利用CNN 提取句子關鍵信息的TextCNN 模型［9］、利用雙向RNN 得到每個詞上下文表示的TextRNN 模型［10］及基于層次注意力機制網絡的HAN 模型［11］等。

2 模型結構

知識蒸餾短文本分類模型主要由兩個子模型組成：教師模型與學生模型。其中教師模型直接學習真實數據標簽，學生模型為結構精簡的小模型，蒸餾模型由學生模型通過學習教師模型的結果并結合真實標簽的分布構建而成。

2.1 教師模型

教師模型（Teacher Model）通常為結構相對復雜的模型，具有很好的泛化能力。本文選用BERT 模型［12］作為教師模型。

雙向編碼模型BERT 采用多層雙向Transformer 編碼器為主體進行訓練，舍棄RNN 等循環神經網絡，采用注意力機制對文本進行建模，可捕捉更長距離的依賴。BERT 使用深而窄的神經網絡，中間層有1 024 個神經元，層數有12層，并采用無監督學習的方式，無需人工干預和標注，使用大規模語料進行訓練，其模型結構如圖1 所示。

Fig.1 The structure of the BERT model圖1 BERT 模型結構

文獻［12］將上文信息和下文信息獨立編碼再進行拼接，但Devlin 等［13］說明了同時編碼上下文信息的重要性。BERT 模型聯合所有層上下文進行訓練，使模型能很好地結合上下文理解語義。預訓練好的模型只需進行參數微調即可快速適應多種類型的下游具體任務。

本文選用哈工大訊飛聯合實驗室（HFL）發布的基于全詞Mask 的中文預訓練模型BERT-wwm-ext1。該預訓練模型收集超大量語料用于預訓練，包括百科、問答、新聞等通用語料，總詞數達到5.4B。BERT-wwm-ext 采用與BERT 相同的模型結構，由12 層Transformer 構成，訓練第一階段（最大長度為128）采用的batchsize 為2 560，訓練1M 步；訓練第二階段（最大長度為512）采用的batchsize 為384，訓練400K步。

為了更直觀測試模型蒸餾效果，本文實驗僅選用1 層全連接層作為分類器，對短文本類別進行分類，并對教師模型BERT 最后4 層進行微調。

2.2 學生模型

盡管教師模型性能良好，但其模型規模往往很大，訓練過程需消耗大量計算資源，甚至由多個模型集成而成。由于教師模型推斷速度慢，對內存、顯存等資源要求高，因此需構建結構相對簡單的學生模型（Student Model）學習教師模型學到的知識。

單獨訓練學生模型往往無法達到與教師模型一樣或相當的效果，因此本文將學生模型與教師模型建立聯系，通過學習教師模型的輸出訓練學生模型。

本文選用單層雙向長短時記憶網絡（BiLSTM）作為學生模型，采用1 層全連接層作分類器，模型結構如圖2 所示。輸入為短文本句向量x，hl和hr分別為雙向LSTM 隱層輸出，預測結果為輸出y。

Fig.2 Student model structure diagram圖2 學生模型結構

在學習上下文相關信息時，通常使用循環神經網絡（Recurrent Neural Network，RNN），然而標準RNN 存儲的上下文信息有限，并在網絡結構較深時存在梯度消失的問題。為了解決這些問題，Hochreiter 等［14］提出了長短時記憶網絡（LSTM），通過訓練可以使LSTM 學習記憶有效信息并遺忘無效信息，更好地捕捉長距離依賴關系。而雙向長短時記憶網絡可從后往前地對信息進行編碼，更好地捕捉雙向語義依賴。

對于單獨訓練學生模型和結合教師模型訓練的知識蒸餾模型，本文對學生模型采取相同的結構，以便進行性能對比。

2.3 知識蒸餾模型

知識蒸餾是一種模型壓縮方法，最早由Hinton 等［15］在計算機視覺領域提出。由于計算資源昂貴，因此本文選用規模更小的模型，消耗更小的計算代價達到期望的性能。但單獨訓練規模小的模型很難達到預期效果，所以將大規模教師模型學習到的細粒度知識遷移至學生模型訓練中。

對于分類問題，本文將真實的標簽數據稱為“硬標簽”，即每1 個數據屬于某類別的概率為1，屬于其他類別的概率為0。然而硬標簽包含的信息量很低，真實數據往往包含一定量其他標簽信息。例如在圖像分類識別的任務中，由于狗和貓具有相似特征，狗被預測為貓的概率遠大于預測為手機的概率。具體來說，1 張長得像貓的狗圖片則蘊含更多信息量，而硬標簽僅給出了這張照片屬于狗這一類別的分類信息。Hinton 將教師模型輸出的softmax 結果作為“軟標簽”，軟標簽有較高的信息熵，學生模型可通過學習軟標簽提高自身泛化能力。

知識蒸餾模型采用教師-學生結構，如圖3 所示，教師模型輸出知識，學生模型接受知識。預訓練教師模型使用的數據集與知識蒸餾模型使用的數據集相同，模型具體實現步驟如下。

Step 1.使用真實數據集D中的硬標簽訓練教師模型T，超參調優得到性能較好的模型。

Step 2.利用訓練好的教師模型T計算軟標簽。

Step 3.結合真實數據集D中的硬標簽以及上一步驟計算得到的軟標簽，訓練學生模型S，損失函數如公式（1）所示。

Step 4.學生模型S的預測與常規方式相同。

Fig.3 Structure diagram of knowledge distillation model圖3 知識蒸餾模型結構

本文選用的知識蒸餾模型損失函數Loss 由兩部分構成。第一部分為硬標簽與學生模型輸出的交叉熵LossCE，第二部分為軟標簽與學生模型輸出logits 的均方差Lossdistill。

其中α為兩部分損失的平衡參數，si為學生模型輸出，yi為真實標簽數據，zt為教師模型輸出的logits，zs為學生模型輸出的logits。

3 實驗設置

3.1 數據集及實驗環境

為驗證該模型，本文使用CLUE 上的短文本分類公開數據集TNEWS 作為實驗數據。該數據集由今日頭條中文新聞標題采集得到，包含380 000 條新聞標題，共有15 個新聞類別。實驗環境如表1 所示。

Table 1 Experimental environment表1 實驗環境

3.2 評價指標

本文使用macroF1值作為模型評價指標。

首先分別計算每個類別精度。

macro精度為所有精度平均值。

同理分別計算每個類別的召回率。

macro召回為所有召回平均值。

最后macroF1計算公式為：

其中，n為類別總數，TPi、FPi和FNi分別表示第i類對應的真正例、假正例和假反例。

3.3 模型參數

3.3.1 教師模型參數

在本文實驗中，教師模型選用在超大量語料上訓練的預訓練模型BERT-wwm-ext，并后接一層全連接層作分類。BERT-wwm-ext 模型共有12 層，隱層含有768 個神經元，使用12 頭自注意力模式。模型采用Adam 優化器進行優化，學習率為5e-4，每句話處理的長度（短填長切）為32。訓練時采用批量處理的方法，批處理大小為64。教師模型共有參數102 424 805 個。

3.3.2 學生模型參數

在本文實驗中，學生模型選用雙向長短時記憶（BiLSTM）模型，也后接一層全連接層作分類。BiLSTM 為單層雙向模型，隱層含有256 個神經元。模型輸入的句向量由其組成的詞的詞向量求和取平均得到，組成句子的詞由結巴分詞工具分詞后得到，詞向量選取用人民日報預訓練好的300 維詞向量［16］。模型使用SGD 作為優化器，學習率為0.05。訓練時采用批量處理的方法，批處理大小為64。學生模型共有參數1 209 093 個。

3.3.3 蒸餾模型參數

在本文實驗中，蒸餾模型聯合教師-學生模型進行訓練，硬標簽采用交叉熵作為損失函數，軟標簽采用均方差（MSE）作為損失函數，平衡參數α選取為0.2。

4 實驗結果及分析

4.1 短文本分類任務

本文使用TNEWS 數據集進行分類實驗，該數據集在各類別中存在非常嚴重的不平衡問題，且本文關注的重點為知識蒸餾模型效果，過多的類別數量也會產生影響。為了防止上述問題對實驗產生影響，本文選擇汽車、文化、教育、游戲、體育5 類數據進行實驗，如表2 所示。

Table 2 Statistics of balanced datasets表2 實驗平衡數據

首先，分別將教師模型與學生模型進行單獨訓練，得到原始模型macroF1結果，再用蒸餾模型對知識進行蒸餾，結果如表3 所示。由表3 可知教師模型BERT-wwm-ext 在微調后準確率可達81.00%，而學生模型BiLSTM 只有75.67%，即教師模型具有更深的網絡層數和更多參數，從原始數據中學習到了更多知識，具有更好的模型泛化能力；而學生模型結構簡單，僅通過超參優化無法達到較好的效果。在教師模型的指導學習下，蒸餾模型macroF1值可達78.83%，比單獨訓練學生模型提高3.16%。可見學生模型不僅可自主地從硬標簽學習知識，還從教師模型獲取了一部分知識。由于模型結構簡單等原因，蒸餾模型分類結果無法超越教師模型，但與單獨訓練學生模型相比，性能明顯提升。

Table 3 Classification results of each model表3 各模型分類效果

各模型在不同類別中的分類結果如圖4 所示。其中，橫坐標為測試數據的不同類別，縱坐標為測試性能指標macroF1值。由圖4 可知，蒸餾模型在教師模型的指導下，在所有5 個類別上分類效果均優于單獨訓練的學生模型，并且在汽車類別上十分接近教師模型分類結果，這表明在一些區分度高的類別上，簡單模型通過一定方式學習后可達到復雜模型效果。所有模型在文化類別中分類效果均相對較差，這是因為文化類語句相對其他類別的語句更加抽象復雜，往往包含其他類別的含義，模型無法學習到深層次特征從而抽象表達這些語句，導致分類結果難以提升。

Fig.4 Test results in various categories圖4 模型在各個類別測試結果

4.2 推理時間分析

本文使用不同模型在相同測試集中進行實驗，分析模型時間性能。教師模型與學生模型在3 000 條測試數據中進行實驗的推理時間對比結果如表4 所示。

Table 4 Runtime to complete one itevation表4 完成1 次迭代的推理時間

從表4 可以看出，學生模型在完成1 次推理的時間遠少于教師模型，所需時間僅有教師模型的1/725，這主要是因為與教師模型相比，學生模型結構相對簡單，模型參數只有教師模型的1/85 倍。知識蒸餾模型在準確率接近教師模型的情況下，推理時間更短，有利于在真實場景（如移動端）的部署。

5 結語

本文針對結構復雜模型難以落地應用的現狀，提出一種基于教師—學生框架的知識蒸餾模型，應用于短文本分類任務。該模型首先預訓練1 個分類性能好、結構復雜的大模型，再將大模型所學知識遷移至結構簡單的小模型中，以此彌補小模型單獨訓練時泛化能力不足的問題。實驗結果表明，知識蒸餾小模型性能顯著改善，同時，模型迭代推理時間大幅縮短，使模型在工業場景中進行應用成為可能。

目前開源中文數據集較少，多為爬蟲獲得，數據集質量較差，人工標記數據又有耗力費時等問題。本文使用的分類器僅使用了1 層全連接層，學生模型也選擇了較為簡單的單層雙向長短時記憶網絡。針對上述問題，下一步工作是尋找合適的生成方式以產生高質量數據，以及如何選擇較為復雜的適用于工業場景的學生模型。