給Bert加速吧！NLP中的知識蒸餾論文 Distilled BiLSTM解讀

論文題目：Distilling Task-Specific Knowledge from BERT into Simple Neural Networks
論文鏈接：https://arxiv.org/pdf/1903.12136.pdf

摘要

在自然語言處理文獻中，神經網絡變得越來越深入和複雜。這一趨勢的苗頭就是深度語言表示模型，其中包括BERT、ELMo和GPT。這些模型的出現和演進甚至導致人們相信上一代、較淺的語言理解神經網絡（例如LSTM）已經過時了。然而這篇論文證明了如果沒有網絡架構的改變、不加入外部訓練數據或其他的輸入特徵，基本的“輕量級”神經網絡仍然可以具有競爭力。文本將最先進的語言表示模型BERT中的知識提煉爲單層BiLSTM，以及用於句子對任務的暹羅對應模型。在語義理解、自然語言推理和情緒分類的多個數據集中，知識蒸餾模型獲得了與ELMo的相當結果，參數量只有ELMo的大約1/100倍，而推理時間快了15倍。

1 簡介

關於自然語言處理研究中，神經網絡模型已經成了主力軍，並且模型結構層出不窮，好像永無止境一樣，這些過程中最開始的神經網絡例如LSTM變得容易被忽視。例如ELMo模型在2018年一些列任務上取得了sota效果，再到雙向編碼表示模型Bert、GPT-2在更多任務上取得了很大提升。

但是如此之大的模型在實踐落地的過程中是存在問題的：

由於參數量特別大，例如 BERT 和 GPT-2，在移動設備等資源受限的系統中是不可部署的。
由於推理時間效率低，它們也可能不適用於實時系統，對於QPS壓測很多場景基本是不過關的。
根據摩爾定律可知，我們需要在一定時間過後重新壓縮模型以及重新評估模型性能。

針對上述問題，本文提出了一種基於領域知識的高效遷移學習方法：

作者將BERT-large蒸餾到了單層的BiLSTM中，參數量減少了100倍，速度提升了15倍，效果雖然比BERT差不少，但可以和ELMo打成平手。
同時因爲任務數據有限，作者基於以下規則進行了10+倍的數據擴充：用[MASK]隨機替換單詞；基於POS標籤替換單詞；從樣本中隨機取出n-gram作爲新的樣本

2 相關工作

關於模型壓縮的背景介紹，大家可以看下李rumor的文章https://zhuanlan.zhihu.com/p/273378905，總結比較精煉和到位，這裏不再重複贅述：

Hinton在NIPS2014[1]提出了知識蒸餾（Knowledge Distillation）的概念，旨在把一個大模型或者多個模型ensemble學到的知識遷移到另一個輕量級單模型上，方便部署。簡單的說就是用小模型去學習大模型的預測結果，而不是直接學習訓練集中的label。

在蒸餾的過程中，我們將原始大模型稱爲教師模型（teacher），新的小模型稱爲學生模型（student），訓練集中的標籤稱爲hard label，教師模型預測的概率輸出爲soft label，temperature(T)是用來調整soft label的超參數。

蒸餾這個概念之所以work，核心思想是因爲好模型的目標不是擬合訓練數據，而是學習如何泛化到新的數據。所以蒸餾的目標是讓學生模型學習到教師模型的泛化能力，理論上得到的結果會比單純擬合訓練數據的學生模型要好。

在BERT提出後，如何瘦身就成了一個重要分支。主流的方法主要有剪枝、蒸餾和量化。量化的提升有限，因此免不了採用剪枝+蒸餾的融合方法來獲取更好的效果。接下來將介紹BERT蒸餾的主要發展脈絡，從各個研究看來，蒸餾的提升一方面來源於從精調階段蒸餾->預訓練階段蒸餾，另一方面則來源於蒸餾最後一層知識->蒸餾隱層知識->蒸餾注意力矩陣。

3 模型方法

本篇論文第一步選擇teacher 模型和student模型，第二步確立蒸餾程序：確立logit-regression目標函數和遷移數據集構建。

3.1 模型選擇

