【论文阅读】SKDBERT: Compressing BERT via Stochastic Knowledge Distillation

2022-2023年论文系列之模型轻量化和推理加速

前言

通过Connected Papers搜索引用PaBEE/DeeBERT/FastBERT的最新工作，涵盖：

• 模型推理加速
• 边缘设备应用
• 生成模型
• BERT模型
• 知识蒸馏

论文目录

1. SmartBERT: A Promotion of Dynamic Early Exiting Mechanism for Accelerating BERT Inference

2. SKDBERT: Compressing BERT via Stochastic Knowledge Distillation

3. COST-EFF: Collaborative Optimization of Spatial and Temporal Efficiency with Slenderized Multi-exit Language Models

4. Knowledge Distillation with Reptile Meta-Learning for Pretrained Language Model Compression

5. Accelerating Inference for Pretrained Language Models by Unified Multi-Perspective Early Exiting

6. A Simple Hash-Based Early Exiting Approach For Language Understanding and Generation

主要内容：近两年(2022-2023年)动态早退的工作进展，粗读motivation和method以及experiment setup。

SKDBERT: Compressing BERT via Stochastic Knowledge Distillation

a. 论文信息

发表会议：AAAI2023

作者：Zixiang Ding, Guoqing Jiang, Shuai Zhang, Lin Guo, Wei Lin

发表单位：美团

开源：None

b. 内容

研究背景：已有工作表明BERT-style语言模型存在参数冗余，例如Voita等人[2019]提出减少BERT的头数量不影响性能。与此同时，许多模型压缩的方法被提出：基于参数共享、基于知识蒸馏、基于剪枝、基于量化和基于NAS，本文聚焦于基于知识蒸馏的方法。

基于 KD 的 BERT-style语言模型压缩方法之间的主要区别是：

• 学习流程：只预训练、只微调、预训练+微调
• 蒸馏目标：软输出概率、嵌入层输出、隐藏状态、自注意力分布和自注意力值的关系

motivation

Wu等人提出多教师模型蒸馏比单教师模型的KD性能更好，但是我们发现多个教师模型的集合并不总是优于单个教师模型，如下表实验结果所示：

Wu, C.; Wu, F.; and Huang, Y. 2021. One teacher is enough? pre-trained language model distillation from multiple teachers. In Findings of the Association for Computational Linguistics: ACL/IJCNLP, 4408–4413.

其原因本文认为有两个：（motivation）

1. 教师模型的集合预测失去了多样性
2. 教师模型和学生模型存在着很大的容量差距影响知识蒸馏的有效性

method

本文提出随机知识蒸馏以获得紧凑的BERT-style语言模型，名为SKDBERT。

SKD关注蒸馏范式，而不是学习流程和蒸馏目标。SKD主要内容：

1. 在每次迭代中，SKD都会从预定义的教师集合中抽取教师模型，该集合由具有多层次能力的多个教师模型组成，以一对一的方式将知识转移到学生模型中。

2. 采样分布在SKD中起着重要作用。本文呈现了三种类型的抽样分布，为每个教师模型分配适当的抽样概率。

SKD的简要过程如下图所示：

• WKD：教师集合中所有教师模型的加权对数用于优化整个训练过程中的学生模型S
• TAKD: 逐步蒸馏，$T_5$为原模型
• SKD：根据特定的概率分布，每个迭代从教师集合中随机抽取教师模型，以一对一的方式蒸馏学生模型

SKD提出的三种类型的采样方式：

• Uniform Distribution(均匀分布)

  每个教师模型的采样概率一致

• Teacher-rank Distribution（按教师等级分布）

  按照教师模型的微调性能决定教师模型的采样概率

• Student-rank Distribution（按学生等级分布）

  相对于每个教师模型蒸馏的学生模型性能决定教师模型的采样概率

SKD学习：特定的采样分布概率下教师模型和学生模型的蒸馏损失

experiment setup

数据集：GLUE（MRPC, RTE, STS-B,QQP,QNLI,MNLI）

实验环境：NVIDIA A100 GPU

SKDBERT(学生模型)设置：

• 4-layer: $TinyBER{T}_4$
• 6-layer: $TinyBER{T}_6$

教师模型：https://github.com/google-research/bert，实现14个BERT模型作为教师模型集合

学习的采样分布：

• SKDBERT4使用T04-T09模型作为教师模型

• SKDBERT6使用T10-T14模型作为教师模型

使用GLUE验证集微调的模型性能：

超参数设置：

result

1、 GLUE-test：与现有方法相比，SKDBERT4和SKDBERT6都达到了最佳的GLUE分数（提交给官方GLUE的评测分数）

SKDBERT将BERT-BASE模型的大小减小了40%，同时保持了99.5%的性能，并且速度提高了100%

2、与更多的BERT-style紧凑的模型对比： SKDBERT展现最优性能，特别是在RTE任务比最好的MiniLM多了4%准确率

3、消融实验：教师模型集合

• 教师集合应该包含强能力教师模型，同时考虑容量差距问题，以实现SKDBERT的性能。
• 性能弱于学生模型的教师模型有助于保护多样性

• 初级教师模型的数量增加对于性能有提升，但是有限（case 2）；高级教师模型的数量增加对于性能的提升是显著的（case 1）。

4、消融实验：采样分布

均匀分布适用于相似容量的教师模型集合，其中每个教师模型的知识有助于提高SKDBERT的性能。此外，教师等级和学生等级分布适用于大容量差距教师模型集合，其中初级教师模型提供的知识容易对SKDBERT产生负面影响。

5、比较WKD、TAKD 和 SKD：

SKD展现在多样性教师模型保护更好和对教师模式容量的敏感性更低的优点，如下表所示：

• WKD 使用所有教师模型的加权对数来为学生模型蒸馏做出更全面的决策，但失去了每个教师模型的多样性（Tran 等人，2020 年）。此外，TAKD在强教师模型方面存在知识消失，其中学生模型只能接受来自最强教师模型的少量知识。
• 对于WKD，小容量的教师模型几乎不会降低其性能，然而，TAKD对教师模型的容量非常敏感。

6、额外学习进程和蒸馏目标的影响： SKDBERT和DA的组合可以实现比TinyBERT更好的性能

本文的实验证明有效解决了上述提出KD的两个问题：

1. 失去多样性：每次迭代，从抽样的教师模型中一对一传递知识
2. 容量差距：利用弱教师模型来填补强教师模型与学生模型之间的容量差距，还可以避免学生模型的性能被弱教师模型限制

c. Summary

SKDBERT使用多教师蒸馏方式，引入新的蒸馏范式即SKD。

SKD范式是固定的采样概率分布，决定固定的教师模型集合中某个模型和学生模型的蒸馏。

• 考虑教师模型的多样性，实现一对一的知识传递
• 考虑教师模型与学生模型的容量差距，避免弱的教师模型限制学生模型的性能。

在GLUE基准测试集实验，实现尺寸减小近一半，速度提高100%。

SKDBERT的缺点在于 固定的采样概率、固定的教师模型集合