CKD知识蒸馏实战：如何在小模型上实现大模型的性能

1次阅读

共计 1838 个字符，预计需要花费 5 分钟才能阅读完成。

在移动端和嵌入式设备上部署深度学习模型时，我们经常面临一个矛盾：大模型性能好但资源消耗高，小模型速度快但精度不足。传统知识蒸馏（Knowledge Distillation, KD）通过让学生模型模仿教师模型的输出分布来缓解这个问题，但它存在两个明显局限：

仅利用最终输出层信息，忽略了中间层特征（feature map）的引导作用
单向知识传递导致学生模型容易过度拟合教师模型的偏好

协作式知识蒸馏（Collaborative Knowledge Distillation, CKD）通过三种机制改进传统方法：

对比损失设计 ：在特征层引入正负样本对比学习，公式如下：
```
L_cont = -log[exp(sim(f_s, f_t)/τ) / (exp(sim(f_s, f_t)/τ) + ∑exp(sim(f_s, f_n)/τ))]
```
其中 f_s/f_t 分别代表学生 / 教师模型特征，f_n 是负样本特征，τ 为温度系数
梯度协同机制 ：教师模型也会根据学生模型的反馈调整参数，形成双向知识流动

方法	CIFAR-10 精度	参数量 (M)	推理时延 (ms)
Teacher	94.2%	23.5	15.3
KD	91.1%	5.8	3.2
FitNets	92.4%	5.8	3.2
CKD(ours)	93.7%	5.8	3.2

class CKDLoss(nn.Module):
    def __init__(self, temp=4.0):
        super().__init__()
        self.temp = temp
        self.kl_div = nn.KLDivLoss(reduction='batchmean')

    def forward(self, student_out, teacher_out, student_feat, teacher_feat):
        # 传统蒸馏损失
        soft_loss = self.kl_div(F.log_softmax(student_out/self.temp, dim=1),
            F.softmax(teacher_out/self.temp, dim=1)
        ) * (self.temp ** 2)

        # 特征对比损失
        norm_s = F.normalize(student_feat, p=2, dim=1)
        norm_t = F.normalize(teacher_feat, p=2, dim=1)
        sim_matrix = torch.mm(norm_s, norm_t.T) / self.temp
        cont_loss = -torch.diag(F.log_softmax(sim_matrix, dim=1)).mean()

        return soft_loss + 0.5 * cont_loss  # 可调整权重

for epoch in range(epochs):
    for x, _ in train_loader:
        # 双模型前向传播
        t_out, t_feat = teacher(x, return_feats=True)
        s_out, s_feat = student(x, return_feats=True)

        # 计算 CKD 损失
        loss = criterion(s_out, t_out, s_feat, t_feat)

        # 反向传播（注意设置 teacher.requires_grad_()）optimizer.zero_grad()
        loss.backward()
        optimizer.step()