介绍

        这一篇论文是知识蒸馏中最经典的论文。论文提出了“蒸馏”的模型训练方式,将知识从大模型转移到更适合部署的小模型。通过使用带温度的softmax函数来软化教师网络的逻辑层输出作为学生网络的监督信息。

        将大模型的泛化能力转移到小型模型的一个明显方法是使用大模型产生的类概率作为训练小模型的“软目标”。对于这个转移阶段,我们可以使用相同的训练集或单独的“转移”集。当大模型是由更简单的模型组成的大型集合时,我们可以使用其单个预测分布的算术或几何平均值作为软目标。当软目标具有高熵时,它们为每个训练案例提供的信息比硬目标多得多,训练案例之间的梯度差异也小得多,因此小模型通常可以在比大模型少得多的数据上训练,并使用更高的学习率。

        对于像MNIST这样的任务,大模型几乎总是以非常高的信心产生正确的答案,关于学习函数的许多信息都存在于软目标中非常小的概率比率中。例如,一个版本的2可能被赋予10-6的3和10-9的7概率,而另一个版本可能相反。这是宝贵的信息,定义了数据上丰富的相似性结构(即它说哪些2看起来像3,哪些看起来像7),但它对传输阶段的跨熵成本函数的影响很小,因为概率非常接近于零。Caruana和他的合作者通过使用logits(对最终softmax的输入)而不是softmax产生的概率作为学习小模型的目标来规避这个问题,他们最大限度地减少了大模型生成的日志与小模型生成的日志之间的平方差异。我们更通用的解决方案,称为“蒸馏”,是提高最终softmax层的温度,直到大的模型产生一组适当的软目标。然后,我们在训练小模型以匹配这些软目标时使用相同的高温。我们稍后表明,匹配繁琐模型的日志实际上是蒸馏的特殊情况。

        用于训练小模型的传输集可能完全由无标签数据[1]组成,或者我们可以使用原始训练集。我们发现,使用原始训练集效果很好,特别是如果我们在目标函数中添加一个小术语,鼓励小模型预测真实目标,并匹配繁琐模型提供的软目标。通常,小模型无法与软目标完全匹配,错误地指向正确答案是有帮助的。

实现方法

神经网络通常通过使用“softmax”输出层生成类概率,该输出层通过将为每个类计算的logit zi转换为概率qi,并将zi与其他logit进行比较。

【模型压缩1503.02531】Distilling the Knowledge in a Neural Network_神经网络

 其中T是通常设置为1的温度。对T使用更高的值会产生更软的类概率分布。

        在最简单的蒸馏形式中,知识通过在转移集上训练知识,并为转移集中的每个情况使用软目标分布,该分布是通过使用高温在软最大值的繁琐模型产生的。训练蒸馏模型时使用相同的高温,但训练后使用1的温度。

        当已知所有或部分传输集的正确标签时,也可以通过训练蒸馏模型生成正确的标签来显著改进这种方法。一种方法是使用正确的标签来修改软目标,但我们发现更好的方法是简单地使用两个不同目标函数的加权平均值。第一个目标函数是与软目标的交叉熵,这种交叉熵在蒸馏模型的软最大值中使用与从繁琐模型生成软目标相同的高温计算。第二个目标函数是带有正确标签的交叉熵。这是在蒸馏模型的软最大值中使用完全相同的日志计算的,但温度为1。我们发现,通常通过对第二个目标函数使用明显较低的权重来获得最佳结果。由于软目标产生的梯度大小缩放为1/T 2,因此在同时使用硬目标和软目标时,必须将其乘以T 2。这确保了在尝试元参数时改变蒸馏时使用的温度,硬目标和软目标的相对贡献大致保持不变。

【模型压缩1503.02531】Distilling the Knowledge in a Neural Network_机器学习_02

对学生网络来说,一部分监督信息来自hard label标签,另一部分来自教师网络提供的soft label。

代码实现

class DistillKL(nn.Module):
def __init__(self, T):
super(DistillKL, self).__init__()
self.T = T

def forward(self, y_s, y_t):
p_s = F.log_softmax(y_s/self.T, dim=1)
p_t = F.softmax(y_t/self.T, dim=1)
loss = F.kl_div(p_s, p_t, size_average=False) * (self.T**2) / y_s.shape[0]
return loss