知识蒸馏的基础概念
·
知识蒸馏(Knowledge Distillation)是一种模型压缩与加速技术,核心思想是把一个大而强的 教师模型(Teacher) 中学到的“软知识”迁移到一个体积更小、推理更快的 学生模型(Student) 上,使学生模型在保持较低计算成本的同时,尽可能接近教师模型的性能。
下面从原理、关键步骤、常见技巧以及应用场景四个方面进行阐述。
1. 原理概述
-
软标签(Soft Targets)
- 教师模型在对输入 x x x 进行前向传播后,会得到一组 logits(未归一化的得分)。
- 通过在 softmax 前加入温度系数
T
>
1
T>1
T>1 进行 温度软化:
p i ( T ) = exp ( z i / T ) ∑ j exp ( z j / T ) p_i^{(T)} = \frac{\exp(z_i/T)}{\sum_j \exp(z_j/T)} pi(T)=∑jexp(zj/T)exp(zi/T) - 当 T T T 较大时,概率分布会变得更平滑,保留了类别之间的相对相似度信息(即“暗知识”),这比硬标签(one‑hot)提供了更丰富的学习信号。
-
蒸馏损失(Distillation Loss)
- 学生模型同样产生 logits 并经过相同温度的 softmax,得到 q i ( T ) q_i^{(T)} qi(T)。
- 采用 KL 散度(或交叉熵)衡量两者分布的差异:
L KD = KL ( p ( T ) ∥ q ( T ) ) = ∑ p i T l o g ( p i T q i T ) \mathcal{L}_{\text{KD}} = \text{KL}\big(p^{(T)} \,\|\, q^{(T)}\big)=\sum p_{i}^{T}log(\frac{p_{i}^{T}}{q_{i}^{T}}) LKD=KL(p(T)∥q(T))=∑piTlog(qiTpiT) - 同时保留对原始硬标签的交叉熵损失 L CE \mathcal{L}_{\text{CE}} LCE。
- 综合损失一般写作:
L = α L CE + ( 1 − α ) T 2 L KD \mathcal{L} = \alpha \,\mathcal{L}_{\text{CE}} + (1-\alpha)\,T^2\,\mathcal{L}_{\text{KD}} L=αLCE+(1−α)T2LKD
其中 α \alpha α 控制两部分的权重, ( T 2 ) (T^{2}) (T2) 用于平衡梯度尺度。
-
学习过程
- 通过最小化上述综合损失,学生模型在学习原始任务的同时,也在模仿教师模型的输出分布,从而捕获教师模型对样本间细粒度关系的理解。
2. 关键步骤
| 步骤 | 说明 |
|---|---|
| ① 训练教师模型 | 使用全量数据、深层网络或集成模型,得到高精度的教师。 |
| ② 生成软标签 | 对训练集(或更大未标注数据)进行前向推理,记录 logits 并进行温度软化。 |
| ③ 设计学生模型 | 结构更轻(如层数更少、宽度更窄),或采用不同的网络单元(MobileNet、TinyBERT 等)。 |
| ④ 定义蒸馏损失 | 组合硬标签交叉熵与软标签 KL 散度,设置温度 T T T 与权重 α \alpha α。 |
| ⑤ 训练学生模型 | 使用梯度下降优化综合损失,通常采用更大的 batch size 以稳定软标签学习。 |
| ⑥ 评估与微调 | 在验证集上评估学生模型性能,必要时调节 T 、 α T、\alpha T、α 或加入额外正则化。 |
3. 常见蒸馏技巧
| 技巧 | 目的 | 说明 |
|---|---|---|
| 多层蒸馏(Layer‑wise Distillation) | 让学生学习教师的中间特征 | 通过额外的 特征匹配损失(如 MSE)对齐对应层的激活或注意力图。 |
| 自蒸馏(Self‑Distillation) | 在同一模型内部进行压缩 | 将模型的深层输出作为软标签,指导浅层或同层的子网络。 |
| 数据增强蒸馏 | 提升学生对输入扰动的鲁棒性 | 对同一样本进行多种增强,教师对所有增强的输出取平均作为软标签。 |
| 对抗蒸馏 | 强化学生对难样本的学习 | 通过生成对抗样本,使教师在这些样本上给出更明确的软标签。 |
| 多教师蒸馏 | 融合多模型知识 | 将多个教师的软标签进行加权平均或采用 集成蒸馏,提升学生的泛化能力。 |
| 跨模态蒸馏 | 将视觉、语言等不同模态的知识迁移 | 例如把视觉 Transformer 的特征蒸馏到文本模型中,实现多模态协同。 |
4. 应用场景
- 移动端/边缘计算:将大型 BERT、ResNet 等模型压缩为 TinyBERT、MobileBERT、MobileNet‑V2 等,满足低功耗、低延迟需求。
- 模型部署加速:在云服务中使用蒸馏模型降低 GPU/CPU 资源占用,提高并发请求处理能力。
- 知识迁移:把在大规模无标签数据上预训练的模型知识迁移到特定任务的轻量模型。
- 多任务学习:通过蒸馏统一多个任务的教师模型,得到一个能够兼顾多任务的学生模型。
- 模型安全:在对抗防御中,蒸馏可以帮助学生模型学习教师对抗样本的稳健输出,提升鲁棒性。
5. 小结
- 核心思想:利用教师模型的软标签(经过温度软化的概率分布)传递暗知识,使学生模型在更小的容量下逼近教师的表现。
- 实现要点:温度 (T)、权重 (\alpha) 的合理设置、适当的损失组合以及可能的中间层对齐。
- 优势:显著降低模型体积与推理成本,几乎不牺牲精度;同时可结合多种技巧提升蒸馏效果。
- 局限:需要先训练好教师模型,且蒸馏过程对超参数较敏感;在极端压缩(如 10× 以上)时仍可能出现性能下降。
ref:mita大模型生成
更多推荐
所有评论(0)