不同于大多数的蒸馏模型中用小模型去拟合大模型的输出,这篇文章将知识定义为神经网络中层与层之间的关系(特征图的内积)
文中提到这样做有三个优点:
- student 网络可以学习的更快
- student 网络可以比 teacher 网络的性能更优
- 适用于迁移学习(teacher 和 student 属于不同的任务
知识的定义
将神经网络的输入和输出看作是问题和答案,之前介绍的 Romero 等人的文章指出让学生模型去拟合教师模型中间层的输出。但是我们知道解决问题的方法可能不止一种,比如图中的最短路径可能有多个解,因此这篇作者就指出让学生模型去学习这种解决问题的方法。这种方法在网络中表现为层与层之间的关系。
数学表示
$$
{G_{i,j}(x;W)=\displaystyle \sum_{s=1}^{h} \sum_{t=1}^{w}\frac{F_{s,t,i}^1(x;W)\times F_{s,t,j}^2(x;W)}{h\times w}}
$$
假设两层特征图的大小分别为 $h\times w \times m$ 和 $h\times w\times n$,那么最后得到的 FSP (flow of the solution process) 矩阵就是两个特征图中的各个通道彼此做点积。我们看这个公式,公式隐含的要求了两个特征图单个通道的大小必须是一样的,而这个性质和残差结构非常的像。事实上,作者用到的网络结构也正是带残差块的。此外这个例子中 FSP 矩阵的大小为 m*n,也就是说 FSP 矩阵的大小仅于通道数有关,而与每个通道的大小无关,这就解决了 Student 网络得到的 FSP 矩阵和 Teacher 的不一致的问题。
损失函数
Student 模型就是在对每对 FSP 矩阵做 L2 损失的基础上进行加权和:
$$
L_{FSP}(W_t,W_s)=\frac {1}{N}\displaystyle \sum_{x}\sum_{i=1}^{n}\lambda_i\times |G_i^T(x;W_t)-G_i^S(x,W_s)|_2
$$
训练流程
与 Romero 等人一致,作者也使用了两阶段的训练方式,完整的算法流程如下:
首先我们用上面所提到的 FSP 损失来训练 Student 网络使其学习 Teacher 模型层与层之间的关系,之后将这个阶段所 Student 学得的参数作为下一阶段的初始化参数。下一个阶段则和普通训练一致,以分类问题为例,损失函数可以选择 Softmax。
第二个阶段模型所做的任务可以和 Teacher 模型不同,这就可以用于迁移学习。此外,在这个初始化参数下,作者用实验证明了模型可以快速的进行优化且达到比教师网络更好的性能。