MLKD-BERT: Multi-level Knowledge Distillation for Pre-trained Language Models

📄 arXiv: 2407.02775v1 📥 PDF

作者: Ying Zhang, Ziheng Yang, Shufan Ji

分类: cs.CL, cs.LG

发布日期: 2024-07-03

💡 一句话要点

提出MLKD-BERT，通过多层次知识蒸馏提升预训练语言模型性能并降低推理时间。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 知识蒸馏 预训练语言模型 BERT 模型压缩 模型加速 多层次知识 注意力机制

📋 核心要点

  1. 现有知识蒸馏方法在关系层面知识挖掘不足，且学生模型注意力头数量设置不够灵活，限制了模型性能和推理效率。
  2. MLKD-BERT通过多层次知识蒸馏，在师生框架下提取更丰富的关系知识，并允许灵活设置学生模型的注意力头数量。
  3. 实验表明，MLKD-BERT在GLUE和抽取式问答任务上优于现有方法，并能在保证性能的同时显著降低推理时间。

📝 摘要（中文）

知识蒸馏是一种有效的预训练语言模型压缩技术。虽然现有的知识蒸馏方法在BERT等典型模型上表现良好，但仍有两个方面可以改进：关系层面的知识可以进一步挖掘以提升模型性能；学生模型的注意力头数量设置可以更加灵活以减少推理时间。因此，我们提出了一种新的知识蒸馏方法MLKD-BERT，在师生框架下蒸馏多层次知识。在GLUE基准和抽取式问答任务上的大量实验表明，我们的方法优于BERT上最先进的知识蒸馏方法。此外，MLKD-BERT可以灵活设置学生模型的注意力头数量，从而在性能损失很小的情况下大幅减少推理时间。

🔬 方法详解

问题定义：论文旨在解决现有BERT知识蒸馏方法在关系层面知识利用不足以及学生模型结构不够灵活的问题。现有方法通常只关注最后一层的输出或注意力权重，忽略了中间层蕴含的关系知识。此外，学生模型的结构通常与教师模型保持一致，限制了模型压缩和推理加速的潜力。

核心思路：论文的核心思路是通过多层次知识蒸馏，将教师模型中不同层次的知识传递给学生模型。具体来说，不仅考虑最后一层的输出，还关注中间层的表示和关系信息。同时，允许学生模型具有更灵活的注意力头数量，从而在性能和效率之间取得更好的平衡。

技术框架：MLKD-BERT采用标准的师生框架。教师模型是预训练的BERT模型，学生模型是需要压缩和加速的模型。蒸馏过程包括多个层次的知识传递：1) 输出层蒸馏：使用传统的softmax输出蒸馏；2) 表示层蒸馏：将教师模型和学生模型中间层的表示进行对齐；3) 关系层蒸馏：利用注意力机制提取关系知识，并进行蒸馏。

关键创新：MLKD-BERT的关键创新在于多层次知识蒸馏策略，特别是关系层蒸馏。通过引入关系层蒸馏，模型能够学习到更丰富的关系知识，从而提升性能。此外，允许学生模型灵活设置注意力头数量，使得模型能够在性能和效率之间进行权衡。

关键设计：在关系层蒸馏中，论文使用注意力矩阵来表示关系知识。具体来说，将教师模型和学生模型的注意力矩阵进行对齐，使用KL散度作为损失函数。此外，论文还设计了一种自适应的注意力头选择策略，根据不同的任务和资源限制，选择合适的注意力头数量。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明，MLKD-BERT在GLUE基准测试中取得了显著的性能提升，优于现有的知识蒸馏方法。例如，在某些任务上，MLKD-BERT的性能提升超过1%。此外，通过灵活设置学生模型的注意力头数量，MLKD-BERT可以在性能损失很小的情况下，将推理时间降低30%以上。这些结果表明，MLKD-BERT是一种有效的BERT模型压缩和加速方法。

🎯 应用场景

MLKD-BERT可应用于各种需要BERT模型进行推理的场景，尤其适用于资源受限的设备或对推理速度有较高要求的应用，如移动设备上的自然语言处理任务、实时对话系统、以及需要快速响应的搜索服务等。通过知识蒸馏，可以在保证模型性能的同时，显著降低模型大小和推理时间，从而提高用户体验和降低部署成本。

📄 摘要（原文）

Knowledge distillation is an effective technique for pre-trained language model compression. Although existing knowledge distillation methods perform well for the most typical model BERT, they could be further improved in two aspects: the relation-level knowledge could be further explored to improve model performance; and the setting of student attention head number could be more flexible to decrease inference time. Therefore, we are motivated to propose a novel knowledge distillation method MLKD-BERT to distill multi-level knowledge in teacher-student framework. Extensive experiments on GLUE benchmark and extractive question answering tasks demonstrate that our method outperforms state-of-the-art knowledge distillation methods on BERT. In addition, MLKD-BERT can flexibly set student attention head number, allowing for substantial inference time decrease with little performance drop.