Huajishe Notes

10-蒸馏与模型压缩

25 min read

10 - 蒸馏与模型压缩

蒸馏（Distillation）是后训练链条的收口步骤，其核心目标是将大模型（Teacher）的能力迁移到小模型（Student），在保持能力的前提下实现压缩部署。蒸馏不仅是简单的模型压缩，更是一种知识迁移过程——Teacher 对数据的多维度理解以软形式传递给 Student，使得 Student 能在关键决策边界上学到比硬标签更丰富的概率分布信息。

1. 蒸馏的基本思想：Teacher → Student

蒸馏的范式可以形式化为：给定一个参数量大、表达能力强的 Teacher 模型 和一个参数量小、推理效率高的 Student 模型 ，通过让 拟合 的输出分布（logits、隐藏层、注意力矩阵等），实现能力迁移。

┌─────────────────────────────────────────────────────────┐
│                      蒸馏流程                            │
│                                                         │
│   Teacher (大模型)                                      │
│       │                                                 │
│       ├── logits / soft targets ─────────────────┐     │
│       ├── hidden states (特征) ──────────────┐   │     │
│       └── attention patterns ──────────┐     │   │     │
│                                        │     │   │     │
│   Student (小模型) ◄─────────────────┘     │   │     │
│       │                                    │   │     │
│       └── 蒸馏损失 (KD Loss) ─────────►  优化器       │
│                                                         │
└─────────────────────────────────────────────────────────┘

蒸馏的三层知识传递：

  1. 输出层蒸馏（Logits Distillation）：让 Student 拟合 Teacher 的输出概率分布。
  2. 表示层蒸馏（Feature Distillation）：让 Student 的隐藏层表示接近 Teacher 的隐藏层表示。
  3. 关系层蒸馏（Relation Distillation）：让 Student 学习 Teacher 对样本间关系、注意力模式的建模能力。

2. 数学推导：蒸馏损失函数

2.1 硬标签损失 vs 软标签损失

传统监督学习的交叉熵损失为：

其中 是硬标签（one-hot 或真实标签）， 是 Student 的预测概率。

蒸馏的核心改进是引入 温度缩放（Temperature Scaling） 的软目标。Teacher 的 logits 通过温度 进行缩放后计算 softmax，得到软标签：

2.2 KL 散度在蒸馏中的作用

蒸馏的软目标损失本质上是最小化 Student 分布与 Teacher 分布之间的 KL 散度（Kullback-Leibler Divergence）：

展开后，忽略常数项，等价于最小化：

为什么温度 会出现？ 因为对数空间中的梯度在温度较高时会更平滑，避免过大的梯度差异导致 Student 过度关注少数高置信度样本。推导如下：

在温度下的梯度为：

梯度会反比于温度 ，因此 KL 散度项乘以 用于补偿温度对梯度的缩放效应，保证不同 下损失函数的梯度尺度大致相当。

2.3 完整蒸馏损失

完整的蒸馏损失通常结合硬标签损失和软标签损失：

其中：

  • ：温度参数，通常 （常用值：2、3、4、5 或 8）
  • ：软标签损失的权重系数
  • 第一项：软目标蒸馏损失（捕捉 Teacher 的暗知识，dark knowledge）
  • 第二项：硬标签监督损失（保证在真实标签上的基本性能）

为什么软标签能传递”暗知识”？ 考虑一个 10 类分类问题。硬标签只告诉 Student”答案是第 3 类”，而 Teacher 的软标签可能显示”第 3 类的概率是 0.7，第 7 类是 0.2，第 1 类是 0.05…”——这种概率分布本身就编码了类别之间的语义相似性关系。例如，如果 Teacher 认为”猫”和”狗”的概率分布高度相关，说明它们在特征空间中是相邻的，这种信息在硬标签中被丢弃了，但蒸馏可以保留。

3. 蒸馏的四种形式

3.1 Logits 蒸馏（Output Distillation）

这是最基础的蒸馏形式，直接让 Student 的输出 logits 拟合 Teacher 的输出 logits：

