Knowledge Distillation

《Distilling the Knowledge in a Neural Network》

几乎改善所有的机器学习算法最简单的方法就是在相同的数据集上使用不同的方法，最后将它们的预测值取平均。但是使用上述方法往往是很耗时耗力的，之前已经有工作证明可以将知识集压缩到单个的模型，从而更易去应用，这篇文章也是继续探究相关的模型压缩问题，从而提出了知识蒸馏（Knowledge distillation）

首先训练一个复杂的大型模型/教师模型，蒸馏就是将这个大型模型学到的知识通过特殊的训练转移到小型模型/学生模型上。

作者认为，学生模型应该学习教师模型的泛化能力，而非数据拟合能力；如果教师模型的泛化能力强，学生模型经过学习训练后，就能够在测试集上表现很好。

抽象地说，知识（Knowledge） 可以看作输入向量如何映射到输出向量

在一般的分类模型训练过程中，模型学习在正确的类别上得到最大的概率，但是不正确的分类上也会得到一些概率。一般而言，我们把它们视作学得不好的地方丢弃不看，但对这些非正确类别的预测概率也能反应模型的泛化能力

例如，一辆宝马车的图片，只有很小的概率被误识别成垃圾车，但是被识别成垃圾车的概率还是比错误识别成胡萝卜的概率高很多倍。这就是泛化能力

那如何利用这些信息呢？将教师模型的泛化能力转移到学生模型的一个方法是：使用教师模型产生的（每个类的）概率作为训练学生模型的"soft targets"

相较而言，如果数据集只提供正确的唯一选项，称之为"hard targets"

对于MNIST这种简单的数据，大型模型的准确率很高，错误标签的概率在 $10^3 \sim 10^6$ 量级。交叉熵损失影响很小，因此，在 softmax 中可以提高蒸馏温度 T，使这些比较小的 logit ，在经过 softmax 后，概率变大变平滑一些

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

在常见的大模型 chat 应用中，温度 T 我们已经看见过很多次了。当 T = 1 是就是常规的 Softmax

训练过程（假设这里选取的 T = 10）：

Teacher 模型：Softmax(T=10)的输出,生成“Soft targets” 。甚至这里的 transfer set 可以使用 unlabeled 数据

Student 模型：

1. 对 Softmax（T = 10）的输出与Teacher 模型的 Softmax（T = 10）的输出求 Loss1

2. 对 Softmax（T = 1）的输出与原始 label 求 Loss2

3. $Loss = Loss_1 + (1/T^2) Loss_2$

因为 soft targets 生成过程中蒸馏后的 softmax 求导会有一个 1/T^2 的系数，这里 hard targets 上面乘以一个系数是为了保持两个 Loss 所产生的影响接近一样（各 50%）

在 MNIST 上的实验表明，Soft targets 训练后可以提高学生模型的分类准确率，将大量知识转移到学生模型。即使在 transfer set 中完全去掉某一个类型（例如“3”的所有样本）学生模型也能分类对大部分样本

在语音识别实验中，由于模型集成效果优于单个模型。一般的思路是训练 10 个模型并集成。那么如何利用单个模型的成本开销，达到10个模型集成的效果？实验表明，蒸馏方法确实能够从训练集中提取更多有用的信息，而不是简单地使用硬标签来训练单个模型。通过使用 10 个模型的集成实现的分类精度的 80% 以上的改进被转移到提炼模型中

• 使用硬目标训练基线模型会导致严重的过拟合，而使用软目标训练的同一模型能够恢复完整训练集中的几乎所有信息，训练不必提前停止

训练专家集群

JFT 是一个 Google 数据集，其中包含 1 亿张带标签的图像和 15000 个标签。基线模型是一个深度卷积神经网络，它已经在大量内核上使用异步随机梯度下降进行了大约六个月的训练。

当类的数量非常大时，我们需要一个在所有数据上训练的通才模型和许多在非常容易混淆的类子集（如不同类型的蘑菇）上训练的 “专家” 模型的集成。

对于一个专家模型，它将不关心的所有标签统一分类为dustbin。每个 specialist 模型都使用 generalist 模型的权重进行初始化。然后，通过训练专家来略微修改这些权重，其中一半示例来自其特殊子集，另一半从训练集的其余部分随机采样。训练后，可以通过将dustbin class的 logit 增加 log(specialist class 被采样的比例) 来纠正有偏差的训练集

为了给专家分配职业，我们将 K-means 聚类算法应用于通才模型预测的协方差矩阵，将经常一起预测的一组类 $S^m$ 用作我们的专家模型之一 m 的目标

对于一个给定的图像 $\mathbf{x}$，分两步进行分类：

1. 对于每个测试样本，根据通用模型找到 n 个最可能的类，称这组类为 k

2. 然后取所有的满足以下的specialists m：其可混淆类的特殊子集 $S^m$ 与 k 有一个非空交集，并将其称为specialists 的active set $A_k$ 。学生模型给出的所有类的完整概率分布为 q，学生模型的loss为以下公式，需要最小化它：

$$KL(\mathbf{p}^g, \mathbf{q}) + \sum_{m \in A_k} KL(\mathbf{p}^m, \mathbf{q})$$

$\mathbf{p}^m \mathbf{p}^g$ 表示specialists或完整模型的概率分布，$\mathbf{p}^m$是 m 的所有specialist类别加上单个dustbin class的分布。优化后的分布 q 即为融合结果，取其中概率最高的类别作为最终预测

对于q，上面这个公式没有一般的封闭式解。还是需要 $\mathbf{q} =softmax(z)$，并用梯度下降来优化 logit z

结果：

• 经过训练的基线全网络开始，专家的训练速度非常快（需要几天而不是几周）

• 我们训练了 61 个专家模型，每个模型有 300 个类别（加上垃圾箱类别），提高了准确率

• 训练独立的专家模型非常容易并行化

• 专家接受数据在其特殊类别中高度丰富。这意味着其训练集的有效大小要小得多，并且它在其特殊类上具有强烈的过度拟合倾向。如果我们不用 dustbin，使用完整的类别softmax，也许会更好地防止过拟合。而且语音识别实验中也证实了：如果用通才的权重初始化专家，用硬目标和非特殊类的软目标训练它，它可以保留几乎所有关于非特殊类的知识。软目标可以由通才提供。作者目前正在探索这种方法。

使用在数据子集上训练的专家与 Mixtures of Experts (MoE) 有一些相似之处。后者使用门控网络来计算将每个样本分配给每个专家的概率。在专家学习处理分配给他们的示例的同时，门控网络正在学习根据专家对该示例的相对判别性表现来选择将每个示例分配给哪些专家。

使用专家的判别性能来确定学习到的分配比简单地对输入向量进行聚类并为每个聚类分配一个专家要好得多，但这使得训练难以并行化：

• 首先，每个专家的加权训练集以一种依赖于所有其他专家的方式不断变化；

• 其次，门控网络需要比较不同专家在同一示例上的性能，以知道如何修改其分配概率。这些困难意味着在包含明显不同子集的巨大数据集的任务中，MoE 很难而且很少使用

训练多个专家模型的并行化要容易得多。我们首先训练一个通才模型，然后使用混淆矩阵来定义专家模型所训练的子集。一旦这些子集被定义，专家模型就可以完全独立地进行训练。在测试时，我们可以使用通才模型的预测来确定哪些专家模型是相关的，并且只需要运行这些相关的专家模型。