知识蒸馏

概述

知识蒸馏（Knowledge Distillation）是将大模型（教师）的知识转移到小模型（学生）的技术，在模型压缩和部署优化中发挥关键作用。从Hinton的经典方法到LLM时代的蒸馏策略，本章系统梳理知识蒸馏的理论与实践。

1. 经典知识蒸馏

1.1 Hinton蒸馏（2015）

核心思想：教师模型的软标签（soft labels）包含类别间关系的"暗知识"（dark knowledge）。

蒸馏损失：

\[ \mathcal{L} = \alpha \mathcal{L}_{\text{CE}}(y, \sigma(z_s)) + (1 - \alpha) T^2 D_{KL}\left(\sigma(z_t / T) \| \sigma(z_s / T)\right) \]

其中：

\(z_s, z_t\)：学生和教师的logits
\(\sigma\)：softmax函数
\(T\)：温度参数（通常 \(T=3\sim20\)）
\(\alpha\)：硬标签和软标签损失的权重
\(T^2\) 因子：补偿温度缩放对梯度的影响

温度的作用：

\[ \sigma(z_i / T) = \frac{\exp(z_i / T)}{\sum_j \exp(z_j / T)} \]

\(T=1\)：标准softmax
\(T \to \infty\)：均匀分布
\(T\) 越大，概率分布越平滑，暴露更多类别间关系

1.2 为什么软标签有效？

教师模型的输出 \([0.7, 0.2, 0.1]\) 比硬标签 \([1, 0, 0]\) 包含更多信息：

类别2比类别3更接近类别1
这种相似性关系是有价值的结构化知识
学生模型学到的不仅是"正确答案"，还有类别间的相对关系

2. 特征蒸馏

2.1 FitNets

FitNets（Romero et al., 2015）：不仅蒸馏输出，还蒸馏中间层特征。

\[ \mathcal{L}_{\text{FitNet}} = \| W_s \cdot F_s - F_t \|_2^2 \]

其中 \(F_s, F_t\) 是学生和教师的中间层特征，\(W_s\) 是维度匹配的投影矩阵。

优势：学生模型学到更好的中间表示。

2.2 注意力蒸馏

Attention Transfer（Zagoruyko & Komodakis, 2017）：蒸馏注意力图。

\[ \mathcal{L}_{\text{AT}} = \sum_l \left\| \frac{A_s^l}{\|A_s^l\|_2} - \frac{A_t^l}{\|A_t^l\|_2} \right\|_2^2 \]

其中注意力图 \(A^l = \sum_c |F^l_c|^2\)（沿通道维度求和的激活值平方）。

2.3 关系蒸馏

RKD（Relational Knowledge Distillation）：蒸馏样本间的关系。

\[ \mathcal{L}_{\text{RKD}} = \sum_{(i,j)} \left(\psi_t(t_i, t_j) - \psi_s(s_i, s_j)\right)^2 \]

其中 \(\psi\) 可以是距离关系或角度关系。

2.4 蒸馏方法总结

方法	蒸馏内容	层级
Hinton KD	Soft logits	输出层
FitNets	中间特征	隐藏层
Attention Transfer	注意力图	注意力层
RKD	样本间关系	表示空间
CRD	对比表示	表示空间
PKD	耐心蒸馏（多层）	多层

3. NLP中的知识蒸馏

3.1 DistilBERT

DistilBERT（Sanh et al., 2019）：BERT的蒸馏版本。

设计：

6层（原BERT 12层），参数减少40%
蒸馏损失 = 软标签 + 硬标签 + 余弦相似度

\[ \mathcal{L} = \alpha \mathcal{L}_{\text{CE}} + \beta \mathcal{L}_{\text{KD}} + \gamma \mathcal{L}_{\text{cos}} \]

结果：保留BERT 97%的性能，速度快60%。

3.2 TinyBERT

TinyBERT（Jiao et al., 2020）：更全面的BERT蒸馏。

四层蒸馏：

Embedding层：\(\mathcal{L}_{\text{emb}} = \text{MSE}(E_s W_e, E_t)\)
注意力层：\(\mathcal{L}_{\text{attn}} = \text{MSE}(A_s, A_t)\)
隐藏层：\(\mathcal{L}_{\text{hid}} = \text{MSE}(H_s W_h, H_t)\)
预测层：\(\mathcal{L}_{\text{pred}} = \text{KD}\)

