[论文解读] LRC-BERT: Latent-representation Contrastive Knowledge Distillation for Natural Language Understanding

Hao Fu, Shaojun Zhou+5 more|arXiv (Cornell University)|Dec 14, 2020

Topic Modeling参考文献 40被引用 5

一句话总结

LRC-BERT 提出了一种新颖的知识蒸馏框架，该框架基于角度距离的对比学习（COS-NCE）来传递 BERT 教师模型的中间层表征，从而在值匹配之外提升知识蒸馏效果。它通过梯度扰动进一步增强了鲁棒性，并采用两阶段训练策略，在 GLUE 上实现了最先进性能，同时提升了泛化能力和稳定性。

ABSTRACT

The pre-training models such as BERT have achieved great results in various natural language processing problems. However, a large number of parameters need significant amounts of memory and the consumption of inference time, which makes it difficult to deploy them on edge devices. In this work, we propose a knowledge distillation method LRC-BERT based on contrastive learning to fit the output of the intermediate layer from the angular distance aspect, which is not considered by the existing distillation methods. Furthermore, we introduce a gradient perturbation-based training architecture in the training phase to increase the robustness of LRC-BERT, which is the first attempt in knowledge distillation. Additionally, in order to better capture the distribution characteristics of the intermediate layer, we design a two-stage training method for the total distillation loss. Finally, by verifying 8 datasets on the General Language Understanding Evaluation (GLUE) benchmark, the performance of the proposed LRC-BERT exceeds the existing state-of-the-art methods, which proves the effectiveness of our method.

研究动机与目标

• 为解决现有 NLP 知识蒸馏方法的局限性，这些方法仅关注输出值匹配，而忽略了中间层表征的结构分布特性。
• 通过引入基于梯度扰动的架构，首次在知识蒸馏中提升训练过程的鲁棒性。
• 通过设计一种两阶段训练策略，优先在早期训练中学习分布结构，从而增强教师到学生模型的知识迁移。
• 通过解耦学生模型结构与教师层约束，实现更灵活的学生模型架构设计，提升压缩灵活性。

提出的方法

• 提出一种新颖的对比损失 COS-NCE，基于基于余弦的噪声对比估计，以捕捉中间层表征之间的角度相似性，建模其结构分布特征。
• 采用两阶段训练策略，总蒸馏损失分阶段优化，初期专注于中间层分布学习，随后再引入预测头损失。
• 在训练过程中对词嵌入层应用梯度扰动，利用梯度值对隐藏状态进行扰动，以提升模型鲁棒性。
• 结合多种损失组件：KL 散度用于软标签，交叉熵用于硬标签，COS-NCE 用于中间层蒸馏，采用自适应加权策略。
• 使用维度变换矩阵，实现灵活的学生网络设计，解耦学生架构与教师模型逐层对齐的约束。
• 采用参数量少于教师模型的学生模型，实现在边缘设备上的高效部署，同时保持高性能。

实验结果

研究问题

• RQ1 基于角度距离的对比学习是否能在 NLP 知识蒸馏中超越中间表征值匹配，实现更优的知识迁移？
• RQ2 在训练过程中引入梯度扰动能否提升知识蒸馏中学生模型的鲁棒性？
• RQ3 两阶段训练策略是否能通过优先学习早期训练中的分布结构，提升蒸馏的有效性？
• RQ4 在标准基准上，该方法与现有蒸馏技术相比，在准确率、鲁棒性和压缩效率方面表现如何？

主要发现

• LRC-BERT 在 GLUE 基准的 8 个数据集上均达到最先进性能，优于现有蒸馏方法。
• 消融实验表明，若移除 COS-NCE 损失，MNLI-m 验证集上的准确率下降 4.0%，证实其在知识迁移中的关键作用。
• 两阶段训练方法相比非分阶段变体性能提升 4.0%，表明早期聚焦分布结构学习可显著增强蒸馏效果。
• 梯度扰动使第二阶段训练损失更稳定，减少波动，提升模型收敛性与鲁棒性。
• COS-NCE 损失能有效捕捉角度相似性特征，即使在 MSE 基线因嵌入差异过大而失效时，仍能实现正确预测。
• 在 CoLA 任务上，模型相比基线蒸馏方法提升 13%，凸显基于角度的蒸馏具有显著优势。