知识蒸馏的基本原理

知识蒸馏 (knowledge distillation)（KD）作为一种有效的模型压缩技术，尤其适用于大型语言模型，其基本原理被讨论。这项技术的主要思想源自 Hinton、Vinyals 和 Dean (2015) 的工作，即将一个熟练但庞大的教师模型所捕获的“知识”迁移给一个更紧凑的学生模型。

本质上，KD 采用这样一种理念，即经过训练的教师模型的输出分布所含信息，比其初始训练期间使用的真实标签更为丰富。考虑一个标准分类任务。训练通常使用独热编码向量 (vector)（硬目标），其中正确类别的概率为 1，所有其他类别为 0。然而，教师模型在训练后会生成一个覆盖所有可能输出的概率分布。例如，在预测序列中的下一个词时，教师模型可能会给正确词语分配高概率，但也会给其他合理（甚至语义相关但错误）的词语分配小的非零概率。这种更丰富的分布，通常被称为“暗知识”，它表明了教师模型所学到的输出之间的相似性和关系。

标准的 KD 方法通过训练学生模型模仿教师模型的软化输出概率来获取这种暗知识。这是通过在 softmax 函数中应用温度缩放参数 (parameter) $T$ 来实现的。对于生成对数 $z$ 的模型，类别 $i$ 的标准 softmax 概率是 $p_i = \frac{\exp(z_i)}{\sum_j \exp(z_j)}$。经过温度缩放后，软化概率 $q_i$ 变为：

$$q_i = \frac{\exp(z_i / T)}{\sum_j \exp(z_j / T)}$$

更高的温度 $T > 1$ 会软化概率分布，增加熵并赋予较小对数值（“暗知识”）更大的权重 (weight)。当 $T \to \infty$ 时，分布趋于均匀；而当 $T=1$ 时，则恢复标准 softmax。

随后，学生模型使用一个损失函数 (loss function)进行训练，该函数促使其软化输出（使用相同温度 $T$ 计算）与教师模型的软化输出匹配。Kullback-Leibler (KL) 散度常用于此目的：

$$L_{KD} = T^2 \cdot D_{KL}(q_{student} || q_{teacher})$$

这里，$q_{student}$ 和 $q_{teacher}$ 分别是学生模型和教师模型产生的软化概率分布。$T^2$ 缩放因子确保梯度在温度变化时量级保持可比。注意，在此过程中，教师模型的参数是固定的；只有学生模型在进行训练。

通常，这种蒸馏损失与一个标准损失函数（例如交叉熵）相结合，该标准损失函数计算学生模型的预测结果（使用 $T=1$）与真实的硬标签 ($y$) 之间的差异。这有助于确保学生模型在原始任务目标上仍能表现良好。最终的损失函数成为一个加权平均：

$$L_{total} = \alpha L_{CE}(p_{student}, y) + (1 - \alpha) L_{KD}(q_{student}, q_{teacher})$$

超参数 (hyperparameter) $\alpha$（通常在 0 到 1 之间）平衡了硬目标损失 ($L_{CE}$) 和软目标蒸馏损失 ($L_{KD}$) 的贡献。为 $T$ 和 $\alpha$ 选择合适的值对于成功的蒸馏非常重要，并且通常需要经验性调整。

使用软目标的基础知识蒸馏过程。教师模型使用温度缩放生成软目标，学生模型通过 KD 损失 ($L_{KD}$) 尝试模仿这些软目标。学生模型也通过标准交叉熵损失 ($L_{CE}$) 从真实值中学习。

教师模型和学生模型不一定需要相同的架构，尽管在蒸馏 LLM 时这很常见（例如，将一个 700 亿参数的 Llama 模型蒸馏到一个 70 亿参数的 Llama 模型）。主要要求是它们的输出层与蒸馏损失的计算兼容。学生模型在软目标中包含的丰富监督信号的引导下，学会逼近教师模型所学习的复杂函数。这通常会使学生模型比仅从头开始使用硬目标进行训练获得更好的泛化能力和表现，从而有效地迁移大型教师模型所学到的归纳偏置 (bias)。

尽管匹配最终输出分布是经典方法，但必须认识到，这仅仅是知识迁移的一个方面。随后的章节将研究更高级的蒸馏目标，这些目标采用中间表示和注意力机制 (attention mechanism)，提供其他途径，让学生模型具备教师模型的能力。