大模型知识蒸馏这个概念,可以理解为一种模型压缩与能力迁移技术。它的基本框架并不复杂:先有一个体积庞大、精度较高的教师模型,再构建一个参数量小得多的学生模型,让学生模型去拟合教师模型在相同输入下的输出分布。与普通监督学习只使用真实标签不同,知识蒸馏还把教师模型输出的概率向量当作训练信号,从而把教师模型学到的隐式知识传递给小模型。

以图像分类为例,假设教师模型对一张手写数字图片给出如下概率:数字3的概率为0.7,数字8的概率为0.2,数字5的概率为0.05。这里就包含了一条硬标签没有的信息:数字3和8在形状上相对接近,模型认为这张图片介于两者之间,但更偏向3。如果学生模型只学习硬标签,它只能知道答案是3,而无法获得关于类别相似度的信息。摄入教师概率分布后,学生模型可以继承这种类间关系,泛化能力通常优于从头训练。
从软标签看知识蒸馏的核心原理
监督学习中最常见的标签是 one-hot 向量,也就是把正确类别标为1、其余类别标为0。这种硬标签只表达正确类别,无法体现类别之间的相似性。例如在自然语言处理任务中,一个句子的情感倾向可能同时包含积极和中性成分;在图像识别中,动物类别里的哈士奇和阿拉斯加在特征空间中也比哈士奇和沙发更接近。教师模型输出的软标签恰好能捕捉这些细节,它给出的概率分布往往在多个类别上都有非零值,这些值的大小反映了教师模型对不同类别之间关系的判断。
但直接使用教师输出的原始 logits 作为软标签存在一个问题:大模型对自己的预测通常非常自信,正确类别的概率会接近1,其他类别的概率被压缩到极小,导致软标签又退化成了近似硬标签。这是知识蒸馏引入温度参数 T 的原因。在 softmax 函数中引入温度后,计算公式变为 softmax(z_i / T),T 越大,输出分布越平滑,类别间的细微差异被放大;T 越小,分布越尖锐。蒸馏时先用较高的温度让教师模型输出一个更“柔软”的分布,学生模型也使用相同温度计算预测分布,两者的差异通过 KL 散度度量,从而传递知识。
直观地说,硬标签告诉学生模型“这是3”,软标签则告诉学生模型“这是3,但它有一点像8,也稍微有一点像5”。后者这种模糊的、分级的相似性信息,是小模型提升泛化能力的重要来源。因此蒸馏不只是让小模型模仿大模型的最终答案,更是让它学习大模型的判断依据和决策边界。
温度系数、损失函数与训练流程
知识蒸馏的损失函数通常由两部分组成。第一部分是学生模型在真实硬标签上的标准交叉熵损失,它保证学生模型不会偏离正确答案太远。第二部分是蒸馏损失,它衡量学生模型与教师模型在同一温度下软输出之间的差异,一般使用 KL 散度。为了让梯度在不同温度下保持合适的尺度,蒸馏损失通常会乘以温度平方 T²。总体损失可以用下面的代码表示。
import torch
import torch.nn as nn
import torch.nn.functional as F
def distillation_loss(student_logits, teacher_logits, labels, T=4.0, alpha=0.7):
# 硬标签损失:学生模型对真实类别的交叉熵
hard_loss = F.cross_entropy(student_logits, labels)
# 软标签损失:教师与学生在温度T下的概率分布差异
soft_student = F.log_softmax(student_logits / T, dim=1)
soft_teacher = F.softmax(teacher_logits / T, dim=1)
distill_loss = F.kl_div(soft_student, soft_teacher, reduction='batchmean') * (T * T)
return alpha * hard_loss + (1 - alpha) * distill_loss
参数 alpha 用于控制硬标签损失和蒸馏损失的比例。当 alpha 设置为0时,学生模型完全依赖教师软标签;alpha 设为1时,则退化为普通监督训练。实际应用中 alpha 通常取0.3到0.9之间,具体值需要根据学生模型容量、教师模型强度以及任务数据量进行调整。温度 T 也不是越大越好,较高的温度可以让学生模型学到更多类别间关系,但会稀释正确类别的监督信号;较低的温度更接近硬标签,但蒸馏效果会减弱。常见设置在2到20之间,文本和图像任务有所差异。
训练流程并不复杂。先加载预训练好的教师模型并固定其参数,教师模型只做前向传播,不参与反向传播。然后初始化学生模型,每个 batch 同时送入教师和学生模型,分别得到两组 logits。之后按照上面的公式计算总体损失,只更新学生模型参数。为了减少显存占用,教师模型的 softmax 输出可以提前离线保存,但这种做法会失去样本间的随机性,通常只在数据量较小或推理成本极高时使用。更好的做法是在训练时动态计算教师输出,这样还能利用数据增强带来的多样性。
还需要注意一个细节:蒸馏损失使用的是 KL 散度,而 KL 散度不是对称的。一般做法是让学生分布去逼近教师分布,也就是 KL(student || teacher)。如果反向使用,梯度方向和优化效果会不同。PyTorch 的 F.kl_div 要求第一个参数是 log 概率,第二个参数是概率,因此代码中先对学生输出做 log_softmax,再对教师输出做 softmax。
大模型蒸馏的典型方式与工程实践
在大模型时代,知识蒸馏已经不只是最后一层输出上的拟合。像 BERT 这样的大型预训练模型参数量动辄几亿,将其蒸馏到小型模型时,仅靠输出层软标签往往不够。常见的扩展方式包括中间层蒸馏和注意力蒸馏。中间层蒸馏让学生模型的某些隐藏层去拟合教师模型对应层的输出,通过均方误差或余弦相似度拉近表示;注意力蒸馏则让学生模型模仿教师模型的注意力权重矩阵,将教师模型在不同词元之间分配注意力的方式传递给学生。这些方法可以单独使用,也可以与输出层蒸馏组合。
选择教师模型时需要评估教师的能力边界。教师模型越强,其软标签质量越高,但教师规模过大也会带来推理开销。尤其要在多个候选模型之间做蒸馏时,每次前向都会消耗计算资源。一种折中方案是先对教师模型做一次推理,把软标签缓存到磁盘,后续训练直接读取。不过缓存软标签会限制数据增强带来的新输入,因此不少工程实现采用在线计算教师输出,同时将教师模型放在混合精度或半精度模式下运行。
另一个实际的挑战是学生模型容量有限时,蒸馏可能无法完全恢复教师模型的精度。甚至在一些任务上,引入过多软标签信息反而会让学生模型更难收敛。此时需要调整教师模型温度、损失权重以及学习率。学习率通常比从头训练时略低,因为学生模型要拟合的分布相对平滑。优化器方面没有特殊限制,AdamW 或 SGD 都可用。下面给出一个简单的训练循环示例。
optimizer = torch.optim.AdamW(student_model.parameters(), lr=2e-5)
teacher_model.eval()
for batch in dataloader:
inputs, labels = batch
with torch.no_grad():
teacher_logits = teacher_model(inputs)
student_logits = student_model(inputs)
loss = distillation_loss(student_logits, teacher_logits, labels, T=4.0, alpha=0.7)
optimizer.zero_grad()
loss.backward()
optimizer.step()
这段代码中教师模型被设置为 eval 模式,并且前向计算放在 torch.no_grad() 里,避免梯度被教师模型的参数图额外占用显存。学生模型保持训练模式,每次迭代同时计算硬标签损失和蒸馏损失。这个流程可以直接用于文本分类、图像分类等任务。
总结来看,知识蒸馏并没有改变学生模型的结构,而是改变训练信号。它把大模型输出中的结构性信息作为额外的监督来源,让小模型在参数规模明显缩小的情况下,尽可能保留大模型的表达能力。在部署边缘设备、移动端或者对延迟敏感的服务时,这种技术几乎成为标配。不过蒸馏效果与教师模型、学生模型、温度、损失权重和数据分布都有关,工程中需要进行小规模实验来确定合适的配置。理解了软标签的来源与温度的作用,再结合中间层蒸馏等手段,就能根据具体任务设计出高效的大模型压缩方案。