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

TL;DR

This work tackles the large computational burden of autoregressive LLMs by proposing a unified knowledge distillation framework anchored in the α-mixture assistant distribution. By introducing the extra design parameter $\alpha$ over the traditional interpolation, AMiD generalizes both $m$-mixture and $e$-mixture approaches and supports arbitrary divergences, achieving better mode coverage or mode seeking as needed. The authors prove the optimality of AMiD under perfect optimization and provide gradient analyses illustrating how $\alpha$ shapes learning dynamics. Extensive experiments on instruction-following and task-specific distillation show AMiD consistently outperforms prior methods, improving stability and generalization across SGOs and model scales. Overall, AMiD offers a flexible, theoretically grounded path to more efficient and robust KD for large language models.

Abstract

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.

Figures (5)

• Figure 1: Visualization of the $\alpha$-mixture assistant distribution family. (a) The $\alpha$-mixture assistant distribution provides a generalized framework for assistant distributions, with prior studies agarwal2024policyko2024distillmshing2025taid recoverable as special cases. (b-f) Illustration of the $\alpha$-mixture assistant distribution where $p=\mathcal{N}(0, {0.5}^2)$, $q_\theta=\mathcal{N}(3, 1^2)$, and $\lambda=0.3$.
• Figure 2: Visualization of the characteristics of the $\alpha$-mixture distillation, AMiD. (a) $\alpha$ determines the geometry of interpolation while $\lambda$ controls the portion between $p$ and $q_\theta$. (b) $r_\theta^{(\alpha,\lambda)}$ can be interpreted as the internal division point in terms of $\alpha$-divergence. Due to the uniqueness, updates of $r_\theta^{(\alpha,\lambda)}$ also affect $q_\theta$. (c) Toy experiment with two-modal $p$ and uni-modal $q_\theta$. $\alpha$ controls the property of $q_\theta^*$ between the mode-covering and mode-seeking, even though we minimize $D_{\mathrm{KL}}(p\|r_\theta^{(\alpha,\lambda)})$.
• Figure 3: ROUGE-L curve on Dolly dataset.
• Figure 4: Performance curve on ROUGE-L (quality) and Self-BLEU (diversity). Colored dashed lines: no-assistant baseline. $\lambda=0.1$.
• Figure 5: Sensitivity analysis for $\lambda$ under fixed $\alpha=-5.0$.

Theorems & Definitions (14)

• Proposition 3.0
• Definition 1: $\alpha$-mixture assistant distribution
• Theorem 3.1
• Proposition 3.1
• Theorem 3.2
• Proposition 3.2
• Proposition A.0
• proof
• Proposition A.0
• proof