[论文解读] ALP-KD: Attention-Based Layer Projection for Knowledge Distillation

Peyman Passban, Yimeng Wu+2 more|arXiv (Cornell University)|Dec 27, 2020

Topic Modeling参考文献 29被引用 8

一句话总结

本文提出 ALP-KD，一种基于注意力的层投影知识蒸馏方法，通过可学习的注意力权重组合多个教师网络层，避免了跳过层的问题。该方法使学生模型能够模仿最终输出和中间表征，从而在将 BERT 模型蒸馏为更小模型时，在 GLUE 基准测试中表现优于现有方法。

ABSTRACT

Knowledge distillation is considered as a training and compression strategy in which two neural networks, namely a teacher and a student, are coupled together during training. The teacher network is supposed to be a trustworthy predictor and the student tries to mimic its predictions. Usually, a student with a lighter architecture is selected so we can achieve compression and yet deliver high-quality results. In such a setting, distillation only happens for final predictions whereas the student could also benefit from teacher's supervision for internal components. Motivated by this, we studied the problem of distillation for intermediate layers. Since there might not be a one-to-one alignment between student and teacher layers, existing techniques skip some teacher layers and only distill from a subset of them. This shortcoming directly impacts quality, so we instead propose a combinatorial technique which relies on attention. Our model fuses teacher-side information and takes each layer's significance into consideration, then performs distillation between combined teacher layers and those of the student. Using our technique, we distilled a 12-layer BERT (Devlin et al. 2019) into 6-, 4-, and 2-layer counterparts and evaluated them on GLUE tasks (Wang et al. 2018). Experimental results show that our combinatorial approach is able to outperform other existing techniques.

研究动机与目标

• 解决现有知识蒸馏方法因架构不匹配而跳过中间教师层的局限性。
• 通过从多个非对应教师层获取监督，而非依赖一对一的层匹配，提升学生模型性能。
• 开发一种组合式、基于注意力的机制，动态加权并融合相关教师层表征用于蒸馏。
• 证明通过注意力组合教师层可实现学生与教师模型之间更优的内部表征对齐。
• 在将深层教师网络（如 12 层 BERT）蒸馏为紧凑学生模型（6 层、4 层和 2 层变体）时，实现最先进的性能。

提出的方法

• 提出一种新颖的基于注意力的层投影（ALP-KD）机制，将多个教师层特征融合为单一加权组合，用于蒸馏。
• 使用可学习的注意力机制，根据各教师层与学生当前层的相关性，分配重要性得分。
• 将组合后的教师表征形式化为多个教师层中间特征的加权和，其中权重通过多头注意力模块计算。
• 通过蒸馏损失将融合后的教师表征集成到学生的训练过程中，比较学生中间输出与教师特征的加权组合。
• 通过支持非相邻或非匹配层的注意力机制，实现灵活配置，消除对逐层对齐的需求。
• 使用组合损失训练学生模型：包括硬标签的交叉熵损失，以及学生输出与注意力融合后的教师特征之间的蒸馏损失。

实验结果

研究问题

• RQ1 通过注意力机制组合多个非对应教师层，能否在知识蒸馏中超越一对一层匹配的性能？
• RQ2 教师表征的注意力融合如何影响学生与教师模型之间中间隐藏状态的对齐？
• RQ3 通过避免跳过问题（即不丢弃任何教师层），是否能带来学生模型性能的可测量提升？
• RQ4 在 NLP 基准测试中，ALP-KD 相较于 PKD 和 RKD 等现有方法，在准确率和表征保真度方面提升程度如何？
• RQ5 ALP-KD 是否能有效应用于将深层模型（如 12 层 BERT）蒸馏为显著更小的架构（如 2 层 BERT），且不造成性能下降？

主要发现

• 在将 12 层 BERT 蒸馏为 6 层、4 层和 2 层学生模型时，ALP-KD 在多个 GLUE 基准任务上优于现有知识蒸馏方法，包括 PKD 和 RKD。
• 注意力机制融合教师层可显著拉近学生与教师模型之间的中间表征距离，如余弦距离可视化所示：ALP-KD 的平均距离为 0.05，而 PKD 为 0.20。
• 通过主成分分析（PCA）的隐藏状态可视化证实，ALP-KD 学生模型学习到的表征与教师模型更为相似，优于基线方法。
• 该方法通过利用所有教师层的信息，有效缓解了跳过问题，避免因架构不匹配而丢弃层。
• 即使在从极深教师网络蒸馏到极窄学生网络时，ALP-KD 仍保持强劲性能，展现出对架构差异的鲁棒性。
• 消融实验确认，注意力机制组合至关重要——简单拼接或均匀加权无法获得具有竞争力的结果。