2015年3月,Geoffrey Hinton在arXiv上发表了一篇只有9页的论文,标题是《Distilling the Knowledge in a Neural Network》。这篇论文没有提出什么新的网络架构,也没有刷新任何榜单,却彻底改变了模型部署的游戏规则。

Hinton在论文中提出了一个看似简单的问题:训练阶段和部署阶段的需求完全不同,为什么我们总是使用相同的模型?

训练时,我们需要从海量数据中提取结构,可以使用大量的计算资源。部署时,我们面临的是完全不同的约束:延迟敏感、内存受限、功耗有限。昆虫有幼虫和成虫两种完全不同的形态来适应不同阶段的需求,神经网络为什么不能?

知识蒸馏(Knowledge Distillation)就是这个问题的答案。

一个反直觉的核心洞察

在传统的分类任务中,模型被训练成对正确答案给出高概率,对错误答案给出低概率。比如一张狗的图片,训练好的模型可能给出这样的预测:

  • 狗:99.9%
  • 猫:0.05%
  • 汽车:0.03%
  • 飞机:0.02%

交叉熵损失函数会告诉模型"狗是正确答案,其他都是错的",于是模型学会了把这个概率推向极端。

但Hinton注意到了一个被忽视的细节:错误答案之间的相对概率差异,包含了重要的信息

一张狗的图片被误认为"狼"的概率(0.05%),远高于被误认为"汽车"的概率(0.001%)。这告诉我们:这只狗看起来有点像狼,但绝不像汽车。这种"错误的分布"——错误答案之间的相对关系——才是模型真正学到的知识。Hinton称之为"暗知识"(Dark Knowledge)。

flowchart LR
    subgraph 传统训练
        A1["输入: 狗的图片"] --> B1["模型输出"]
        B1 --> C1["狗: 99.9%<br/>狼: 0.05%<br/>猫: 0.03%<br/>汽车: 0.02%"]
        C1 --> D1["交叉熵损失<br/>只关注正确答案"]
    end
    
    subgraph 知识蒸馏
        A2["输入: 狗的图片"] --> B2["教师模型输出"]
        B2 --> C2["软标签分布<br/>包含'暗知识'"]
        C2 --> D2["学生模型学习<br/>完整的概率分布"]
    end

传统的交叉熵训练完全忽略了这些信息。因为正确答案的概率已经接近1,其他所有概率都接近0,梯度几乎为零。这些"错误但有意义"的信息,被训练目标无情地抛弃了。

知识蒸馏的核心思想就是:让学生模型学习教师模型的完整输出分布,而不仅仅是正确答案

温度参数:软化概率分布的艺术

既然暗知识隐藏在低概率区域,为什么不让这些概率"显形"?

这就是温度参数(Temperature)的作用。标准的softmax函数定义为:

$$p_i = \frac{\exp(z_i)}{\sum_j \exp(z_j)}$$

引入温度$T$后,公式变为:

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

当$T=1$时,这就是标准的softmax。当$T>1$时,概率分布变得更加平滑(softer)。

graph LR
    subgraph "标准Softmax T=1"
        A1["猫: 95%"] --> B1["狗: 4%"]
        B1 --> C1["鸟: 1%"]
    end
    subgraph "软化Softmax T=5"
        A2["猫: 60%"] --> B2["狗: 30%"]
        B2 --> C2["鸟: 10%"]
    end

温度越高,概率分布越均匀,原本被压缩的小概率值被"放大"了。以一个具体的数字例子说明:

假设模型的原始logits输出为$[10, 5, 1]$(分别对应猫、狗、鸟):

  • 当$T=1$时:softmax输出约为$[0.993, 0.007, 0.000]$——几乎所有的概率都集中在第一个类别
  • 当$T=5$时:softmax输出约为$[0.630, 0.300, 0.070]$——概率分布变得平滑,原本被忽略的"狗"和"鸟"的概率变得可见
