4. 知识蒸馏

知识蒸馏，说白了就是「大模型教小模型」。

你训练了一个 7B 参数的 VLM，效果不错，但部署到边缘设备上根本跑不动。怎么办？让这个 7B 模型当老师，把它的「知识」压缩到一个 300M 的小模型里。这就是蒸馏的核心思路。

我在实际项目中遇到过，直接用小模型从头训练，精度差了 15 个点。但用蒸馏方案，差距能缩小到 2-3 个点。这笔账，划算。

4.1 教师-学生框架

框架本身不复杂。一个训练好的大模型（教师），一个待训练的小模型（学生）。学生模仿教师的输出行为，而不是直接学习原始数据标签。

核心流程：
  1. 教师推理：用教师模型对训练数据做前向推理，保存 logits 或中间特征
  2. 学生训练：学生模型同时学习真实标签和教师模型的输出
  3. 损失加权：用超参数 α 平衡硬标签损失和软标签损失

这里有个关键点：教师模型在推理时要关掉 dropout 和 batch norm 的 training 模式。我刚开始做蒸馏时没注意这个，教师输出的 logits 抖动很大，学生学得一塌糊涂。

# 伪代码：蒸馏训练循环
for batch in dataloader:
    images, texts, labels = batch
    
    # 教师推理（不计算梯度）
    with torch.no_grad():
        teacher_logits = teacher_model(images, texts)
    
    # 学生推理
    student_logits = student_model(images, texts)
    
    # 硬标签损失（交叉熵）
    loss_hard = ce_loss(student_logits, labels)
    
    # 软标签损失（KL散度）
    loss_soft = kl_div(
        log_softmax(student_logits / T),
        softmax(teacher_logits / T)
    ) * (T ** 2)
    
    # 总损失
    loss = alpha * loss_hard + (1 - alpha) * loss_soft
    loss.backward()
    optimizer.step()
温度 T 的直觉：T 越大，softmax 输出的概率分布越平滑，小模型能学到更多类别间的相对关系。我一般从 T=4 开始调，效果不好再降。

4.2 Logit 蒸馏与特征蒸馏

蒸馏到底「蒸」什么？主要有两条路：蒸输出（logit 蒸馏）和蒸中间层（特征蒸馏）。

维度 Logit 蒸馏 特征蒸馏
对齐目标 最后一层 logits 中间层特征图
实现难度 低，直接算 KL 散度 高，需要对齐维度
信息量 类别间关系 空间/语义结构
适用场景 分类任务 检测/分割/多模态

Logit 蒸馏：学生直接模仿教师的输出概率分布。好处是简单，坏处是信息瓶颈——教师的所有知识都被压缩到最后一层向量里。

特征蒸馏：让学生模仿教师的中间层特征。比如在视觉 backbone 的某个 stage，让学生的特征图和教师的特征图尽可能相似。

我做过一个实验：只用 logit 蒸馏，学生模型在 VQA 任务上掉了 4 个点。加上中间层的视觉特征蒸馏后，只掉 1 个点。为什么？因为 VLM 的视觉编码器学到了很多空间结构信息，这些信息在 logits 层已经被高度抽象了，学生根本学不到。

注意维度对齐：教师和学生的特征图尺寸通常不一样。需要用 1x1 卷积或 adapter 做投影，把学生特征映射到教师特征的空间。我曾经忘了加这个投影层，loss 直接崩了。

4.3 VLM 中的多模态蒸馏策略

VLM 比纯视觉模型复杂得多。它有视觉编码器、文本编码器、多模态融合模块。蒸馏时，每个模块都可以单独蒸馏，也可以联合蒸馏。

我总结了几种常见策略：

  • 视觉分支蒸馏：只蒸馏视觉 backbone 的特征。适合视觉编码器较大的场景（比如 ViT-L → ViT-B）
  • 文本分支蒸馏：蒸馏文本编码器的输出。适合文本侧较重的模型（比如用了大 LLM 做文本编码）
  • 跨模态对齐蒸馏：让学生模仿教师的图文相似度矩阵。这是 VLM 特有的蒸馏方式
  • 联合蒸馏：同时蒸馏多个模块，用加权损失组合

举个例子。我在做图文检索模型的蒸馏时，发现只蒸馏视觉分支，学生学到了图像特征，但图文对齐能力很差。后来加了跨模态对齐蒸馏——让学生和教师的图文相似度矩阵尽可能一致——效果直接提升了 8 个点。

# 跨模态对齐蒸馏损失
def cross_modal_distill_loss(student_sim, teacher_sim, T=4):
    """
    student_sim: [batch, batch] 学生模型的图文相似度矩阵
    teacher_sim: [batch, batch] 教师模型的图文相似度矩阵
    """
    # 对相似度矩阵做 softmax（按行）
    teacher_prob = F.softmax(teacher_sim / T, dim=1)
    student_prob = F.log_softmax(student_sim / T, dim=1)
    
    # KL 散度
    loss = F.kl_div(student_prob, teacher_prob, reduction='batchmean')
    return loss * (T ** 2)
我的经验：多模态蒸馏时，视觉分支和跨模态对齐的权重可以设大一点（0.5-0.7），文本分支权重小一点（0.1-0.3）。因为文本编码器通常已经比较轻量了。

4.4 蒸馏损失函数设计

损失函数是蒸馏的灵魂。选对了，事半功倍；选错了，学生学到的全是噪声。

常用的损失函数有几种：

  • KL 散度：最常用，适合 logit 蒸馏。衡量两个概率分布的差异
  • MSE 损失：适合特征蒸馏。直接算特征图的欧氏距离
  • 余弦相似度：适合对齐蒸馏。关注方向而非幅值
  • 对比损失：适合跨模态蒸馏。拉近正样本对，推远负样本对

我踩过一个坑：用 MSE 做特征蒸馏时，教师和学生的特征幅值差异很大，导致 loss 被幅值主导，忽略了方向信息。后来改用余弦相似度 + 温度缩放，效果好多了。

损失函数 公式 适用场景 注意事项
KL 散度 DKL(Pt || Ps) Logit 蒸馏 需要温度 T 平滑分布
MSE ||Ft - Fs||² 特征蒸馏 注意特征幅值对齐
余弦相似度 1 - cos(Ft, Fs) 对齐蒸馏 对幅值不敏感
对比损失 -log(exp(simpos/τ) / Σexp(sim/τ)) 跨模态蒸馏 需要构造正负样本对
损失组合技巧：我习惯用加权和的形式。比如 L = λ1Lkl + λ2Lmse + λ3Lcontrast。λ 的初始值可以设 1.0、0.5、0.5，然后根据验证集效果调整。

嗯，最后说一句。蒸馏不是万能的。如果教师模型本身就有偏差，学生学到的也是偏差。我见过有人拿一个有 bug 的教师模型蒸馏，学生模型反而比教师还差。所以，蒸馏前先确认教师模型的质量，这是基本前提。

VLM 知识蒸馏流程 教师模型 7B 参数 VLM 学生模型 300M 参数 VLM 输入数据（图文对） 教师 Logits + 中间特征 学生 Logits + 中间特征 蒸馏损失函数 L = α·Lhard + (1-α)·Lsoft 反向传播（仅学生） 关键点 教师模型冻结不更新梯度 | 温度 T 控制软标签平滑度 | 可同时蒸馏 logits 和中间特征 多模态 VLM 需额外蒸馏跨模态对齐矩阵
← 上一章 📖 返回目录 下一章 →