给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) = s o f t m a x (W h)$ ,其中 $W∈Rk∗dW\in R^{k *d}$ 是softmax权重矩阵，k是类别个数。对于文本对任务，我们可以直接两个文本输入到Bert提取特征，然后收入到softmax进行分类。

对于“student”模型，本文选择的是BiLSTM和一个非线性分类器。如下图所示：

主要流程是将文本词向量表示，输入到BiLSTM，选取正向和反向最后时刻的隐藏层输出并进行拼接，然后经过一个relu输出，输入到softmax得到最后的概率。

3.2 蒸馏目标

$yi=softmax(z)=exp(wiTh)∑iexpWjThy_{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=α∗LCE+(1−α)Ldistill=−α∑itilog(yiS)−(1−α)∣∣Z(B)−Z(S)∣∣22L=\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 BertClassificationUSE_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_dimself.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_seqself.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 usedloss.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()

参考链接

【经典简读】知识蒸馏(Knowledge Distillation) 经典之作
【论文笔记】Distilling Task-Specific Knowledge from BERT into Simple Neural Networks
知识蒸馏论文选读（二）