graph TB
    subgraph "Logits输入 [10, 5, 1]"
        L1["猫: 10"] --> S1["Softmax"]
        L2["狗: 5"] --> S1
        L3["鸟: 1"] --> S1
    end
    
    S1 --> T1["T=1: [99.3%, 0.7%, 0.0%]"]
    S1 --> T2["T=2: [91.9%, 7.6%, 0.5%]"]
    S1 --> T5["T=5: [63.0%, 30.0%, 7.0%]"]
    S1 --> T10["T=10: [47.4%, 34.2%, 18.4%]"]
    
    style T1 fill:#ffcccc
    style T5 fill:#ccffcc
    style T10 fill:#ccccff

这就是温度参数的魔法:它让原本被压制的"暗知识"浮出水面。

Hinton在原始论文中给出了一个精妙的类比:高温度下的软目标,就像是把一张被揉皱的纸展平——原本折叠在角落里的信息,现在清晰可见了。

蒸馏损失函数:两条学习路径

知识蒸馏的损失函数由两部分组成:

$$\mathcal{L} = \alpha \cdot \mathcal{L}_{\text{hard}} + (1-\alpha) \cdot \mathcal{L}_{\text{soft}}$$

其中:

  • $\mathcal{L}_{\text{hard}}$:学生模型预测与真实标签之间的交叉熵损失
  • $\mathcal{L}_{\text{soft}}$:学生模型与教师模型软化输出之间的KL散度损失
  • $\alpha$:平衡两个损失的权重参数
flowchart TB
    subgraph "知识蒸馏损失计算"
        Input["输入数据"] --> Teacher["教师模型<br/>(冻结)"]
        Input --> Student["学生模型<br/>(训练中)"]
        
        Teacher --> SoftTeacher["软标签<br/>Softmax(z/T)"]
        Student --> SoftStudent["学生软输出<br/>LogSoftmax(z/T)"]
        Student --> HardOutput["硬标签预测<br/>Softmax(z)"]
        
        SoftTeacher --> KLLoss["KL散度损失<br/>× T²"]
        SoftStudent --> KLLoss
        HardOutput --> CELoss["交叉熵损失"]
        TrueLabel["真实标签"] --> CELoss
        
        KLLoss --> TotalLoss["总损失 = α×CE + (1-α)×KL"]
        CELoss --> TotalLoss
    end

KL散度损失的计算方式为:

$$\mathcal{L}_{\text{soft}} = T^2 \cdot \text{KL}(p_{\text{teacher}}^T \| p_{\text{student}}^T)$$

这里有一个容易被忽略的细节:为什么要乘以$T^2$?

原因是softmax输出的梯度会随着温度的升高而缩小。当$T$增大时,概率分布变得更加均匀,梯度也随之变小。乘以$T^2$可以补偿这种缩小效应,确保软标签损失和硬标签损失的梯度量级在同一水平。

Hinton在论文中证明了一个有趣的结论:当温度足够高时,蒸馏损失等价于最小化教师和学生logits之间的均方误差。这提供了一个直观的理解:蒸馏本质上是让学生模型的输出尽可能接近教师模型的输出

PyTorch实现:从理论到代码

完整的知识蒸馏实现并不复杂。以下是一个可以直接使用的PyTorch实现:

import torch
import torch.nn as nn
import torch.nn.functional as F

class DistillationLoss(nn.Module):
    """
    知识蒸馏损失函数
    
    Args:
        temperature: 温度参数,默认为4.0
        alpha: 硬标签损失权重,默认为0.7
    """
    def __init__(self, temperature=4.0, alpha=0.7):
        super().__init__()
        self.temperature = temperature
        self.alpha = alpha
        self.ce_loss = nn.CrossEntropyLoss()
        self.kl_loss = nn.KLDivLoss(reduction='batchmean')
    
    def forward(self, student_logits, teacher_logits, labels):
        # 硬标签损失:学生模型预测与真实标签的交叉熵
        hard_loss = self.ce_loss(student_logits, labels)
        
        # 软标签损失:KL散度
        # 注意:学生使用log_softmax,教师使用softmax
        soft_student = F.log_softmax(student_logits / self.temperature, dim=-1)
        soft_teacher = F.softmax(teacher_logits / self.temperature, dim=-1)
        
        # 乘以T^2补偿梯度缩放
        soft_loss = self.kl_loss(soft_student, soft_teacher) * (self.temperature ** 2)
        
        # 加权组合
        return self.alpha * hard_loss + (1 - self.alpha) * soft_loss


