Dual-Space Knowledge Distillation with Key-Query Matching for Large Language Models with Vocabulary Mismatch

📄 arXiv: 2603.22056v1 📥 PDF

作者: Stella Eva Tsiapali, Cong-Thanh Do, Kate Knill

分类: cs.CL

发布日期: 2026-03-23

备注: Accepted at ICASSP 2026


💡 一句话要点

提出DSKD-CMA-GA,通过生成对抗学习解决LLM蒸馏中词表不匹配问题。

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

关键词: 知识蒸馏 大型语言模型 词表不匹配 生成对抗学习 跨模型注意力 文本生成 模型压缩

📋 核心要点

  1. 大型语言模型部署成本高昂,知识蒸馏是降低成本的有效方法,但跨词表蒸馏仍具挑战。
  2. 论文提出DSKD-CMA-GA,利用生成对抗学习,对齐不同模型生成的键和查询分布。
  3. 实验表明,该方法在文本生成质量上取得了一致的ROUGE-L增益,尤其是在分布外数据上。

📝 摘要(中文)

大型语言模型(LLM)在各种语言任务中取得了最先进的(SOTA)性能,但由于其规模和资源需求,部署成本很高。知识蒸馏(KD)通过训练较小的学生模型来模仿较大的教师模型来解决这个问题,从而在不显著降低性能的情况下提高效率。双空间知识蒸馏与跨模型注意力(DSKD-CMA)已经成为LLM之间进行KD的一种SOTA方法,但其内部工作原理在很大程度上仍然不透明。在这项工作中,我们通过手动token对齐探测和热图可视化,系统地分析了DSKD-CMA的注意力机制,揭示了其优点和局限性。在此基础上,我们引入了一种基于生成对抗(GA)学习的新方法DSKD-CMA-GA,以解决由不同模型计算出的键和查询之间的不匹配分布问题。实验表明,文本生成质量方面有适度但一致的ROUGE-L增益,特别是在分布外数据上(平均+0.37),缩小了跨tokenizer和同tokenizer KD之间的差距。

🔬 方法详解

问题定义:论文旨在解决大型语言模型知识蒸馏中,当教师模型和学生模型使用不同tokenizer时,由于词表不匹配导致的性能下降问题。现有方法,如DSKD-CMA,在跨模型注意力机制上存在局限性,键(Key)和查询(Query)的分布存在差异,影响蒸馏效果。

核心思路:论文的核心思路是通过生成对抗学习(Generative Adversarial Learning)来对齐教师模型和学生模型生成的键(Key)和查询(Query)的分布。通过对抗训练,使得学生模型生成的键和查询更接近教师模型的分布,从而提高知识迁移的效率。

技术框架:DSKD-CMA-GA方法建立在DSKD-CMA的基础上,主要包括以下几个阶段:1) 使用教师模型和学生模型分别生成文本表示;2) 通过跨模型注意力机制计算注意力权重;3) 使用生成对抗网络(GAN)对齐键和查询的分布;4) 使用知识蒸馏损失函数训练学生模型。

关键创新:该方法最重要的技术创新点是引入了生成对抗学习来解决跨词表知识蒸馏中的键和查询分布不匹配问题。与现有方法相比,DSKD-CMA-GA能够更有效地对齐不同模型生成的表示,从而提高知识迁移的质量。

关键设计:DSKD-CMA-GA的关键设计包括:1) 使用生成器(Generator)将学生模型生成的键和查询映射到教师模型的分布;2) 使用判别器(Discriminator)区分学生模型和教师模型生成的键和查询;3) 使用对抗损失(Adversarial Loss)训练生成器,使其生成的键和查询更接近教师模型的分布;4) 使用知识蒸馏损失(Knowledge Distillation Loss)训练学生模型,使其模仿教师模型的行为。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,DSKD-CMA-GA在文本生成质量方面取得了适度但一致的ROUGE-L增益,尤其是在分布外数据上,平均提升了0.37。该方法缩小了跨tokenizer和同tokenizer KD之间的差距,表明其在解决词表不匹配问题上的有效性。这些结果验证了生成对抗学习在知识蒸馏中的应用潜力。

🎯 应用场景

该研究成果可应用于各种需要模型压缩和加速的场景,例如移动设备上的自然语言处理、边缘计算和低资源语言的机器翻译。通过知识蒸馏,可以将大型语言模型的知识迁移到小型模型,从而在资源受限的环境中实现高性能的语言处理能力。此外,该方法对于解决多语言模型中的词表对齐问题也具有潜在价值。

📄 摘要(原文)

Large language models (LLMs) achieve state-of-the-art (SOTA) performance across language tasks, but are costly to deploy due to their size and resource demands. Knowledge Distillation (KD) addresses this by training smaller Student models to mimic larger Teacher models, improving efficiency without significant performance loss. Dual-Space Knowledge Distillation with Cross-Model Attention (DSKD-CMA) has emerged as a SOTA method for KD between LLMs with distinct tokenizers, yet its internal workings remain largely opaque. In this work, we systematically analyse the attention mechanism of DSKD-CMA through manual token alignment probing and heatmap visualisations, revealing both strengths and limitations. Building on this, we introduce a novel method, DSKD-CMA-GA, based on Generative Adversarial (GA) learning, to address the mismatched distributions between the keys and queries computed from distinct models. Experiments show modest but consistent ROUGE-L gains in text generation quality, particularly on out-of-distribution data (+0.37 on average), narrowing the gap between cross- and same-tokenizer KD.