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

会议: ICLR 2026
arXiv: 2510.15982
代码: https://github.com/aailab-kaist/AMiD
领域: 模型压缩 / 知识蒸馏
关键词: 知识蒸馏, 辅助分布, α-混合, f-散度, LLM压缩

一句话总结¶

提出α-mixture assistant distribution及统一蒸馏框架AMiD，通过引入新设计变量α（控制教师-学生分布插值路径的几何形状）泛化了现有辅助分布方法（m-mixture和e-mixture为α=±1的特例），并证明了在任意散度和α下的最优性保证，在多个LLM蒸馏基准上取得SOTA性能。

研究背景与动机¶

领域现状：LLM知识蒸馏通过对齐教师-学生的token级分布来压缩模型。近期研究引入了"辅助分布"——教师和学生分布的混合体——来缓解容量差距和近零概率导致的训练不稳定性。
现有痛点：(a) 现有辅助分布方法（GKD/DistiLLM用算术平均即m-mixture，TAID用几何平均即e-mixture）各自独立提出，缺乏统一框架；(b) 辅助分布的设计与散度的选择耦合在一起，搜索空间被人为限制；(c) α（控制插值路径几何）被固定为±1，未被探索。
核心矛盾：LLM的高维输出空间中大量概率接近零，导致密度比估计不稳定；同时教师-学生容量差距使直接对齐困难。辅助分布是解决这两个问题的关键，但现有设计不够通用。
本文要解决什么：建立辅助分布和散度的统一理论框架，发现新的、更好的辅助分布形式。
切入角度：用信息几何中的广义 \(f_\alpha\)-均值统一现有辅助分布——m-mixture和e-mixture分别对应算术均值(α=-1)和几何均值(α=1)，而α可以取任意实数值。
核心idea一句话：用广义 \(f_\alpha\)-均值将辅助分布从两种离散选择扩展为一族连续参数化的分布，并证明任意α和散度下的蒸馏最优性。

方法详解¶

整体框架¶

AMiD引入两个控制变量：α（控制插值路径几何）和λ（控制插值位置）。辅助分布定义为：\(\tilde{r}_\theta^{(\alpha,\lambda)}(z) = (\lambda p(z)^{\frac{1-\alpha}{2}} + (1-\lambda) q_\theta(z)^{\frac{1-\alpha}{2}})^{\frac{2}{1-\alpha}}\)（α≠1时），归一化后得到有效概率分布。蒸馏目标为最小化 \(D(p, r_\theta^{(\alpha,\lambda)})\) 或 \(D(q_\theta, r_\theta^{(\alpha,\lambda)})\)，D可以是任意散度。

关键设计¶

α-mixture辅助分布族
做什么：用参数α控制教师-学生分布插值路径的几何形状
核心思路：α=-1时为算术均值（m-mixture，直线路径）；α=1时为几何均值（e-mixture，对数空间直线）；α=3时为调和均值；其他α值给出新的插值路径。Theorem 3.2证明 \(r^{(\alpha,\lambda)}\) 是在α-散度意义下p和q的内分点（测地线上的点）
设计动机：α<1时support为并集（mode-covering），α≥1时support为交集（mode-seeking）——这直接影响蒸馏行为
最优性保证（Theorem 3.4）
做什么：证明对任意正则散度D和任意α，AMiD的最优解等价于p=q_θ
核心思路：如果辅助分布与教师完全匹配则p必须等于q_θ——插值点与一端重合时必须与另一端重合
设计动机：保证引入辅助分布不改变蒸馏的最终目标
梯度分析与mode-covering/seeking控制（Proposition 3.5）
做什么：分析α如何影响f-散度下的梯度行为
核心思路：梯度中出现加权项 \(w = \frac{(1-\lambda)q_\theta^{\frac{1-\alpha}{2}}}{\lambda p^{\frac{1-\alpha}{2}} + (1-\lambda)q_\theta^{\frac{1-\alpha}{2}}}\)，α较大时在p>q_θ区域w更大→mode-covering；α较小时在p<q_θ区域w更大→mode-seeking
设计动机：即使使用固定散度D，仍可通过α控制质量-多样性tradeoff

训练策略¶

兼容任意散度和数据策略，推荐α-β散度+λ=0.1
α<1用于模式覆盖，α≥1用于模式聚焦
支持自适应α调度

实验关键数据¶

主实验——GPT-2 XL→GPT-2蒸馏（指令跟随ROUGE-L）¶

方法	Dolly	Self-Inst	Vicuna	Super NI	Avg
GKD (α=-1)	24.58	11.78	14.60	22.84	~18
DistiLLM (α=-1)	~25	~12	~15	~23	~19
TAID (α=1)	~25	~12	~15	~23	~19
AMiD	最佳	最佳	最佳	最佳	最佳

消融——α的影响¶

α=-1到α=1之间的中间值（如α=0）在多数任务上表现最佳，说明现有方法用的端点值错过了最优区域。toy实验验证了α控制mode-covering/seeking的理论预测。

关键发现¶

α和λ是正交的设计维度——λ控制"走多远"，α控制"走哪条路径"
不同任务的最优α不同，但中间值通常优于端点值
AMiD训练更稳定，得益于辅助分布缓解近零概率问题

亮点与洞察¶

信息几何视角的统一极其优雅——用广义均值+α-散度内分点定理将零散方法统一为连续参数族
α与λ的正交性是核心洞察——之前所有工作只调λ不调α，错过了重要的设计维度
Proposition 3.5的梯度分析将mode-covering/seeking的直觉形式化

局限性 / 可改进方向¶

α的最优选择仍需实验调参，缺乏自动化机制
实验主要在GPT-2级别（0.1B-1.5B），大规模LLM验证不足
归一化常数 \(Z_r\) 增加计算开销
未与非KD压缩方法对比

评分¶

新颖性: ⭐⭐⭐⭐⭐ 信息几何驱动的统一框架，α作为新设计维度很深刻
实验充分度: ⭐⭐⭐⭐ 多任务+消融+toy验证完整，但模型规模偏小
写作质量: ⭐⭐⭐⭐⭐ 理论推导严谨，图示极其直观
价值: ⭐⭐⭐⭐ 为LLM知识蒸馏的辅助分布设计提供了统一理论基础