Sinkhorn Distance Minimization for Knowledge Distillation
作者: Xiao Cui, Yulei Qin, Yuting Gao, Enwei Zhang, Zihan Xu, Tong Wu, Ke Li, Xing Sun, Wengang Zhou, Houqiang Li
分类: cs.LG, cs.CL
发布日期: 2024-02-27
备注: Accepted by COLING 2024
💡 一句话要点
提出Sinkhorn知识蒸馏以解决现有蒸馏方法的不足
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 知识蒸馏 Sinkhorn距离 自然语言处理 模型压缩 高维空间 散度度量 深度学习
📋 核心要点
- 现有知识蒸馏方法在教师与学生分布重叠较少时，无法有效提供监督，导致性能下降。
- 提出Sinkhorn知识蒸馏（SinKD），利用Sinkhorn距离精确评估教师与学生分布的差异，采用批量重构方法。
- 在GLUE和SuperGLUE上进行的实验表明，SinKD在多种LLMs架构上均优于现有最先进的方法，表现出更好的可比性和有效性。
📝 摘要（中文）
知识蒸馏（KD）已被广泛应用于压缩大型语言模型（LLMs）。现有的KD方法探讨了多种散度度量，包括Kullback-Leibler（KL）、反Kullback-Leibler（RKL）和Jensen-Shannon（JS）散度。然而，由于这些度量的假设和定义的局限性，当教师和学生之间的分布重叠较少时，它们无法提供有效的监督。本文展示了KL、RKL和JS散度分别存在模式平均、模式崩溃和模式低估的问题，导致基于logits的KD在多样化的NLP任务中效果不佳。我们提出了Sinkhorn知识蒸馏（SinKD），利用Sinkhorn距离确保教师和学生分布之间差异的细致和准确评估。此外，借助Sinkhorn度量的特性，我们摆脱了限制每对教师-学生样本感知散度的逐样本KD，提出了一种批量重构方法，以捕捉高维空间中分布的几何复杂性。在GLUE和SuperGLUE上的全面评估显示，我们在各种编码器、编码器-解码器和解码器架构的LLMs上优于现有最先进的方法。
🔬 方法详解
问题定义：本文旨在解决现有知识蒸馏方法在教师与学生分布重叠较少时的监督不足问题，现有的KL、RKL和JS散度存在模式平均、崩溃和低估等缺陷，影响了蒸馏效果。
核心思路：提出Sinkhorn知识蒸馏（SinKD），利用Sinkhorn距离的特性，提供对教师和学生分布之间差异的细致评估，避免逐样本的限制，采用批量处理来捕捉分布的几何特性。
技术框架：SinKD的整体架构包括数据预处理、Sinkhorn距离计算、批量重构和模型训练四个主要模块。首先对教师和学生模型的输出进行处理，然后计算Sinkhorn距离，最后通过批量重构优化模型。
关键创新：SinKD的核心创新在于引入Sinkhorn距离作为新的散度度量，克服了传统散度在处理分布差异时的不足，特别是在样本重叠较少的情况下。
关键设计：在损失函数中引入Sinkhorn距离，并设计了适应高维空间的批量重构策略，确保了模型在训练过程中能够有效捕捉分布的几何特性。
🖼️ 关键图片
📊 实验亮点
在GLUE和SuperGLUE的实验中，SinKD在各类大型语言模型上均表现出显著的性能提升，相较于最先进的方法，整体性能提升幅度达到X%（具体数据待补充），验证了其有效性和优越性。
🎯 应用场景
该研究的潜在应用领域包括自然语言处理、机器翻译和对话系统等。通过改进知识蒸馏技术，SinKD能够有效提升大型语言模型的性能，具有广泛的实际价值和未来影响，尤其是在资源受限的环境中。
📄 摘要（原文）
Knowledge distillation (KD) has been widely adopted to compress large language models (LLMs). Existing KD methods investigate various divergence measures including the Kullback-Leibler (KL), reverse Kullback-Leibler (RKL), and Jensen-Shannon (JS) divergences. However, due to limitations inherent in their assumptions and definitions, these measures fail to deliver effective supervision when few distribution overlap exists between the teacher and the student. In this paper, we show that the aforementioned KL, RKL, and JS divergences respectively suffer from issues of mode-averaging, mode-collapsing, and mode-underestimation, which deteriorates logits-based KD for diverse NLP tasks. We propose the Sinkhorn Knowledge Distillation (SinKD) that exploits the Sinkhorn distance to ensure a nuanced and precise assessment of the disparity between teacher and student distributions. Besides, profit by properties of the Sinkhorn metric, we can get rid of sample-wise KD that restricts the perception of divergence in each teacher-student sample pair. Instead, we propose a batch-wise reformulation to capture geometric intricacies of distributions across samples in the high-dimensional space. Comprehensive evaluation on GLUE and SuperGLUE, in terms of comparability, validity, and generalizability, highlights our superiority over state-of-the-art methods on all kinds of LLMs with encoder-only, encoder-decoder, and decoder-only architectures.