特点：

  • 实现最简单，只需要 Teacher 的一次前向传播
  • 适合 Student 和 Teacher 结构相近的场景（如同架构的不同规模模型）
  • 对于结构差异大的模型（如 Transformer 到 LSTM），logits 维度不对应，难以直接对齐

3.2 隐藏层蒸馏（Feature Distillation）

当 Teacher 和 Student 结构不同时（如大 Transformer 到小 Transformer，或到 RNN），不能在 logits 层面直接对齐，需要在中间表示层进行蒸馏。

假设 Teacher 第 层的隐藏状态为 ，Student 第 层的隐藏状态为 ，使用适配层（projection layer）将 Student 表示映射到 Teacher 空间：

常见的适配方式：

  • 线性投影：，适用于
  • 注意力迁移（Attention Transfer）：让学生学习 Teacher 的注意力权重分布
  • Bottleneck 方式：在 Student 中间层引入瓶颈层，强迫其产生与 Teacher 某层维度匹配的表示

ResNet 到 ResNet 的隐藏层蒸馏：

Teacher: ResNet-101 (2048-d hidden) → Student: ResNet-50 (2048-d hidden)
Student 的某些中间层直接拟合 Teacher 对应层的输出

Teacher: ResNet-152 → Student: ResNet-56
FitAt: 学生第 56 层的输出直接拟合教师第 152 层

3.3 注意力蒸馏（Attention Distillation）

Transformer 模型中的注意力矩阵包含丰富的结构信息。假设第 层的注意力矩阵为 （ 为序列长度）：

注意力蒸馏损失：

注意：由于注意力矩阵是概率分布，使用 KL 散度而非 MSE 更合适。

3.4 序列级蒸馏（Sequence-Level Distillation）

在生成任务（翻译、摘要）中，不仅要让学生在 token 级别模仿 Teacher，还可以在序列级别进行蒸馏：

  • Beam Search 对齐：让 Student 在 Teacher 的最优 Beam 路径上进行训练
  • 最小风险训练（MRT）：直接优化序列级 reward：
  • SFT + KD 联合：先用 SFT 数据做监督微调，再用蒸馏数据做 KD 微调

4. 蒸馏与后训练链条的闭环

蒸馏不是孤立的步骤，而是与整个后训练链条紧密耦合。以下是蒸馏在完整流程中的位置：

预训练模型
    │
    ▼
┌─────────────────┐
│   PEFT (LoRA等) │ ◄── 可在蒸馏前作为 Teacher 高效微调
└─────────────────┘
    │
    ▼
┌─────────────────┐
│ SFT (监督微调)  │ ◄── 蒸馏 Teacher 常用流程
└─────────────────┘
    │
    ▼
┌─────────────────┐
│ RLHF / DPO      │ ◄── 对齐能力的关键环节
└─────────────────┘
    │
    ▼
┌─────────────────┐
│   蒸馏 (KD)      │ ◄── 收口步骤：将能力迁移到 Student
└─────────────────┘
    │
    ▼
压缩后的小模型 ───► 部署

4.1 蒸馏与 SFT 的串接

SFT（Supervised Fine-Tuning）阶段让 Teacher 模型在高质量数据集上学习人类期望的输出格式和内容。蒸馏时，Teacher 已经在 SFT 阶段对齐了人类偏好，Student 通过蒸馏继承这一对齐能力。

流程变体：

变体流程适用场景
SFT → KD先 SFT Teacher，再蒸馏标准流程，Teacher 已对齐
KD → SFT先蒸馏中间模型，再 SFT蒸馏后对齐能力丢失，需要再对齐
SFT + KD 交替分阶段交替训练能力衰减严重的场景

4.2 蒸馏与 RLHF / DPO 的串接

RLHF（Reinforcement Learning from Human Feedback）和 DPO（Direct Preference Optimization）用于对齐模型的输出与人化偏好。蒸馏在此的作用是将 Teacher 的对齐知识迁移到 Student。

关键问题：蒸馏能保留多少对齐能力？

能保留的场景：

  • Teacher 的对齐信号主要体现在输出分布上（logits 层面）
  • 使用 KL 散度损失可以保留类别间的相对偏好关系