對於“teacher”模型，本文選擇Bert去做微調任務，比如文本分類，文本對分類等。對文本分類，可以直接將文本輸入到bert，拿到cls輸出直接softmax，可以得到每個標籤概率: $y (B) = softmax(Wh)$ ,其中 $W\in R^{k *d}$ 是softmax權重矩陣，k是類別個數。對於文本對任務，我們可以直接兩個文本輸入到Bert提取特徵，然後收入到softmax進行分類。

對於“student”模型，本文選擇的是BiLSTM和一個非線性分類器。如下圖所示：

主要流程是將文本詞向量表示，輸入到BiLSTM，選取正向和反向最後時刻的隱藏層輸出並進行拼接，然後經過一個relu輸出，輸入到softmax得到最後的概率。

3.2 蒸餾目標

$y_{i}=softmax(z)=\frac{exp(w_{i}^{T}h)}{\sum_{i}exp{W_{j}^{T}h}}$
其中 $w_{i}$ 是權重矩陣 $W$ 的第i行， $z$ 等於 $w^Th$

蒸餾的目標就是爲了最小化student模型與teacher模型的平方誤差MSE:
$L_{distill}=||Z(B)-Z(S)||_{2}^{2}$
其中 $Z(B)$ 和 $Z(S)$ 分類代表teacher和student模型的logit輸出

最終蒸餾模型的訓練函數可以將MSE損失和交叉熵損失結合起來：
$L=\alpha *L_{CE}+(1-\alpha)L_{distill}\\ =-\alpha\sum_{i}t_{i}log(y_{i}^{S})-(1-\alpha)||Z^(B)-Z^(S)||_{2}^{2}$

3.3 數據增強

用[MASK]隨機替換單詞：“I loved the comedy.”變成“I [MASK] the comedy”
基於POS標籤替換單詞；“What do pigs eat?” 變成“How do pigs eat?”
從樣本中隨機取出n-gram作爲新的樣本

4 實驗結果

本文采用的數據集爲SST-2、MNLI、QQP
實驗結果如下：

推理更加快：

5 蒸餾代碼

https://github.com/qiangsiwei/bert_distill

# coding:utf-8
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from keras.preprocessing import sequence
import pickle
from tqdm import tqdm
import numpy as np
from transformers import BertTokenizer
from utils import load_data
from bert_finetune import BertClassification


USE_CUDA = torch.cuda.is_available()
if USE_CUDA: torch.cuda.set_device(0)
FTensor = torch.cuda.FloatTensor if USE_CUDA else torch.FloatTensor
LTensor = torch.cuda.LongTensor if USE_CUDA else torch.LongTensor
device = torch.device('cuda' if USE_CUDA else 'cpu')

class RNN(nn.Module):
    def __init__(self, x_dim, e_dim, h_dim, o_dim):
        super(RNN, self).__init__()
        self.h_dim = h_dim
        self.dropout = nn.Dropout(0.2)
        self.emb = nn.Embedding(x_dim, e_dim, padding_idx=0)
        self.lstm = nn.LSTM(e_dim, h_dim, bidirectional=True, batch_first=True)
        self.fc = nn.Linear(h_dim * 2, o_dim)
        self.softmax = nn.Softmax(dim=1)
        self.log_softmax = nn.LogSoftmax(dim=1)

    def forward(self, x):
        embed = self.dropout(self.emb(x))
        out, _ = self.lstm(embed)
        hidden = self.fc(out[:, -1, :])
        return self.softmax(hidden), self.log_softmax(hidden)


class Teacher(object):
    def __init__(self, bert_model='bert-base-chinese', max_seq=128, model_dir=None):
        self.max_seq = max_seq
        self.tokenizer = BertTokenizer.from_pretrained(bert_model, do_lower_case=True)
        self.model = torch.load(model_dir)
        self.model.eval()

    def predict(self, text):
        tokens = self.tokenizer.tokenize(text)[:self.max_seq]
        input_ids = self.tokenizer.convert_tokens_to_ids(tokens)
        input_mask = [1] * len(input_ids)
        padding = [0] * (self.max_seq - len(input_ids))
        input_ids = torch.tensor([input_ids + padding], dtype=torch.long).to(device)
        input_mask = torch.tensor([input_mask + padding], dtype=torch.long).to(device)
        logits = self.model(input_ids, input_mask, None)
        return F.softmax(logits, dim=1).detach().cpu().numpy()


