给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() 

参考链接

推荐阅读更多精彩内容