可能丢失的场景：

  • 对齐信号隐藏在复杂的奖励模型中，logits 无法完全表达
  • RLHF 中的对抗性训练产生的非分布内（out-of-distribution）能力
  • DPO 中的偏好对比关系无法通过单样本蒸馏传递

DPO + 蒸馏的联合训练：

其中 保留偏好对比关系， 项保留 Teacher 的概率分布信息。

4.3 蒸馏与量化（Quantization）

蒸馏可以与量化结合，形成量化感知蒸馏（Quantization-Aware Distillation, QAD）：

  • 在蒸馏训练过程中，将 Teacher 和 Student 的计算都量化为低精度
  • Student 在量化约束下学习 Teacher 的全精度知识
  • 这种方法特别适合 INT8、INT4 部署场景

无校准数据时的蒸馏量化：

其中 是量化操作， 和 分别是 Student 和 Teacher 的全精度版本。

4.4 蒸馏与剪枝（Pruning）

蒸馏可以弥补剪枝带来的能力损失：

  1. 剪枝：移除不重要的权重或神经元
  2. 蒸馏：用原始 Teacher 指导剪枝后模型恢复能力

这形成修剪-蒸馏迭代循环：

原始大模型
    │
    ▼ (剪枝)
稀疏小模型
    │
    ▼ (蒸馏)
能力恢复的小模型
    │
    ▼ (迭代剪枝)
更稀疏的小模型
    │
    ▼ (迭代蒸馏)
最终压缩模型

4.5 蒸馏与 PEFT

PEFT（Parameter-Efficient Fine-Tuning）方法（如 LoRA、Adapter）在蒸馏场景下有特殊应用：

  • Teacher 端：对 Teacher 高效微调以适配特定任务，再进行蒸馏
  • Student 端：对 Student 高效微调以适应 Teacher 的输出分布
  • 双向适配：LoRA 适配层可以附加在任何位置，便于蒸馏损失的反向传播
┌──────────────────────────────────────────┐
│ Teacher (大模型, 全参数或 LoRA 微调)      │
│   └── logits + hidden + attention        │
└──────────────────────────────────────────┘
                    │
                    ▼
┌──────────────────────────────────────────┐
│ Student (小模型, LoRA 适配层训练)         │
│   └── 蒸馏损失通过 LoRA 适配层反传        │
└──────────────────────────────────────────┘

5. 工程流程

5.1 Teacher 数据生成

Teacher 数据的质量直接决定蒸馏效果。常见策略：

生成方式描述优缺点
SFT 数据集复用直接用 SFT 阶段的数据让 Teacher 推理最简单，但可能过拟合 Teacher 自身
伪标签生成用 Teacher 在未标注数据上生成软标签规模大，但噪声较多
增强数据用 Teacher 在原始数据上生成多样化的增强数据覆盖面广，计算开销大
合成数据用其他大模型生成高质量问答数据质量可控，依赖外部模型
人类标注数据直接使用人类标注数据质量最高，规模受限

关键原则：Teacher 生成的数据的分布要与部署场景匹配。蒸馏是知识迁移，不是知识创造——Teacher 不会的东西，Student 也不会学到。

5.2 Student 训练组织

蒸馏训练通常分为三阶段：

阶段 1：热启动

用硬标签（标准 SFT 数据）训练 Student 几个 epoch，确保 Student 具备基本能力。

阶段 2：蒸馏训练（核心）

切换到蒸馏损失，同时使用硬标签和软标签：

阶段 3：微调校准

在特定任务数据或偏好数据上做最后微调，可选择是否继续保留蒸馏损失。

5.3 混合硬标签、软标签和偏好数据

对于需要同时保留通用能力和特定对齐的场景，可以设计多目标损失的混合策略：

数据混合比例的经验法则：

  • 硬标签数据：蒸馏数据：偏好数据 ≈ 1:2:1（可根据任务调整）
  • 硬标签数据用于保持基础能力
  • 蒸馏数据用于从 Teacher 迁移知识
  • 偏好数据用于保持对齐

动态权重调度：在训练初期增大 （保证不遗忘），训练后期增大 （专注蒸馏 Teacher 知识）。

5.4 如何评测压缩前后能力差异