def train_student(bert_model_dir="/data0/sina_up/dajun1/src/doc_dssm/sentence_bert/bert_pytorch",
                  teacher_model_path="./model/teacher.pth",
                  student_model_path="./model/student.pth",
                  data_dir="data/hotel",
                  vocab_path="data/char.json",
                  max_len=50,
                  batch_size=64,
                  lr=0.002,
                  epochs=10,
                  alpha=0.5):

    teacher = Teacher(bert_model=bert_model_dir, model_dir=teacher_model_path)
    teach_on_dev = True
    (x_tr, y_tr, t_tr), (x_de, y_de, t_de), vocab_size = load_data(data_dir, vocab_path)

    l_tr = list(map(lambda x: min(len(x), max_len), x_tr))
    l_de = list(map(lambda x: min(len(x), max_len), x_de))

    x_tr = sequence.pad_sequences(x_tr, maxlen=max_len)
    x_de = sequence.pad_sequences(x_de, maxlen=max_len)

    with torch.no_grad():
        t_tr = np.vstack([teacher.predict(text) for text in t_tr])
        t_de = np.vstack([teacher.predict(text) for text in t_de])

    with open(data_dir+'/t_tr', 'wb') as fout: pickle.dump(t_tr,fout)
    with open(data_dir+'/t_de', 'wb') as fout: pickle.dump(t_de,fout)

    model = RNN(vocab_size, 256, 256, 2)

    if USE_CUDA: model = model.cuda()
    opt = optim.Adam(model.parameters(), lr=lr)
    ce_loss = nn.NLLLoss()
    mse_loss = nn.MSELoss()
    for epoch in range(epochs):
        losses, accuracy = [], []
        model.train()
        for i in range(0, len(x_tr), batch_size):
            model.zero_grad()
            bx = Variable(LTensor(x_tr[i:i + batch_size]))
            by = Variable(LTensor(y_tr[i:i + batch_size]))
            bl = Variable(LTensor(l_tr[i:i + batch_size]))
            bt = Variable(FTensor(t_tr[i:i + batch_size]))
            py1, py2 = model(bx)
            loss = alpha * ce_loss(py2, by) + (1-alpha) * mse_loss(py1, bt)  # in paper, only mse is used
            loss.backward()
            opt.step()
            losses.append(loss.item())
        for i in range(0, len(x_de), batch_size):
            model.zero_grad()
            bx = Variable(LTensor(x_de[i:i + batch_size]))
            bl = Variable(LTensor(l_de[i:i + batch_size]))
            bt = Variable(FTensor(t_de[i:i + batch_size]))
            py1, py2 = model(bx)
            loss = mse_loss(py1, bt)
            if teach_on_dev:
                loss.backward()             
                opt.step()
            losses.append(loss.item())
        model.eval()
        with torch.no_grad():
            for i in range(0, len(x_de), batch_size):
                bx = Variable(LTensor(x_de[i:i + batch_size]))
                by = Variable(LTensor(y_de[i:i + batch_size]))
                bl = Variable(LTensor(l_de[i:i + batch_size]))
                _, py = torch.max(model(bx, bl)[1], 1)
                accuracy.append((py == by).float().mean().item())
        print(np.mean(losses), np.mean(accuracy))
    torch.save(model, student_model_path)


if __name__ == "__main__":
    train_student()

給Bert加速吧！NLP中的知識蒸餾論文 Distilled BiLSTM解讀

摘要

1 簡介

2 相關工作

3 模型方法

3.1 模型選擇

3.2 蒸餾目標

3.3 數據增強

4 實驗結果

5 蒸餾代碼

參考鏈接

[轉帖]cpupower

今天，昨天，近七天，近30天，近90天，js封裝

STI比賽任務二：【答案檢驗基線方案以及思路分享】

百度搜索首屆技術創新挑戰賽：搜索模型推理優化

文本分類微調技巧實戰2.0

Kaggle 專利匹配比賽賽後總結

【CCF2022】Web攻擊檢測與分類識別 baseline

https://yachay.unat.edu.pe/blog/index.php?comment_area=format_blog&comment_component=blog&comment_co

linux以太網驅動總結