def train_with_distillation(teacher, student, dataloader, optimizer, 
                            temperature=4.0, alpha=0.7, device='cuda'):
    """
    使用知识蒸馏训练学生模型
    """
    criterion = DistillationLoss(temperature, alpha)
    teacher.eval()  # 教师模型冻结
    
    for epoch in range(num_epochs):
        for inputs, labels in dataloader:
            inputs, labels = inputs.to(device), labels.to(device)
            
            # 教师模型推理(不计算梯度)
            with torch.no_grad():
                teacher_logits = teacher(inputs)
            
            # 学生模型推理
            student_logits = student(inputs)
            
            # 计算蒸馏损失
            loss = criterion(student_logits, teacher_logits, labels)
            
            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
sequenceDiagram
    participant D as 数据
    participant T as 教师模型
    participant S as 学生模型
    participant L as 损失函数
    participant O as 优化器
    
    loop 每个batch
        D->>T: 输入数据
        T->>T: 前向传播(无梯度)
        T-->>L: 教师logits
        D->>S: 输入数据
        S->>S: 前向传播
        S-->>L: 学生logits
        D-->>L: 真实标签
        L->>L: 计算蒸馏损失
        L->>O: 返回损失
        O->>S: 更新学生权重
        Note over T: 教师权重不变
    end

这个实现中有几个关键点需要注意:

教师模型必须处于评估模式。教师模型的权重是固定的,我们只更新学生模型。使用torch.no_grad()上下文管理器可以避免不必要的计算图构建。

学生使用log_softmax,教师使用softmax。这是因为KLDivLoss期望的输入是log概率,而不是概率。这是一个常见的陷阱,直接使用softmax会导致错误的梯度计算。

温度参数的选择需要实验。Hinton在原始论文中推荐$T=2$到$T=20$的范围,但最优值取决于具体的任务和模型。一般来说,模型越小,最优温度越低;类别数越多,最优温度可能越高。

三种蒸馏范式:响应、特征与关系

知识蒸馏不仅仅只有一种形式。根据迁移的知识类型,可以分为三类:

graph TB
    subgraph "知识蒸馏三种范式"
        R["响应蒸馏<br/>Response-based"] --> R1["迁移最终输出"]
        R1 --> R2["实现简单<br/>应用最广"]
        
        F["特征蒸馏<br/>Feature-based"] --> F1["迁移中间层特征"]
        F1 --> F2["需要架构对齐<br/>效果更好"]
        
        Rel["关系蒸馏<br/>Relation-based"] --> Rel1["迁移特征间关系"]
        Rel1 --> Rel2["最灵活<br/>适合异构模型"]
    end

响应蒸馏(Response-based Distillation)

这是最经典的蒸馏方式,直接让学生模型模仿教师模型的最终输出。上文讨论的温度参数和软标签都属于这一类。

响应蒸馏的优点是实现简单,只需要访问模型的最终输出层。缺点是只能迁移"结果"知识,无法传递教师模型的内部表示。

特征蒸馏(Feature-based Distillation)

特征蒸馏的核心思想是:让学生模型的中间层特征图尽可能接近教师模型。

这需要解决一个架构问题:教师和学生模型的中间层维度可能不同。比如教师模型有2048个通道,学生模型只有512个通道。

flowchart LR
    subgraph 教师模型
        TI["输入"] --> TC1["Conv 2048ch"]
        TC1 --> TC2["Conv 2048ch"]
        TC2 --> TO["输出"]
    end
    
    subgraph 学生模型
        SI["输入"] --> SC1["Conv 512ch"]
        SC1 --> SC2["Conv 512ch"]
        SC2 --> SO["输出"]
    end
    
    TC1 -.->|"特征蒸馏"| SC1
    TC2 -.->|"特征蒸馏"| SC2
    TO -.->|"响应蒸馏"| SO
    
    style TC1 fill:#ffcccc
    style SC1 fill:#ccffcc

常见的解决方案有三种:

# 方案1:添加适配层(Adaptor)
class StudentWithAdaptor(nn.Module):
    def __init__(self, student_channels, teacher_channels):
        super().__init__()
        self.student = StudentModel()
        self.adaptor = nn.Conv2d(student_channels, teacher_channels, 1)
    
    def forward(self, x):
        features = self.student.extract_features(x)
        adapted_features = self.adaptor(features)
        return adapted_features

