[论文解读] MiniLMv2: Multi-Head Self-Attention Relation Distillation for Compressing Pretrained Transformers

Wenhui Wang, Hangbo Bao+3 more|arXiv (Cornell University)|Dec 31, 2020

Topic Modeling参考文献 45被引用 15

一句话总结

本文提出 MiniLMv2，一种知识蒸馏方法，通过多头自注意力关系蒸馏来压缩预训练 Transformer 模型，无需教师模型与学生模型的注意力头数量匹配。通过蒸馏多个关系头中查询、键和值向量之间的细粒度缩放点积关系，该方法在单语和多语任务上均达到最先进性能，其中从 BERT Large 压缩得到的 6×768 模型在性能上超越 BERT Base，同时推理速度提升 2 倍。

ABSTRACT

We generalize deep self-attention distillation in MiniLM (Wang et al., 2020) by only using self-attention relation distillation for task-agnostic compression of pretrained Transformers. In particular, we define multi-head self-attention relations as scaled dot-product between the pairs of query, key, and value vectors within each self-attention module. Then we employ the above relational knowledge to train the student model. Besides its simplicity and unified principle, more favorably, there is no restriction in terms of the number of student's attention heads, while most previous work has to guarantee the same head number between teacher and student. Moreover, the fine-grained self-attention relations tend to fully exploit the interaction knowledge learned by Transformer. In addition, we thoroughly examine the layer selection strategy for teacher models, rather than just relying on the last layer as in MiniLM. We conduct extensive experiments on compressing both monolingual and multilingual pretrained models. Experimental results demonstrate that our models distilled from base-size and large-size teachers (BERT, RoBERTa and XLM-R) outperform the state-of-the-art.

研究动机与目标

• 解决知识蒸馏中学生模型与教师模型必须具有相同注意力头数量的限制。
• 通过利用细粒度自注意力关系而非仅注意力分布，提升模型压缩效果。
• 识别适合蒸馏的大模型最优教师层，尤其针对深层模型。
• 通过任务无关蒸馏，在单语和多语设置中实现最先进性能。
• 通过解耦注意力头数量与教师模型的关系，实现灵活的学生模型设计。

提出的方法

• 该方法将多头自注意力关系定义为每个自注意力模块内查询、键和值向量对之间的缩放点积操作。
• 引入关系头机制，将注意力头输出拼接后分割为用户指定数量的关系头，从而实现不同头数模型之间的对齐。
• 通过最小化教师模型与学生模型之间 Q-Q、K-K 和 V-V 关系矩阵的 L2 损失来执行知识蒸馏。
• 该方法允许从教师模型的任意层进行蒸馏，而不仅限于最后一层，实验表明对于深层模型，上中层通常能获得更优性能。
• 采用统一的训练目标，结合掩码语言建模与关系蒸馏，对学生模型无架构限制。
• 关系头数量无需是学生或教师头数的倍数，从而实现灵活且细粒度的知识迁移。

实验结果

研究问题

• RQ1 自注意力关系蒸馏是否能在不强制教师与学生模型具有相同注意力头数量的前提下，提升模型压缩效果？
• RQ2 从大模型的上中层（如 24 层中的第 12 层）蒸馏知识是否比使用最后一层获得更优性能？
• RQ3 增加关系头数量如何影响蒸馏后学生模型的性能？
• RQ4 多头自注意力关系蒸馏是否能在单语和多语任务上均优于现有知识蒸馏方法？
• RQ5 在蒸馏过程中，使用更多自注意力关系与计算成本之间存在何种权衡？

主要发现

• 从 BERT Large 蒸馏得到的 6×768 模型在 GLUE 基准上性能超越 BERT Base，同时推理速度提升 2.0 倍。
• 从 RoBERTa Large 蒸馏得到的学生模型在使用显著更少训练样本的情况下，性能优于 RoBERTa Base。
• 使用大模型（如 BERT Large 和 XLM-R Large）的上中层（如 24 层中的第 12 层）进行蒸馏，性能优于使用最后一层。
• 增加关系头数量可提升性能，更细粒度的知识迁移有助于更精确地模仿教师模型的自注意力模块。
• 同时蒸馏 Q-Q、K-K 和 V-V 关系的效果优于仅蒸馏 Q-K 和 V-V 关系，且在大多数 GLUE 任务上优于 MobileBERT，同时推理更快。
• 该方法具有高度灵活性与泛化能力，可在无学生模型架构限制的前提下，从多种大型模型（如 BERT、RoBERTa、XLM-R）中进行蒸馏。