knowledge distillation论文阅读之：Teacher-Class Network: A Neural Network Compression Mechanism

Abstract

• 本文提出了一种 teacher-class 网络结构，由一个 teacher 网络和多个 student 网络构成
• 本文的核心思路是将 teacher 网络的密集的特征表示分成多个部分，每个部分指导一个 student 网络学习，最终把这些训练好的 student 网络的特征表示进行合并。
• 相比于 KD 一直以来使用 logits 层的做法，本文方法使用 logits 层之前的网络学习的密集的特征表示；与使用 logits 层相比，本文提出的这种方法可以更好地解决各类问题，而不仅仅局限于分类问题。
• 和其他 single-student 的方法相比，本文方法有很好的表现。当我们达到接近 teacher 网络的表现的时候，参数和运算量只有 teacher 网络的 1/30 到1/10

Methodology

我们提出的方法利用了多名学生的优势，并将密集的表示知识从教师传递给几个学生。

Extracting dense representation from the teacher

• 神经网络在各种任务中通常会生成密集的特征表示 $d$，在分类任务中，这个特征表示被送入特定做分类任务的 logits 层 $z_i$，然后被送入 softmax 函数进行分类的概率输出，从而实现分类任务。这个时候，$z_i$ 会进行如下运算，然后产生对应每个种类的概率分布 $\hat{y}$：
  

  • $c$ 是全部的种类的数量
  • $T$ 是温度参数，用来控制标签的“软度”分布

• 通常来讲，teacher-student 模式的任务是最小化 teacher 网络产生的 soft target $\hat{y}$ 以及 student 网络产生的 soft target 之间的 cross-entropy 损失。

• 然而， 由于这些从 logits 层 $z_i$ 中通过 softmax 函数最终产生的概率分布 $\hat{y}$ 只能代表输入 inputs 的种类信息；学习这些 soft target 会限制 student 网络只能解决 “分类” 问题

• 众所周知，整个数据集通过神经网络学习的信息是非常丰富的，而且往往会被储存在密集的特征表示 $d$ 中；这个结论可以被一些典型的迁移学习场景所证实：移除 logit 层 $z_i$，仅利用密集特征层 $d$ 中的知识，通过转移相关任务[15]中的知识来学习新任务。

• 因此，我们将知识转移的目标重新定义为训练一个小的学生模型来模仿一个大的预训练教师网络的密集特征 $d$。换句话说，目标是使 教师网络的密集特征 向量与 学生网络产生的密集特征向量 之间的重构误差最小，如下图所示：

  

  • $m$ 是数据集中用于训练的 sample，
  • $d$ 和 $\hat{d}$ 分别是 teacher 网络 和 student 网络的密集特征
  • $d$ 是 teacher 网络在 logit 层之前的上一层的输出结果

• 一旦 student 网络学会了如何重构密集的特征表示，输出层（以分类任务举例，输出层是 logits 层和 softmax 层）就可以被引入来完成迁移学习中相关的任务。这个输出层可以简单地使用带有预训练权重的 teacher 网络的输出层。这句话的意思是：让 student 学会了密集的特征表示之后，直接可以嫁接在 teacher 网络的输出层之前，用 student 网络提取的密集特征 + teacher 网络的输出层来完成特定的任务。

Learning dense representation using n students

• 对于 teacher-student 模式的网络结构，尝试将全部的知识蒸馏到一个学生网络的做法会让简单的网络难以承受，多个 student 的网络可以被用来分担这种压力，并且模仿 teacher 网络。
• 一个之前的研究尝试使用集成的二元分类器，即：每个 logit 只包含一个特定类的分类信息，当面对1000类的分类情况时，这种方法要使用 1000 个 student
• 和这种方式形成鲜明对比，我们的方法可以将 teacher 网络中的密集特征表示 $d$ 划分成任意多个知识块，划分知识块的方法既可以通过均等的把 $d$ 分成 $n$ 个部分，也可以通过奇异值分解来完成
• 但是奇异值分解来划分 $d$ 的方法在数据集很大的情况下会非常麻烦，不容易实现，所以我们采用了第一种划分的方式。
• 把 $d$ 分为 $n$ 个部分， $d=[d_1,d_2,...,d_n{]}^{{}^{\prime }}$，所有的 $d_k$ 可以被重新生成 $d$，而且所有的 $d_k$ 会被分配给不同的 $studen{t}_k$
• 假设现有 $n$ 个 student 网络，$S_k^n$ 是第 $k$ 个 student，在数学上，将 teacher 网络的知识迁移到 $n$ 个 student 网络的过程可以被如下定义：
  
  • $\widehat{d_k}$ 表示第 $k$ 个 student 网络的知识蒸馏过程
  • 注意：每个 student 学习的都是 teacher 网络的一个部分（$d_k$）
• 每个 student 网络学习属于他自己的那一个 $d_k$ 部分，通过最小化均方误差的方式如下：
  
  • $L_{kd}^k$ 是第 $k$ 个 student 学习的 knowledge distillation loss
  • $m$ 是一个训练样本的总数
  • 下图显示了 $n$ 个学生如何通过最小化知识蒸馏损失($L_{k}d$)来学习密集表示的特征表示
    

Combining learned chunks of knowledge

• 经过 $n$ 个 student 分别学习 $d_k$，产生的 $n$ 个 $\widehat{d_k}$ 可以最终重新组合成 $\hat{d}$；这个过程被定义如下；
  

  • $[...{]}^{{}^{\prime }}$ 是聚合所有 student 网络训练结果的过程；
  • $X$ 是 input data，可以是文字或者图像

• 理想情况下，向量 $\hat{d}$ 应该尽可能的接近从 teacher 网络中产生的密集的特征表示 $d$ ；并且 $\hat{d}$ 应该可以用来解决 teacher 网络也能够解决的问题。

• 下图表示了分别从 teacher 网络和 4 个 student 网络使用 MNIST 和 Fashion MNIST 数据集训练得到的密集特征表示；并用 t-SNE 可视化出来
  

  • 可以看出，教师与学生网络的类别特定簇在学习空间中被很好地分割；这代表了 student 网络训练出的密集特征表示可以很好地模拟 teacher 网络的表现

• 既然 student 网络现在已经集体学习了 teacher 网络的密集特征表示，那么，当他们被送到teacher 网络的输出层（对于分类任务：输出层是 logits + softmax层）时，他们应该给出与教师网络相同的结果(如果解决了教师正在解决的相同问题)；因此，以分类问题为例，softmax 函数可以产生如下的 概率分布：
  

  • $g$ 代表了 teacher 网络的输出层（logits + softmax 层），他作用在所有 student 网络提取并聚合的密集特征表示 $\hat{d}$ 上
  • ${\theta }_g$ 代表输出层的权重参数，这是 teacher 网络预先训练好的权值

• 虽然 $d$ 和 $\hat{d}$ 只有很小的差距，但我们仍然希望修正这个误差，提高 student 网络的表现；方式是：保持 student 网络 nontrainable，然后微调输出层的参数。因此，在进行分类时，只有最后一层输出层可以采用交叉熵损失函数进行优化：
  

  • $L_{class}(y,\hat{y})={\sum }_{i=0}^N{y}_{i}log(\widehat{y_i})$ 是分类误差函数，
  • $\hat{y}$ 是通过 student 模型结合 teacher 的输出层给出的输出，$y$ 是 groundtruth 标签