Co-training and Co-distillation for Quality Improvement and Compression of Language Models

📄 arXiv: 2311.02849v2 📥 PDF

作者: Hayeon Lee, Rui Hou, Jongpil Kim, Davis Liang, Hongbo Zhang, Sung Ju Hwang, Alexander Min

分类: cs.CL, cs.AI

发布日期: 2023-11-06 (更新: 2023-11-07)

备注: Findings of EMNLP 2023

💡 一句话要点

提出协同训练与协同蒸馏以提升语言模型质量与压缩效率

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 知识蒸馏 语言模型 协同训练 模型压缩 性能提升 自然语言处理 深度学习

📋 核心要点

  1. 现有的知识蒸馏方法在压缩大型语言模型时，通常导致小模型性能低于原模型，无法兼顾性能与推理速度。
  2. 本文提出的CTCD框架通过协同训练和相互蒸馏知识，旨在同时提升小模型和大模型的性能。
  3. 实验结果显示，CTCD框架下的小模型在GLUE基准测试中超越了原大型模型，提升幅度达到1.66，验证了其有效性。

📝 摘要（中文）

知识蒸馏（KD）通过将大型预训练语言模型（PLMs）的知识转移到较小模型中，从而压缩计算开销，适用于资源受限或实时场景。然而，大多数小模型的性能未能超越原始大型模型，导致在提升推理速度时牺牲了性能。为了解决这一问题，本文提出了协同训练与协同蒸馏（CTCD）框架，通过共同训练两个模型并相互蒸馏知识，来同时提升性能和推理速度。CTCD框架基于两个重要发现：1）在协同训练过程中，从小模型向大模型蒸馏知识可以提升大模型的性能；2）大模型性能的提升进一步促进了小模型的表现。CTCD框架具有与现有技术（如架构设计或数据增强）结合的潜力，能够替代单向KD方法，实现更进一步的性能提升。大量消融实验表明，CTCD的有效性，且通过CTCD蒸馏的小模型在GLUE基准测试中超越原始大型模型1.66的显著幅度。

🔬 方法详解

问题定义：本文旨在解决现有知识蒸馏方法在压缩大型语言模型时，导致小模型性能不足的问题。现有方法往往牺牲性能以提高推理速度，无法实现两者的平衡。

核心思路：CTCD框架的核心思路是通过协同训练两个模型，并在此过程中相互蒸馏知识，从而实现小模型和大模型的共同提升。这样的设计使得小模型的知识可以反过来增强大模型的性能，形成良性循环。

技术框架：CTCD框架包括两个主要阶段：协同训练和知识蒸馏。在协同训练阶段，两个模型同时进行训练，互相影响；在知识蒸馏阶段，小模型向大模型传递知识，同时大模型也为小模型提供反馈。

关键创新：CTCD框架的创新在于其双向蒸馏机制，区别于传统的单向蒸馏方法。通过这种双向蒸馏，两个模型能够互相促进，提升整体性能。

关键设计：在CTCD中，采用了特定的损失函数来平衡两个模型的训练目标，并设计了适应性参数设置，以确保在协同训练过程中，两个模型能够有效地互相学习。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

CTCD框架在GLUE基准测试中表现出色，蒸馏后的小模型性能超越原大型模型1.66，显示出显著的性能提升。这一结果证明了CTCD在语言模型压缩与性能提升方面的有效性。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和实时翻译等场景，尤其是在资源受限的设备上，如移动端和嵌入式系统。CTCD框架的有效性将推动小型语言模型在实际应用中的广泛使用，提升用户体验。

📄 摘要（原文）

Knowledge Distillation (KD) compresses computationally expensive pre-trained language models (PLMs) by transferring their knowledge to smaller models, allowing their use in resource-constrained or real-time settings. However, most smaller models fail to surpass the performance of the original larger model, resulting in sacrificing performance to improve inference speed. To address this issue, we propose Co-Training and Co-Distillation (CTCD), a novel framework that improves performance and inference speed together by co-training two models while mutually distilling knowledge. The CTCD framework successfully achieves this based on two significant findings: 1) Distilling knowledge from the smaller model to the larger model during co-training improves the performance of the larger model. 2) The enhanced performance of the larger model further boosts the performance of the smaller model. The CTCD framework shows promise as it can be combined with existing techniques like architecture design or data augmentation, replacing one-way KD methods, to achieve further performance improvement. Extensive ablation studies demonstrate the effectiveness of CTCD, and the small model distilled by CTCD outperforms the original larger model by a significant margin of 1.66 on the GLUE benchmark.