两阶段：通用蒸馏（预训练阶段） + 任务蒸馏（微调阶段）

3.3 MiniLM

MiniLM（Wang et al., 2020）：蒸馏自注意力中的Q-K关系和V-V关系。

\[ \mathcal{L} = D_{KL}(A_t^L \| A_s^l) + D_{KL}(V_t^L {V_t^L}^\top \| V_s^l {V_s^l}^\top) \]

4. LLM时代的知识蒸馏

4.1 LLM蒸馏的挑战

挑战	描述
规模差距	教师动辄100B+，学生通常1-7B
白盒/黑盒	很多API模型无法获取logits
任务多样性	LLM是通用模型，不是单任务
涌现能力	小模型可能无法复现涌现行为

4.2 白盒蒸馏

当可以获取教师模型的logits时：

标准方法：

\[ \mathcal{L} = \lambda \mathcal{L}_{\text{KD}}(p_t, p_s) + (1-\lambda) \mathcal{L}_{\text{NTP}}(y, p_s) \]

MiniLLM：用反向KL散度替代正向KL：

\[ \mathcal{L}_{\text{MiniLLM}} = D_{KL}(p_s \| p_t) \]

反向KL鼓励学生集中在教师概率高的区域（mode-seeking），避免过度分散。

4.3 黑盒蒸馏

当只能获取教师的文本输出时：

数据生成：用教师生成高质量训练数据
Alpaca/Vicuna方式：用GPT-4生成指令数据训练小模型
Self-Instruct：用教师自动生成指令-响应对

4.4 代表性工作

方法	教师	学生	策略
Alpaca	text-davinci-003	LLaMA-7B	黑盒数据蒸馏
Vicuna	GPT-4	LLaMA-13B	黑盒对话蒸馏
Orca	GPT-4	LLaMA-13B	解释增强蒸馏
MiniLLM	GPT-2 XL	GPT-2	白盒反向KL
GKD	PaLM	-	在线蒸馏

5. 自蒸馏（Self-Distillation）

5.1 概念

模型自己作为教师蒸馏自己：

Born-Again Networks：训练多代，每代用上一代作为教师
DINO/DINOv2：EMA教师 + 学生的自蒸馏
Noisy Student：自训练 + 噪声

5.2 LLM中的自蒸馏

自我改进（Self-Improve）：模型生成数据训练自己
STaR：自我教授推理
宪法AI（Constitutional AI）：自我评判和改进

6. 实践指南

6.1 蒸馏超参数

超参数	推荐值	说明
温度 \(T\)	3-20	越大越平滑
\(\alpha\)	0.5-0.9	软标签权重
学生深度	教师的1/2-2/3	太浅效果差
学生宽度	教师的1/2-3/4	宽度比深度重要

6.2 蒸馏效果经验

教师-学生差距不宜过大（否则学不好）
多步蒸馏（Teacher → TA → Student）可能更好
数据质量比数量重要
特征蒸馏通常比logits蒸馏效果好

7. 总结

graph TD
    A[知识蒸馏] --> B[Logits蒸馏]
    A --> C[特征蒸馏]
    A --> D[关系蒸馏]
    A --> E[数据蒸馏]

    B --> B1[Hinton KD]
    B --> B2[MiniLLM]

    C --> C1[FitNets]
    C --> C2[TinyBERT]

    D --> D1[RKD]
    D --> D2[CRD]

    E --> E1[Alpaca]
    E --> E2[Vicuna]

核心要点：

经典蒸馏利用软标签中的暗知识
特征蒸馏传递中间表示
LLM蒸馏面临规模和黑盒挑战
黑盒蒸馏（数据生成）是LLM时代的主流
自蒸馏是一个不断增长的方向

参考文献

Hinton et al., "Distilling the Knowledge in a Neural Network," 2015
Romero et al., "FitNets: Hints for Thin Deep Nets," ICLR 2015
Sanh et al., "DistilBERT, a distilled version of BERT," NeurIPS Workshop 2019
Jiao et al., "TinyBERT: Distilling BERT for Natural Language Understanding," EMNLP 2020
Gu et al., "MiniLLM: Knowledge Distillation of Large Language Models," ICLR 2024