# 方案2:使用池化降维
teacher_features_pooled = F.adaptive_avg_pool2d(teacher_features, 
                                                 output_size=student_features.shape[2:])

# 方案3:计算余弦相似度而非逐元素匹配
cosine_loss = 1 - F.cosine_similarity(student_features, teacher_features, dim=1).mean()

特征蒸馏的核心损失函数通常使用均方误差或余弦嵌入损失:

$$\mathcal{L}_{\text{feature}} = \frac{1}{HW} \sum_{i,j} (f_s^{i,j} - f_t^{i,j})^2$$

其中$f_s$和$f_t$分别是学生和教师的特征图。

关系蒸馏(Relation-based Distillation)

关系蒸馏的思路更进一步:我们不要求学生的特征完全匹配教师,而是要求特征之间的"关系结构"保持一致。

比如,输入两张图片A和B,教师模型提取的特征向量分别为$f_t^A$和$f_t^B$,学生模型提取的特征向量分别为$f_s^A$和$f_s^B$。关系蒸馏要求:

$$\text{sim}(f_s^A, f_s^B) \approx \text{sim}(f_t^A, f_t^B)$$
graph TB
    subgraph "教师模型"
        T1["输入A"] --> TF1["特征 f_t^A"]
        T2["输入B"] --> TF2["特征 f_t^B"]
        TF1 --> TR["关系: sim(f_t^A, f_t^B)"]
        TF2 --> TR
    end
    
    subgraph "学生模型"
        S1["输入A"] --> SF1["特征 f_s^A"]
        S2["输入B"] --> SF2["特征 f_s^B"]
        SF1 --> SR["关系: sim(f_s^A, f_s^B)"]
        SF2 --> SR
    end
    
    TR -.->|"匹配关系结构"| SR

这种方法的优点是:学生不需要精确复制教师的特征,只需要保持相同的"结构"。这对于架构差异大的教师-学生对特别有用。

大模型蒸馏的实践:从DistilBERT到DeepSeek R1

DistilBERT:保留97%性能的轻量BERT

DistilBERT是将知识蒸馏应用于Transformer架构的开创性工作。它的架构非常简单:直接删除BERT-base的一半层,然后用蒸馏训练剩下的部分。

graph LR
    subgraph "BERT-base (教师)"
        B1["Embedding"] --> BL1["Layer 1"]
        BL1 --> BL2["Layer 2"]
        BL2 --> BL3["..."]
        BL3 --> BL4["Layer 12"]
        BL4 --> BO["Output"]
    end
    
    subgraph "DistilBERT (学生)"
        D1["Embedding"] --> DL1["Layer 1"]
        DL1 --> DL2["Layer 2"]
        DL2 --> DL3["Layer 3"]
        DL3 --> DL4["Layer 6"]
        DL4 --> DO["Output"]
    end
    
    BO -.->|"输出蒸馏"| DO

但DistilBERT的蒸馏策略很巧妙:

  1. 输出蒸馏:学生模型学习教师模型的logits(温度$T=2$)
  2. 隐藏层蒸馏:学生模型的隐藏层状态与教师模型对齐
  3. 注意力蒸馏:学生模型的注意力图与教师模型对齐

损失函数是三者的加权和:

$$\mathcal{L} = \alpha_{\text{mlm}} \mathcal{L}_{\text{mlm}} + \alpha_{\text{cos}} \mathcal{L}_{\text{cos}} + \alpha_{\text{distil}} \mathcal{L}_{\text{distil}}$$

实验结果令人印象深刻:DistilBERT保留了BERT-base 97%的性能,同时参数量减少了40%,推理速度提升了60%。

TinyBERT:两阶段蒸馏策略

TinyBERT提出了一个更系统的蒸馏框架:预训练阶段蒸馏 + 任务特定蒸馏

timeline
    title TinyBERT两阶段蒸馏
    section 预训练阶段
        通用知识迁移 : 词嵌入蒸馏<br/>隐藏层蒸馏<br/>注意力蒸馏
    section 任务特定阶段
        任务知识迁移 : 针对具体任务<br/>精细微调蒸馏

