[论文解读] Extract then Distill: Efficient and Effective Task-Agnostic BERT Distillation

Cheng Chen, Yichun Yin+5 more|arXiv (Cornell University)|Apr 24, 2021

Topic Modeling参考文献 30被引用 4

一句话总结

本文提出了一种名为提取后蒸馏（Extract Then Distill, ETD）的方法，这是一种灵活且高效的无任务依赖BERT蒸馏方法，通过重用预训练教师模型的参数来初始化更小的学生模型。通过基于重要性评分或随机选择提取并重用最关键的神经元（前馈网络、注意力头、隐藏状态），ETD将训练成本降低了70%，同时在性能上优于基线蒸馏方法。

ABSTRACT

Task-agnostic knowledge distillation, a teacher-student framework, has been proved effective for BERT compression. Although achieving promising results on NLP tasks, it requires enormous computational resources. In this paper, we propose Extract Then Distill (ETD), a generic and flexible strategy to reuse the teacher's parameters for efficient and effective task-agnostic distillation, which can be applied to students of any size. Specifically, we introduce two variants of ETD, ETD-Rand and ETD-Impt, which extract the teacher's parameters in a random manner and by following an importance metric respectively. In this way, the student has already acquired some knowledge at the beginning of the distillation process, which makes the distillation process converge faster. We demonstrate the effectiveness of ETD on the GLUE benchmark and SQuAD. The experimental results show that: (1) compared with the baseline without an ETD strategy, ETD can save 70\% of computation cost. Moreover, it achieves better results than the baseline when using the same computing resource. (2) ETD is generic and has been proven effective for different distillation methods (e.g., TinyBERT and MiniLM) and students of different sizes. The source code will be publicly available upon publication.

研究动机与目标

• 为解决无任务依赖BERT蒸馏带来的高计算成本问题，该问题通常需要在大规模无监督语料上进行大量训练。
• 实现对任意尺寸的学生模型（包括宽度和深度与教师模型不同的模型）的高效且有效的蒸馏。
• 开发一种通用的初始化策略，通过重用教师模型参数来加速收敛并提升性能。
• 证明基于剪枝的参数提取方式在重用参数方面优于随机初始化或简单层复制方法。

提出的方法

• ETD提出两种变体：ETD-Rand（随机神经元提取）和ETD-Impt（基于重要性的神经元提取），用于宽度方向上的参数重用。
• 该方法在教师模型的Transformer层中执行均匀的层选择，以确保广泛的知识迁移。
• 通过计算前馈网络神经元、注意力头和隐藏状态的重要性评分，指导选择性参数提取。
• 提取的参数在微调前用于初始化学生模型，从而减少对大规模蒸馏训练的需求。
• 该方法与现有的蒸馏框架（如TinyBERT和MiniLM）兼容。
• 提取过程高效，计算开销可忽略不计。

实验结果

研究问题

• RQ1 通过结构化提取重用预训练教师模型参数，能否提升无任务依赖BERT蒸馏的效率？
• RQ2 基于重要性的参数选择是否在初始化学生模型时优于随机或均匀选择？
• RQ3 ETD是否可应用于任意宽度和深度的学生模型，而不仅限于与教师模型相同的结构？
• RQ4 ETD在保持或提升模型性能的同时，能在多大程度上降低计算成本？
• RQ5 层选择策略（均匀、顶部、底部）的选择如何影响蒸馏的收敛速度和最终性能？

主要发现

• 与基线相比，ETD将蒸馏的计算成本降低了70%，且未牺牲性能。
• 在相同计算资源下，ETD在GLUE和SQuAD基准测试中均优于基线蒸馏方法。
• ETD-Impt（基于重要性评分选择参数）在所有实验中均持续优于ETD-Rand和随机初始化。
• 消融实验表明，即使是最不重要的神经元也具有正向贡献，而反转重要性排序会导致性能下降。
• 均匀层选择优于顶部或底部层选择，因为它能确保浅层到深层的均衡知识迁移。
• 该方法在不同蒸馏框架（包括TinyBERT和MiniLM）中均具有通用性和有效性，并适用于各种尺寸的学生模型。