arXiv:2510.15982cs.LGcs.AI2025-10被引 3

提出新型知识蒸馏框架AMiD,提升小模型训练稳定性和性能。

AMiD: Knowledge Distillation for LLMs with $α$-mixture Assistant Distribution

  • 引入可调节参数α的混合辅助分布,拓展蒸馏空间。
  • 在多个数据集上实现更优性能与更强训练稳定性。
  • 适合追求高效模型压缩与稳定训练的研究者使用。

自回归大语言模型在众多任务中表现卓越,但计算和内存开销巨大。知识蒸馏通过分布对齐将大模型知识迁移到小模型以缓解此问题。以往研究虽提出多种差异度量,但高维输出导致的零概率问题及容量差距仍带来训练不稳等根本挑战。近期部分工作尝试引入辅助分布,但缺乏对插值路径与发散关系的系统分析。本文提出α-混合辅助分布,构建了一类新的辅助分布家族,并设计统一框架AMiD,其通过引入可调参数α实现分布设计的连续扩展,同时基于最优性理论推广了适用的散度家族。大量实验表明,AMiD通过利用更广阔且理论完备的辅助分布空间,在多个基准上实现了更优性能与更强训练稳定性。代码已开源:https://github.com/aailab-kaist/AMiD。

原文摘要 · Abstract (English)

Autoregressive large language models (LLMs) have achieved remarkable improvement across many tasks but incur high computational and memory costs. Knowledge distillation (KD) mitigates this issue by transferring knowledge from a large teacher to a smaller student through distributional alignment. Previous studies have proposed various discrepancy metrics, but the capacity gap and training instability caused by near-zero probabilities, stemming from the high-dimensional output of LLMs, remain fundamental limitations. To overcome these challenges, several approaches implicitly or explicitly incorporating assistant distribution have recently been proposed. However, the past proposals of assistant distributions have been a fragmented approach without a systematic investigation of the interpolation path and the divergence. This paper proposes $α$-mixture assistant distribution, a novel generalized family of assistant distributions, and $α$-mixture distillation, coined AMiD, a unified framework for KD using the assistant distribution. The $α$-mixture assistant distribution provides a continuous extension of the assistant distribution by introducing a new distribution design variable $α$, which has been fixed in all previous approaches. Furthermore, AMiD generalizes the family of divergences used with the assistant distributions based on optimality, which has also been restricted in previous works. Through extensive experiments, we demonstrate that AMiD offers superior performance and training stability by leveraging a broader and theoretically grounded assistant distribution space. We release the code at https://github.com/aailab-kaist/AMiD.

知识蒸馏大模型压缩分布对齐

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。