在预训练阶段,学生模型从教师模型学习通用的语言表示。在任务特定阶段,学生模型针对具体任务(如分类、问答)进行微调蒸馏。

TinyBERT的关键创新是引入了嵌入层蒸馏:学生模型的词嵌入向量需要与教师模型的词嵌入向量对齐。由于学生和教师的嵌入维度不同,TinyBERT使用一个可学习的线性投影矩阵:

$$\mathcal{L}_{\text{embd}} = \|W_e E_s - E_t\|_2^2$$

其中$W_e$是可学习的投影矩阵,$E_s$和$E_t$分别是学生和教师的嵌入矩阵。

DeepSeek R1:推理能力的蒸馏

2025年初发布的DeepSeek R1展示了一个新的蒸馏范式:推理链蒸馏

传统的知识蒸馏主要迁移分类或预测能力,但DeepSeek R1展示了一个更高级的目标——迁移推理能力本身。

DeepSeek R1的核心发现是:让大模型生成详细的思维链(Chain-of-Thought),然后用这些思维链训练小模型,可以让小模型获得接近大模型的推理能力。

flowchart TB
    subgraph "传统蒸馏"
        Q1["问题"] --> T1["教师模型"]
        T1 --> A1["答案"]
        A1 --> S1["学生学习"]
    end
    
    subgraph "推理链蒸馏"
        Q2["问题"] --> T2["教师模型"]
        T2 --> COT["思维链<br/>一步步推理过程"]
        COT --> A2["答案"]
        A2 --> S2["学生学习完整推理"]
        COT --> S2
    end

实验数据显示:

训练方式 GSM8K准确率 平均推理长度
直接预测答案 29% N/A
人类专家思维链 68% 280字符
R1合成思维链 87% 2000字符

R1生成的思维链比人类专家更长、更详细,这反而带来了更好的性能。这揭示了一个重要的洞见:好的教师不仅要给出正确答案,还要展示完整的推理过程

温度参数如何选择?

温度参数的选择是知识蒸馏中最关键的超参数决策之一。虽然没有万能公式,但有一些经验法则可以参考:

graph TD
    subgraph "温度选择指南"
        Small["模型容量差距大"] --> LowT["选择较低温度<br/>T = 2-4"]
        Large["模型容量相近"] --> HighT["选择较高温度<br/>T = 4-10"]
        ManyClass["类别数多"] --> MidT["中等温度<br/>T = 3-6"]
        FewClass["类别数少"] --> HighT2["较高温度<br/>T = 5-10"]
    end

模型容量差距越大,温度应该越低。当学生模型比教师模型小很多时,它无法复制教师模型的完整知识。过高的温度会让软标签过于平滑,反而增加了学习难度。Hinton在论文中提到,当学生模型很小(每层只有30个神经元)时,$T=2.5$到$T=4$的效果最好。

类别数越多,温度可以越高。在类别数很多的任务中(比如ImageNet有1000个类),即使温度为1,概率分布也会相对平滑。此时可以适当降低温度。

典型范围是$T=2$到$T=10$。Hinton原始论文使用$T=2$到$T=20$,但后续研究表明,对于大多数任务,$T=3$到$T=5$是一个好的起点。

一个实用的调参策略是:

  1. 先固定$\alpha=0.5$,尝试$T \in \{2, 4, 6, 8\}$
  2. 找到最优温度后,微调$\alpha \in \{0.3, 0.5, 0.7, 0.9\}$

值得注意的是,温度参数和权重参数$\alpha$之间存在交互效应。高温度使得软标签更平滑,此时应该增加软标签损失的权重(降低$\alpha$)。

知识蒸馏与其他压缩技术的权衡

知识蒸馏不是唯一的模型压缩技术。它需要与量化(Quantization)和剪枝(Pruning)一起考虑:

graph TD
    subgraph "模型压缩技术对比"
        Q["量化<br/>精度换速度<br/>实现简单"]
        P["剪枝<br/>移除冗余<br/>需要微调"]
        D["蒸馏<br/>知识迁移<br/>需要教师"]
    end
    
    Q --> C["组合使用效果更佳"]
    P --> C
    D --> C
    
    C --> Best["推荐组合:<br/>蒸馏 + 量化"]

