[论文解读] Knowledge Distillation from Internal Representations

Gustavo Aguilar, Ling Yuan+4 more|arXiv (Cornell University)|Oct 8, 2019

Topic Modeling参考文献 20被引用 23

一句话总结

本文提出从 BERT 类模型的内部表征进行知识蒸馏，训练学生网络以匹配大型教师模型的软标签和隐藏层激活。通过在多个内部层使用注意力图和隐藏状态的 KL 散度进行蒸馏，该方法显著提升了泛化能力，并在 GLUE 基准数据集上优于标准知识蒸馏。

ABSTRACT

Knowledge distillation is typically conducted by training a small model (the student) to mimic a large and cumbersome model (the teacher). The idea is to compress the knowledge from the teacher by using its output probabilities as soft-labels to optimize the student. However, when the teacher is considerably large, there is no guarantee that the internal knowledge of the teacher will be transferred into the student; even if the student closely matches the soft-labels, its internal representations may be considerably different. This internal mismatch can undermine the generalization capabilities originally intended to be transferred from the teacher to the student. In this paper, we propose to distill the internal representations of a large model such as BERT into a simplified version of it. We formulate two ways to distill such representations and various algorithms to conduct the distillation. We experiment with datasets from the GLUE benchmark and consistently show that adding knowledge distillation from internal representations is a more powerful method than only using soft-label distillation.

研究动机与目标

• 为解决标准知识蒸馏的局限性，即仅传递输出概率，可能无法保留教师模型的内部语言知识。
• 通过不仅传递软标签，还传递来自大型教师模型的内部表征（如注意力图和隐藏状态），提升模型压缩效果。
• 通过使学生模型的内部行为与教师模型对齐，而非仅匹配其输出分布，确保学生模型的泛化能力与教师模型相似。
• 开发一种方法，使学生模型性能超越仅通过软标签蒸馏所能达到的水平，尤其在低资源或受限推理环境中。

提出的方法

• 该方法不仅在最终输出概率上进行知识蒸馏，还在教师模型多个层的中间隐藏表征上进行蒸馏。
• 使用 KL 散度匹配教师模型与学生模型在每一层的注意力概率，以捕捉高层次的语言抽象特征。
• 在选定层上，对教师模型与学生模型的隐藏状态应用独立的 KL 散度损失，实现内部表征的结构对齐。
• 该方法支持渐进式和堆叠式知识蒸馏策略，可灵活地将多个教师层压缩为单个学生层。
• 学生模型通过交叉熵损失（在硬标签上）、软标签上的蒸馏损失，以及注意力图和隐藏状态上的额外蒸馏损失联合训练。
• 与先前方法（如 FitNet）不同，该方法避免引入用于表征对齐的额外参数，而是直接在相同空间中匹配表征。

实验结果

研究问题

• RQ1 从大型教师模型蒸馏内部表征（如注意力图和隐藏状态）是否能提升小型学生模型的泛化能力和性能，超越标准知识蒸馏？
• RQ2 通过对齐学生模型的内部行为与教师模型（而不仅输出概率），是否能带来更稳健和准确的预测，尤其是在模糊或具有挑战性的样本上？
• RQ3 在错误模式和与教师预测的一致性方面，内部表征蒸馏与标准 KD 相比如何？
• RQ4 该方法能否在不引入额外参数或增加推理复杂度的前提下，有效压缩深层 Transformer 层中编码的语言知识？
• RQ5 在不同层（如早期 vs. 晚期）进行蒸馏对最终学生模型性能和泛化能力有何影响？

主要发现

• 使用内部表征蒸馏训练的学生模型在 QQP 开发集上达到了教师模型 97.9% 的正确预测率（36,191 / 36,967），而标准 KD 的学生模型为 95.7%（35,401）。
• 内部蒸馏的学生模型犯了 776 次错误，而标准 KD 学生模型犯了 1,566 次错误，表明其与教师模型的决策边界对齐更好。
• 在 QQP 数据集上，内部蒸馏的学生模型正确识别了教师模型正确而标准 KD 学生模型失败的所有情况，例如表 4 中的样本 3 和 4。
• 当教师模型正确时，该方法使错误预测数量比标准 KD 减少 50% 以上，表明其具有更优的一致性和泛化能力。
• 在教师模型犯错的情况下（如表 4 中的样本 1），内部蒸馏的学生模型表现更谨慎（概率为 0.4221，而标准 KD 为 0.9999），表明其减少了过度自信，校准性更好。
• 该方法在所有四个 GLUE 基准数据集上均持续优于标准 KD，展示了在准确率和与教师行为对齐方面的稳健提升。