评测维度评测指标评测方式
通用能力MMLU、HumanEval、GSM8K标准基准测试
任务能力任务特定的 Accuracy / F1任务测试集
对齐能力人类评估、GPT-4 评估偏好对比
校准能力Expected Calibration Error (ECE)置信度 vs 准确率
泛化能力OOD 数据集表现分布外测试
效率提升延迟、吞吐量、显存占用部署基准测试

关键结论：蒸馏后的小模型在通用基准上通常能达到 Teacher 的 95-99%，但在分布外数据和对抗样本上差距会扩大——这说明蒸馏传递的是 Teacher 在训练分布内的知识，对分布边界的感知能力较弱。

6. 常见问题

6.1 模式坍塌（Mode Collapse）

表现：Student 只学习了 Teacher 输出的部分模式，忽略其他有效输出。

原因：

  • 温度 设置过低（），Teacher 的输出趋向 one-hot，Student 被迫模仿极端分布
  • 过大，Student 过度依赖 Teacher 的软标签，忽视硬标签的多样性信号

解决方案：

  • 使用适当的温度（）
  • 采用 逐渐冷却（Gradual Cooling）：训练初期用高温，后期降低温度
  • 加入 多样性正则项，鼓励 Student 探索非 Teacher 主导的模式

6.2 过拟合 Teacher

表现：Student 在训练数据上表现接近 Teacher，但在测试数据上差距明显。

原因：

  • 蒸馏训练轮次过多，Student 记忆了 Teacher 在特定数据上的输出噪声
  • Student 容量不足，强行拟合 Teacher 的高维输出导致过拟合

解决方案：

  • 使用独立的验证集监控蒸馏损失，当验证损失不再下降时提前停止
  • 在蒸馏损失中加入 权重衰减（Weight Decay）
  • 采用 拒绝采样（Rejection Sampling）：只蒸馏 Teacher 置信度高的样本

6.3 风格迁移不足

表现：Student 学到了 Teacher 的内容知识，但输出风格、格式与 Teacher 差距大。

原因：

  • 硬标签监督信号不足，Student 未充分学习期望的输出格式
  • 蒸馏数据中 Teacher 的风格变异未传递给 Student

解决方案：

  • 在硬标签数据中包含高质量的格式示例
  • 加入 风格一致性损失：
  • 使用 Chain-of-Thought 数据让 Teacher 展示推理过程

7. 小模型部署的收益与局限

收益

收益维度具体表现
延迟降低参数量减少 50%，推理延迟通常降低 40-60%（与硬件和结构相关）
显存节省7B 模型 INT8 量化后可在 16GB 显存的消费级 GPU 部署
吞吐量提升相同硬件下可服务更多并发请求
能源效率每次推理的能耗显著降低，适合边缘部署
成本下降云端部署的 token 成本可降低 5-10 倍

局限

局限维度具体表现
能力上限Student 的参数量决定了其能力上限，蒸馏无法让 Student 超越 Teacher
分布边界在训练分布外的泛化能力下降更明显
长尾能力Teacher 的长尾知识（低频但重要的知识）容易被蒸馏忽略
对齐漂移如果 Teacher 在 RLHF 阶段产生了对抗性能力，蒸馏可能丢失
多模态挑战跨模态蒸馏（如图像到文本）的知识对应关系难以建立

经验法则：如果 Teacher 在某项任务上的准确率是 ，Student 在相同任务上经过蒸馏后能达到 是合理的期望。如果差距超过 10%，通常说明蒸馏策略或数据存在问题。

8. 蒸馏训练流程图

开始
  │
  ▼
┌──────────────────────────────────────┐
│ 步骤 1: 准备 Teacher 和 Student      │
│ - Teacher: 微调后的大模型 (可选 PEFT) │
│ - Student: 初始化的小模型             │
│ - 对齐 Student/Teacher 架构           │
└──────────────────────────────────────┘
  │
  ▼
┌──────────────────────────────────────┐
│ 步骤 2: 生成蒸馏数据                 │
│ - 选择 Teacher 置信度高的样本         │
│ - 用 Teacher 推理生成软标签           │
│ - 收集硬标签 SFT 数据                │
└──────────────────────────────────────┘
  │
  ▼
