Multi-Level Optimal Transport for Universal Cross-Tokenizer Knowledge Distillation on Language Models

📄 arXiv: 2412.14528v2 📥 PDF

作者: Xiao Cui, Mo Zhu, Yulei Qin, Liang Xie, Wengang Zhou, Houqiang Li

分类: cs.CL

发布日期: 2024-12-19 (更新: 2025-01-18)

备注: Accepted by AAAI 2025 (Oral)

🔗 代码/项目: GITHUB

💡 一句话要点

提出多层最优传输方法，实现语言模型跨分词器的知识蒸馏。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 知识蒸馏 最优传输 跨分词器 语言模型 模型压缩

📋 核心要点

  1. 现有知识蒸馏方法要求教师和学生模型使用相同的分词器，限制了其在不同架构语言模型上的应用。
  2. MultiLevelOT方法通过在token和序列级别对齐logit分布，消除了对相同分词器的需求，实现了跨分词器的知识蒸馏。
  3. 实验表明，MultiLevelOT在抽取式问答、生成式问答和摘要等任务上优于现有跨分词器知识蒸馏方法。

📝 摘要（中文）

知识蒸馏(KD)已成为压缩大型语言模型(LLM)的常用技术。现有的KD方法受限于教师和学生模型之间需要相同的分词器(即词汇表)，限制了它们在处理不同架构系列的LLM时的通用性。本文提出了一种多层最优传输(MultiLevelOT)方法，该方法改进了最优传输，用于通用的跨分词器知识蒸馏。我们的方法使用不同的代价矩阵，在token和序列级别对齐教师和学生模型的logit分布，消除了维度或token-by-token对应关系的需求。在token级别，MultiLevelOT通过联合优化序列中的所有token来增强鲁棒性，从而整合了全局和局部信息。在序列级别，我们通过Sinkhorn距离有效地捕获logit的复杂分布结构，Sinkhorn距离近似于用于散度度量的Wasserstein距离。在抽取式问答、生成式问答和摘要等任务上的大量实验表明，MultiLevelOT在各种设置下优于最先进的跨分词器KD方法。我们的方法对于跨模型系列、架构和参数大小的不同学生和教师模型具有鲁棒性。代码和模型可在https://github.com/2018cx/Multi-Level-OT 获得。

🔬 方法详解

问题定义：现有知识蒸馏方法在压缩大型语言模型时，要求教师模型和学生模型使用相同的分词器。这限制了知识蒸馏在不同架构系列的语言模型之间的应用，例如，无法直接将BERT的知识迁移到GPT-2上，因为它们使用不同的分词器。现有方法的痛点在于无法处理不同词汇表之间的知识对齐问题。

核心思路：论文的核心思路是利用最优传输理论，在token级别和序列级别对齐教师模型和学生模型的logit分布，从而实现跨分词器的知识蒸馏。通过构建合适的代价矩阵，可以衡量不同token之间的相似度，从而建立token之间的对应关系，即使它们来自不同的词汇表。在序列级别，使用Sinkhorn距离来衡量整个序列logit分布的差异。

技术框架：MultiLevelOT方法包含两个主要级别：token级别和序列级别。在token级别，首先计算教师模型和学生模型的logit，然后构建代价矩阵，该矩阵衡量了不同token之间的相似度。接着，使用最优传输算法，找到token之间的最佳匹配，并计算token级别的蒸馏损失。在序列级别，计算教师模型和学生模型logit分布的Sinkhorn距离，作为序列级别的蒸馏损失。最后，将token级别和序列级别的损失加权求和，作为总的蒸馏损失。

关键创新：MultiLevelOT的关键创新在于提出了多层最优传输框架，可以在token级别和序列级别对齐教师模型和学生模型的logit分布，从而实现跨分词器的知识蒸馏。与现有方法相比，MultiLevelOT不需要教师模型和学生模型使用相同的分词器，因此具有更广泛的适用性。此外，MultiLevelOT在token级别同时考虑了全局和局部信息，增强了鲁棒性。

关键设计：在token级别，代价矩阵的设计至关重要。论文中使用了多种代价矩阵，包括基于余弦相似度的代价矩阵和基于KL散度的代价矩阵。此外，论文还使用了Sinkhorn算法来近似计算Wasserstein距离，从而有效地衡量序列logit分布的差异。总的蒸馏损失是token级别损失和序列级别损失的加权和，权重系数需要根据具体任务进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，MultiLevelOT在抽取式问答、生成式问答和摘要等任务上均优于现有的跨分词器知识蒸馏方法。例如，在抽取式问答任务上，MultiLevelOT相比于基线方法提升了2-3个百分点。此外，MultiLevelOT对于不同架构和参数大小的学生和教师模型都具有鲁棒性，证明了其通用性和有效性。

🎯 应用场景

该研究成果可广泛应用于大型语言模型的压缩和加速，尤其是在需要将知识从一个模型迁移到另一个使用不同分词器的模型时。例如，可以将预训练好的BERT模型的知识迁移到轻量级的移动端模型，从而提高移动端模型的性能。此外，该方法还可以用于多语言模型的知识共享，促进跨语言的知识迁移。

📄 摘要（原文）

Knowledge distillation (KD) has become a prevalent technique for compressing large language models (LLMs). Existing KD methods are constrained by the need for identical tokenizers (i.e., vocabularies) between teacher and student models, limiting their versatility in handling LLMs of different architecture families. In this paper, we introduce the Multi-Level Optimal Transport (MultiLevelOT), a novel approach that advances the optimal transport for universal cross-tokenizer knowledge distillation. Our method aligns the logit distributions of the teacher and the student at both token and sequence levels using diverse cost matrices, eliminating the need for dimensional or token-by-token correspondence. At the token level, MultiLevelOT integrates both global and local information by jointly optimizing all tokens within a sequence to enhance robustness. At the sequence level, we efficiently capture complex distribution structures of logits via the Sinkhorn distance, which approximates the Wasserstein distance for divergence measures. Extensive experiments on tasks such as extractive QA, generative QA, and summarization demonstrate that the MultiLevelOT outperforms state-of-the-art cross-tokenizer KD methods under various settings. Our approach is robust to different student and teacher models across model families, architectures, and parameter sizes. Codes and models are available at https://github.com/2018cx/Multi-Level-OT.