链接：https://arxiv.org/abs/2110.08536

 

导读

有效提升模型推理速度的方式是对 SOTA 的 Transformer 模型进行压缩，然而当前的推理速度仍然有待提高。本文中，作者对 RoBERTa-Large 进行了模型蒸馏，得到了一个基于 DAN 架构的、具有十亿级参数、稀疏激活的学生模型。实验表明，该模型在六个文本分类任务上保持了教师模型 97% 的表现，同时在 CPU 和 GPU 上的推理速度都提升了将近600倍。进一步的调查表明，本文提出的 pipeline 在隐私保护和域泛化设置中也很有效。

 

贡献

大型预训练 Transformer 模型进来在 NLP 领域取得了很大的成功，它们的大量推理成本推进了对模型压缩工作的研究。知识蒸馏（Knowledge Distillation）是当下比较主流的一种方式，主要思路是将教师模型（性能良好的大规模模型）的知识尽可能多的传递给学生模型（小模型）。例如 DistillBERT [1] 是 BERT 经过蒸馏得到的，大小是 BERT 的60%，推理速度则提高了60%。

然而，对于有大容量或低延迟要求的推理任务，这种加速可能仍然不够。在本文中，作者考虑通过引入稀疏蒸馏提炼出浅层、稀疏激活和过参数化的学生模型，来进一步推动推理速度的提升。与CV一般的使用“更小、更快、更便宜”学生模型的习惯相悖，本文探索了设计空间的一个新领域：使用比教师模型还大几倍的学生模型。

本文使用的学生模型基于 DAN（Deep Averaging Network） 网络 [1]，其采用简单的架构，将输入句子中的 n-gram 映射到 embedding 并进行平均池化聚合，然后再使用多个线性层进行分类。通过选择 n-gram 词汇表和 embedding 维度，DAN 可以扩展到数十亿个参数。同时，由于在训练和推理期间 DAN 被稀疏激活，因此成本可以保持在较低水平。DAN 的一个弱势在于无法像自注意力一样计算 long-range contexts，但是根据研究表明，在某些任务上 DAN 可以取得与自注意力相当的结果。

本文在六个文本分类任务上进行了实验，观察到生成的学生模型保留了97%的 RoBERTa-Large 教师模型表现。同时，作者还进一步考虑了隐私保护设置（即在蒸馏期间无法访问特定于任务的数据）和域泛化设置（即学生模型被应用并适应新数据域），发现本文的方法给 pipeline 带来了持续的提高。

 

方法

本文使用了模型蒸馏。其中教师模型选取的是在训练集上 fine-tune 好的  RoBERTa-Large 模型；学生模型基于深度平均网络 DAN，但作者对其进行了修改：不再将单词作为操作单元，而是使用 n-grams 进行代替。模型结果如下图所示，首先从输入的句子x中提取出所有的 n-grams，转化为 embeddings 后做均值化，再输入到两个全连接层中得到 logits，并通过 softmax 操作输出最终概率。该方法可以通过预处理 embeddings 来降低复杂度，但不能处理 n-gram 以外的信息。

在蒸馏时，作者同时在大规模未标注数据集C和下游任务训练集Dtrain上对齐教师模型和学生模型的输出（最小化输出概率的 KL 散度），得到的模型称为 DAN(KD)，之后可以继续在Dtrain进行进一步的 fine-tune，得到的模型称为 DAN(KD+FT)，最后这一步是可选择性的。划分 n-gram 时，作者使用了 sklearn 库中的 CountVectorizer 方法来计算每个下游任务数据集上的词汇表，并选取最频繁的V个词汇。具体地，作者将 n-gram 范围设置为1-4，并设置|V|=1,000,000。

 

实验

实验选取了六个文本分类的数据集，并选取了相应的蒸馏过程使用的未标注数据集C。具体信息如下表所示：

作者分别选取了三类模型，作为本文模型的对比实验：（1）不加入模型蒸馏，直接训练学生模型；（2）使用不基于 DAN 的学生模型结构，例如 Bi-LSTM 或 CNN；（3）直接对已有的压缩模型进行 fine-tune，例如 DistilBERT [2] 和 MobileBERT [3]。

最终得到的实验结果和分析如下：

如下表所示，在5个数据集上，教师模型和学生模型的表现均小于3%，这说明简单的 n-gram model 能力被低估了。并且模型蒸馏帮助弥补了至少一半的表现差距，尤其在 TREC 数据集上。这是一个46分类数据集，且训练集只有5452个样本。然而，本文提出的方法有效地加强了监督，并有助于解决模型训练期间的稀疏性问题。

由于设计的 DAN 模型简单且稀疏，所以预期能取得较高的加速比。在本文中，所有推理速度测试均设置了32的 batch-size。GPU 推理使用一个 Quadro RTX 8000 GPU，CPU 推理使用 56 个 Intel Xeon CPU E5-2690 v4 CPU。如下表可以观察到，DAN 学生模型有很高的性能和推理效率，在 IMDB 数据集上加速比最为显著，达到了600倍的加速，作者推断这可能是因为 IMDB 数据集平均输入长度有 300 tokens，使用自注意力机制的 transformer 架构复杂度是输入长度的平方级别，而 DAN 架构复杂度仅为线性级别，所以加速效果显著。

根据下表所示，本文的 pipeline 难以对复杂语言现象进行建模，因为 DAN 架构只关注 n-gram，并且无法对单词之间的更高级别交互进行建模。

 

参考

[1] Mohit Iyyer, Varun Manjunatha, Jordan Boyd-Graber, and Hal Daumé III. Deep unordered composition rivals syntactic methods for text classification. In ACL, pp. 1681–1691, 2015

[2]Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of BERT: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01108, 2019.

[3] https://mp.weixin.qq.com/s/Xtfn1hcT0krME1b5GUWrVg

[4]Zhiqing Sun, Hongkun Yu, Xiaodan Song, Renjie Liu, Yiming Yang, and Denny Zhou. Mobilebert: a compact task-agnostic BERT for resource-limited devices. arXiv preprint arXiv:2004.02984, 2020.

[5]https://paperswithcode.com/paper/what-are-expected-queries-in-end-to-end

[6]https://paperswithcode.com/paper/supmae-supervised-masked-autoencoders-are

[7]https://paperswithcode.com/paper/cogvideo-large-scale-pretraining-for-text-to

[8]https://hub.baai.ac.cn/view/16180

内容中包含的图片若涉及版权问题，请及时与我们联系删除