┌──────────────────────────────────────┐
│ 步骤 3: 热启动 Student (可选)        │
│ - 用硬标签数据训练 1-2 epochs        │
│ - 确保 Student 基本能力               │
└──────────────────────────────────────┘
  │
  ▼
┌──────────────────────────────────────┐
│ 步骤 4: 蒸馏训练主循环               │
│ ┌────────────────────────────────┐   │
│ │ for epoch in range(num_epochs):  │   │
│ │   for batch in dataloader:       │   │
│ │     - 计算 KL 散度损失           │   │
│ │     - 计算 CE 损失 (可选)         │   │
│ │     - 反向传播 + 优化器更新        │   │
│ │   - 验证集评估                    │   │
│ │   - Early stopping 检查           │   │
│ └────────────────────────────────┘   │
└──────────────────────────────────────┘
  │
  ▼
┌──────────────────────────────────────┐
│ 步骤 5: 能力校准与微调              │
│ - 在任务数据上微调 (可选)            │
│ - DPO/RLHF 偏好对齐 (可选)          │
│ - 验证对齐能力是否保留              │
└──────────────────────────────────────┘
  │
  ▼
┌──────────────────────────────────────┐
│ 步骤 6: 压缩后处理                 │
│ - 量化 (INT8/INT4)                  │
│ - 剪枝 (如有需要)                   │
│ - 评测与基准测试                    │
└──────────────────────────────────────┘
  │
  ▼
结束 → 部署

9. 蒸馏超参建议表

超参数建议值范围说明
温度 2 ~ 8温度越高，软标签越平滑；生成任务可用更高温度
（软标签权重）0.5 ~ 0.9软标签权重越高，越依赖 Teacher 知识；初期可设高
蒸馏 Epochs3 ~ 20根据 Student 大小和数据量调整；大 Student 可更长
Batch Size8 ~ 128与标准 SFT 相同，注意显存约束
学习率1e-5 ~ 5e-5通常比 SFT 学习率略低
学习率调度Cosine / Linear Decay与 SFT 相同
权重衰减0.01 ~ 0.1防止过拟合 Teacher
Early Stopping验证集损失不再下降 2-3 个 epoch防止过拟合
Teacher 采样策略Top-p=0.9, Temperature=1.0生成蒸馏数据时 Teacher 用标准采样
学生初始化随机初始化 / Teacher 蒸馏初始化随机初始化 + 热启动是稳妥方案
蒸馏数据量Teacher 训练数据的 1x ~ 5x蒸馏数据可以重复使用多次
温度调度前期高温、后期低温可采用线性或指数降温

不同场景下的超参配置

场景其他调整
通用 NLP 任务2 ~ 40.7标准配置
代码生成3 ~ 50.8代码风格一致性损失
数学推理4 ~ 80.9CoT 数据蒸馏
对话系统2 ~ 30.6 ~ 0.7加入偏好数据混合
多模态 (Vision-Language)3 ~ 50.8跨模态注意力蒸馏

10. 总结

蒸馏是后训练链条中收口能力迁移的关键环节。它以 Teacher 的软输出为信号源，通过 KL 散度等损失函数将大模型的能力以概率分布的形式传递给小模型。与其将蒸馏单纯理解为”压缩”，不如理解为知识迁移——Teacher 对世界的多维度理解（类别关系、特征结构、注意力模式）以软形式编码在输出分布中，Student 通过学习这一分布来获得类似的世界认知。

蒸馏与 SFT 的关系是继承与增强：SFT 定义了”正确答案是什么”，蒸馏定义了”正确答案为什么比其他答案更可能”。蒸馏与 RLHF/DPO 的关系是能力保留：对齐训练产生的偏好信号需要通过蒸馏传递给压缩模型。蒸馏与量化、剪枝、PEFT 的关系是互补压缩：这些技术可以从不同角度实现模型压缩，蒸馏负责在压缩后恢复甚至提升能力。

蒸馏的本质局限在于：Teacher 不会的，Student 更不会。因此，提升 Teacher 质量、优化 Teacher 的对齐能力，是蒸馏出高质量 Student 的前提。

Graph View