社区计算机视觉课程文档

基于 Vision Transformer 的知识蒸馏

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

基于 Vision Transformer 的知识蒸馏

我们将学习知识蒸馏,这是 distilGPTdistilbert 背后的方法,这两个模型是 Hugging Face Hub 上下载次数最多的模型!

据推测,我们都有过这样的老师,他们只是简单地提供正确答案,然后测试我们以前从未见过的问题,这类似于机器学习模型的监督学习,我们提供带标签的数据集进行训练。然而,我们可以不让模型在标签上训练,而是追求 知识蒸馏 作为一种替代方案,以获得一个更小的模型,其性能可以与更大的模型相媲美,并且速度更快。

知识蒸馏背后的直觉

想象一下,你被给予了这道多项选择题

Multiple Choice Question

如果有人只是告诉你,“答案是德拉科·马尔福,”这并没有教会你很多关于每个角色与哈利·波特的相对关系。

另一方面,如果有人告诉你,“我非常确信不是罗恩·韦斯莱,我有点确信不是纳威·隆巴顿,我非常确信德拉科·马尔福,”这会给你一些关于这些角色与哈利·波特关系的信息!这正是根据知识蒸馏范式传递给我们学生模型的信息类型。

在神经网络中提炼知识

在论文 Distilling the Knowledge in a Neural Network 中,Hinton 等人介绍了被称为知识蒸馏的训练方法,其灵感来自于昆虫,尤其是。正如昆虫从幼虫期过渡到为不同任务优化的成虫期一样,大规模机器学习模型最初可能很笨重,就像幼虫一样,用于从数据中提取结构,但可以将它们的知识提炼成更小、更高效的模型以进行部署。

知识蒸馏的本质是使用来自教师网络的预测 logits 将信息传递给更小、更高效的学生模型。我们通过重写损失函数来做到这一点,使其包含蒸馏损失,这鼓励学生模型在输出空间上的分布近似于教师模型的分布。

蒸馏损失的公式如下

Distillation Loss

KL 损失是指教师和学生输出分布之间的 Kullback-Leibler 散度。然后,学生模型的总体损失被公式化为这种蒸馏损失与ground-truth 标签上的标准交叉熵损失之和。

要查看此损失函数在 Python 中的实现以及 Python 中完全完成的示例,请查看本节的 notebook

Open In Colab

利用知识蒸馏进行边缘设备部署

随着 AI 模型部署在边缘设备上,知识蒸馏变得越来越重要。部署大型模型,例如大小为 1 GB 且延迟为 1 秒的模型,对于实时应用程序是不切实际的,因为计算和内存要求很高。这些限制主要归因于模型的大小。因此,该领域已接受知识蒸馏,这是一种将模型参数减少 90% 以上而性能下降最小的技术。

知识蒸馏的后果(好与坏)

1. 熵增益

在信息论的背景下,熵类似于物理学中的对应物,它衡量系统内的“混沌”或无序程度。在我们的场景中,它量化了分布包含的信息量。考虑以下示例

  • 哪个更难记住:[0, 1, 0, 0] 还是 [0.2, 0.5, 0.2, 0.1]

第一个向量 [0, 1, 0, 0] 更容易记住和压缩,因为它包含的信息较少。这可以表示为第二个位置的“1”。另一方面,[0.2, 0.5, 0.2, 0.1] 包含更多信息。在此基础上,假设,例如,我们在 ImageNet 上训练了一个 80M 参数网络,然后将其(如前所述)蒸馏成一个 5M 参数的学生模型。我们会发现教师模型输出中包含的熵远低于学生模型的熵。这意味着即使学生模型的输出是正确的,它也比教师的输出更混乱。这归结为一个简单的事实:教师的额外参数帮助它更容易区分类别,因为它提取了更多特征。这种对知识蒸馏的看法非常有趣,并且正在积极研究以减少学生的熵,无论是通过将其用作损失函数还是通过应用受物理学启发的类似指标(例如能量)。

2. 一致的梯度更新

模型通过最小化损失函数并通过梯度下降更新其参数来迭代学习。考虑一组参数 P = {w1, w2, w3, ..., wn},它们在教师模型中的作用是在检测到 A 类样本时激活。如果一个模棱两可的样本类似于 A 类但属于 B 类,则模型在错误分类后会进行激进的梯度更新,从而导致不稳定。相反,蒸馏过程,使用教师模型的软目标,在训练期间促进更稳定和一致的梯度更新,从而为学生模型带来更平滑的学习过程。

3. 在未标记数据上训练的能力

教师模型的存在使得学生模型可以在未标记数据上进行训练。教师模型可以为这些未标记样本生成伪标签,学生模型然后可以使用这些伪标签进行训练。这种方法显着增加了可用的训练数据量。

4. 视角的转变

深度学习模型通常在这样的假设下进行训练:提供足够的数据将使它们能够逼近函数 F,该函数准确地表示了潜在现象。然而,在许多情况下,数据稀缺使得这种假设不切实际。传统方法包括构建更大的模型并迭代微调它们以实现最佳结果。相比之下,知识蒸馏转变了这种视角:鉴于我们已经有了一个训练有素的教师模型 F,目标变为使用更小的模型 f 来逼近 F

< > 在 GitHub 上更新