量化将模型权重从高精度浮点数转换为低精度整数(如INT8、INT4)。它的优点是实现简单,几乎所有框架都支持;缺点是精度损失不可避免。

剪枝移除模型中不重要的连接或神经元。它的优点是可以显著减少参数量;缺点是剪枝后的模型通常需要特殊的稀疏矩阵运算支持才能获得加速。

蒸馏通过教师模型指导学生模型学习。它的优点是可以迁移"知识"而非仅仅是压缩;缺点是需要一个预先训练好的教师模型。

Apple的研究发现一个有趣的结论:量化加蒸馏的效果优于剪枝。当三者组合使用时,蒸馏和量化的组合效果最好,其次是蒸馏和剪枝。

实际部署中,一个常见的策略是:

  1. 先用蒸馏训练一个较小的学生模型
  2. 对学生模型进行量化
  3. (可选)对量化后的模型进行轻量剪枝

蒸馏的局限性与陷阱

知识蒸馏并非万能药,它有明确的适用边界:

graph TB
    subgraph "知识蒸馏的局限"
        L1["容量差距问题<br/>学生太小无法承载知识"]
        L2["错误继承<br/>教师错误会被学生学习"]
        L3["架构差异<br/>不同架构需要特殊处理"]
        L4["数据依赖<br/>需要高质量训练数据"]
    end
    
    L1 --> S1["解决方案: 降低温度"]
    L2 --> S2["解决方案: 使用教师助手"]
    L3 --> S3["解决方案: 关系蒸馏"]
    L4 --> S4["解决方案: 数据增强"]

容量差距问题。当学生模型太小,无法承载教师模型的知识时,蒸馏效果会大打折扣。这时候降低温度可能会有帮助,因为它减少了软标签中的信息量。

教师模型的错误会被继承。如果教师模型对某些样本的预测是错误的,学生模型也会学习这些错误。一种缓解方法是使用"教师助手"(Teacher Assistant)——一个中等大小的模型作为中间层,过滤掉一些错误。

架构差异带来的挑战。当学生和教师的架构差异很大时(比如教师是Transformer,学生是CNN),特征蒸馏和关系蒸馏比响应蒸馏更有效。

训练数据的依赖。蒸馏需要教师模型在训练数据上生成软标签。如果训练数据质量差,蒸馏的效果也会受限。

写在最后

知识蒸馏的核心思想可以用一句话概括:不要教学生正确答案,要教学生如何思考

传统的监督学习像是把答案直接告诉学生,而知识蒸馏则像是让学生观察老师是如何分析问题的。老师对"错误答案"的分析——“这个选项虽然错,但比那个错得更少”——这些细微的判断才是真正宝贵的知识。

从Hinton 2015年的论文到今天DeepSeek R1的推理蒸馏,知识蒸馏已经走过了十年。它从一个模型压缩技术,演变成了知识迁移的通用框架。在模型规模不断膨胀的今天,这项技术的重要性只会越来越高。

对于实践者,记住几点:

  • 温度参数是核心,从$T=4$开始尝试
  • 损失函数记得乘以$T^2$
  • 模型差距大时考虑使用中间层蒸馏
  • 蒸馏加量化是高效的组合

知识蒸馏的本质,不是压缩,而是传承。它是大模型时代连接"象牙塔"与"落地应用"的桥梁。


参考资料

  1. Hinton, G., Vinyals, O., & Dean, J. (2015). Distilling the Knowledge in a Neural Network. arXiv:1503.02531
  2. Sanh, V., et al. (2019). DistilBERT: a distilled version of BERT. arXiv:1910.01108
  3. Jiao, X., et al. (2020). TinyBERT: Distilling BERT for Natural Language Understanding. Findings of EMNLP 2020
  4. DeepSeek-AI. (2025). DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. arXiv:2501.12948
  5. PyTorch Documentation. Knowledge Distillation Tutorial. https://docs.pytorch.org/tutorials/beginner/knowledge_distillation_tutorial.html
  6. Gou, J., et al. (2021). Knowledge Distillation: A Survey. International Journal of Computer Vision
  7. Apple Machine Learning Research. Combining Compressions for Multiplicative Size Scaling. https://machinelearning.apple.com/research/combining-compressions