[论文解读] Contrastive Distillation on Intermediate Representations for Language Model Compression

Siqi Sun, Zhe Gan+4 more|arXiv (Cornell University)|Sep 29, 2020

Topic Modeling参考文献 44被引用 5

一句话总结

本文提出 CoDIR，一种对比蒸馏框架，通过在中间表示上引入对比目标，提升压缩语言模型中的知识蒸馏效果，在 GLUE 基准上以更小、更快的模型超越现有方法。

ABSTRACT

Existing language model compression methods mostly use a simple L2 loss to distill knowledge in the intermediate representations of a large BERT model to a smaller one. Although widely used, this objective by design assumes that all the dimensions of hidden representations are independent, failing to capture important structural knowledge in the intermediate layers of the teacher network. To achieve better distillation efficacy, we propose Contrastive Distillation on Intermediate Representations (CoDIR), a principled knowledge distillation framework where the student is trained to distill knowledge through intermediate layers of the teacher via a contrastive objective. By learning to distinguish positive sample from a large set of negative samples, CoDIR facilitates the student's exploitation of rich information in teacher's hidden layers. CoDIR can be readily applied to compress large-scale language models in both pre-training and finetuning stages, and achieves superb performance on the GLUE benchmark, outperforming state-of-the-art compression methods.

研究动机与目标

• 为解决传统知识蒸馏依赖隐藏表示之间独立 $L_2$ 损失的问题，该方法无法捕捉 BERT 中间层的结构依赖关系。
• 探索对比学习在大规模语言模型知识蒸馏中的应用，特别是挖掘中间表示中的高阶关系。
• 设计适用于预训练和微调阶段的有效采样策略，其中标签不可用或稀疏。
• 通过将标准 $L_2$ 蒸馏替换为能区分同一样本（一致）与不同样本（不一致）表示对的对比目标，提升模型效率与性能。
• 证明平均池化表示在蒸馏中优于 [CLS] 标记，尤其在下游微调阶段表现更优。

提出的方法

• CoDIR 对教师与学生网络的中间隐藏表示应用对比损失，使学生能够区分来自同一输入的表示对（一致对）与来自不同输入的大量表示对（不一致对）。
• 对比目标被表述为对比学习损失，通过拉近同一样本的表示（正样本对）并推远不同样本的表示（负样本对）来实现。
• 内存库存储当前批次的所有表示，以在训练期间实现高效的负样本采样，尤其在预训练阶段（标签不可用）中尤为重要。
• 在微调阶段，使用标签不同的样本构建不一致对以最大化差异；在预训练阶段，则使用同一小批量内的随机配对作为代理。
• 该方法在预训练阶段结合了对比蒸馏、标准知识蒸馏（logits 上的 KL 散度）以及掩码语言建模（MLM）损失。
• 使用所有标记表示的平均池化作为蒸馏的句子级嵌入，替代 [CLS] 标记，实证结果表明该方法更有效。

实验结果

研究问题

• RQ1 与基于 $L_2$ 的标准蒸馏相比，对中间表示应用对比学习是否能提升压缩语言模型中的知识蒸馏效果？
• RQ2 在预训练和微调阶段，如何设计有效的负样本采样策略，尤其在预训练阶段标签不可用时？
• RQ3 使用平均池化表示替代 [CLS] 标记是否能提升压缩模型中的蒸馏性能？
• RQ4 负样本数量对语言模型压缩背景下对比蒸馏性能的影响如何？
• RQ5 所提出的 CoDIR 框架是否能在 GLUE 上实现最先进性能，同时相比现有压缩方法减小模型尺寸并降低推理时间？

主要发现

• CoDIR 在 GLUE 基准上超越最先进压缩方法，相比标准 MLM 预训练提升 1.9%，相比仅使用 KD 的预训练提升 1.0%。
• 学生模型大小仅为 BERT-base 的一半，推理速度提升 2 倍，且对 GPU 的需求更低。
• 使用平均池化表示作为蒸馏目标，GLUE 平均得分为 83.0，优于基于 [CLS] 标记的蒸馏（平均得分为 82.5）。
• 将负样本数量从 100 增加到 1000 在多数 GLUE 任务上提升了性能，其中 MNLI 和 QNLI 任务的增益最大。
• 对比损失对性能有显著贡献：在 KD 基础上加入 CRD 损失，相比标准 KD 提升平均 GLUE 得分 1.0%，相比原始 MLM 预训练提升 1.9%。
• 模型性能在较小数据集（如 CoLA、MRPC、RTE）上表现出更高方差，标准差约为 1.5，表明对随机种子初始化较为敏感。