神经网络模型压缩之Knowledge Distillation

神经网络模型压缩之Knowledge Distillation

1. Background

随着深度学习理论的普及和PC计算能力的提高, Neural Network也日益趋于复杂化–越来越深和越来越大, 比如VGG系列,深度均在10层以上, 残差网络的深度更是从数十到一百以上. 复杂的网络模型带来的是超强的规则学习能力, 虽然以现在动辄几十或上百多机GPU并行的计算能力而言, 这样的巨大模型并不是拦路虎, 但是考虑到当今微小智能设备的普及, smart phone, 可穿戴设备, 智能手表等, 要在这些嵌入式设备中放入动辄几十上百M的模型是不现实的. 就算能放下, 其高昂的computation cost也是难以承受的. 那么有没有办法既保证模型能力又能大幅削减cost呢? 答案是肯定的. 在深度学习领域有一个专门的小分支研究model compression, 即模型压缩. 模型压缩方法可以大致分为4类: 参数共享(Parameter sharing), 网络裁剪(Network pruning), 矩阵分解(Matrix decomposition), 以及我们在这里要说的知识蒸馏(Knowledge Distillation), 也称为Dark knowledge.

在讲解Knowledge Distillation之前,我们先来讲讲NN的拟合能力问题. 我们知道, 现在的深度模型都是由一层层的网络叠加而成, 而每层网络又由多个神经元组成, 这些神经元中通常包含一个非线性的激活函数(sigmoid, Relu, tanh等), 这些激活函数的输出值通过每层的权重矩阵W进行加权组合, 成为下一层神经元的输入. 而在多年前早有大牛在数学上证明了一个具有有限个神经元的单一隐藏层的BP网络就可以拟合任意目标函数(Universal approximation theorem, 简单数据拟合演示Demo), 那么我们的NN真的需要这么大吗? 参考文献[4]

中的CIFAR-10实验, 通过pre-train和teach-student的训练方法, 作者仅仅使用了2.5M的参数就达到了接近state-of-the-art的效果, 而使用的网络参数数目比传统方法采用的模型小得多. 那么问题就来了,为什么在实践中直接使用original的数据中使用很难训练出一个效果差不多的轻量级模型? 原因可能在于:大模型具有更好的Inductive bias, 使之能够学习到更好的层次结构; 现有的regularization方法更适用于深层模型([6]). 那么我们如何在现有的外部条件下训练出一个小的高精度模型呢? Knowledge Distillation就提供了一条很好的思路—-使用大模型(Teacher)来teaching小模型(Student). 有人可能会问了, 就算大模型精度再高, predict acc也不可能达到100%的准确度(至少是非常难的), 为啥不直接使用更为准确的training dataset? 这里有几个原因:

1.Teacher的输出往往带有一些额外的信息, 即Dark Knowledge. 手写体识别的classifier为例, 该classifier的作用是识别0-9的手写体数字, 其最后输出的是一个样本属于10个数字的distribution.

比如样本x的输出为[1e-10, 1e-10, 0.98, 0.01, 1e-10, 0.009, 1e-10, 1e-10, 1e-10, 1e-10] (依次对应数字0-9),

根据输出distribution, 最大概率为数字2, 故模型对x的预测标签为2. 需要注意的是, 3和5对应的概率虽然小但远大于除2以外的数字概率, 这说明模型认为x虽然代表数字2, 但它也有点像数字2和5. 这种附加的信息其实反映出classifier在学习到了某种规律(函数), 通过该规律来计算x与不同数字的相似度. 相比original training data中生硬的one-hot标签, 这种具有额外信息的soft label携带更多的数据结构信息.

2.相对于one-hot label, 由teacher提供的soft label更容易学习, 相比于输出概率集中在一个class上的one-hot, teacher模型将概率分散在了多个class上, 这样的好处是, 使模型的预测更加的贴近实际的情况(比如在第1点中提到的)

3.通常在人工标注的训练数据中都存在错误标记, 而训练良好的Teacher模型的predict结果可以很好的”纠正”这些错误.

2. Algorithm

讲了这么多, 该怎么做呢? 在这一节,我们就来看看如何训练. 关于模型的训练, 读者最先想到的应该是直接拟合Teacher的输出, 即把这个学习问题当成一个regression, 使用平方损失函数:
L=12T∑Ti=1[ft(x(i))−fs(x(i);W)]2 L = 1 2 T ∑ i = 1 T [ f t ( x ( i ) ) − f s ( x ( i ) ; W ) ] 2
看到这里, 大家可能会有个疑问, 这个f是模型softmax层输出的distribution值吗? 显然, 直接使用这个值会有一个问题, 因为Teacher模型毕竟是根据one-hot label进行训练的, 它的输出自然会很接近one-hot, 是一个非常”sharp”的分布, 即在predict的class上概率值高度集中(通常是0.99+), 其它class上值很小, 如果直接用来计算loss, dark Knowledge会被模型忽略. 那怎么办呢? 将大尺度的值进行缩放最常用的办法就是对数运算, 因此, 直接取隐藏层的logit输出 z(i) z ( i ) 是一个不错的选择:
L=12T∑Ti=1[zt(x(i))−fs(x(i);W)]2 L = 1 2 T ∑ i = 1 T [ z t ( x ( i ) ) − f s ( x ( i ) ; W ) ] 2
当然了, 对于loss函数,我们也可以用 KL(pt||ps) K L ( p t | | p s ) , 但是根据文献[2], 这样做的效果不如前者, 原因很可能也是因为其输出distribution的sharp问题. 那么有没有办法改进呢? 在文献[3]中在softmax中加入了一个额外的temperature参数 τ≥1 τ ≥ 1 :
pτ(x)=softmaxτ(x)=ezi(x)/τ∑jezj(x)/τ p τ ( x ) = s o f t m a x τ ( x ) = e z i ( x ) / τ ∑ j e z j ( x ) / τ
参数 τ τ 越大,输出的distribution就越分散平滑, 否则就越sharp, 当为1时就退化为普通的softmax. 另外, 为了引入原始的标签信息, 在[2]中对loss function进行了优化:
L=KL(ptrue||ps)+λKL(pτt||pτs) L = K L ( p t r u e | | p s ) + λ K L ( p t τ | | p s τ )
其中KL为cross-entropy loss function, λ λ 是可调整的参数, 式中第一部分的 pt p t 通过普通的softmax计算得出(即 τ=1 τ = 1 ), 在[2]中也提到了使用引入对原始标签的loss能够极大的提升distilled模型的性能.

最后一个问题就是, 对于student模型, 原则上来说在广度和深度上都要比Teacher更小. 熟悉神经网络的读者都知道, 网络的深度对模型建模能力的影响比广度要大, 在同样的参数数目下, 越深的网络拥有更好的拟合能力. 因此,我们的student模型可以有较大的深度, 同时为了加速训练速度, 可以引入hint, 即使用Teacher的若干层对student进行pretraining, 详细可以参考文献[4], 除此之外,为了增加模型的泛化能力, 还可以对Teacher模型的输出logit进行加噪(文献[5])

最后, 附上chengshengchan的开源代码, 该代码实现了文献[3]的Distillation以及[5]中的加噪:
https://github.com/chengshengchan/model_compression

---

参考文献:

[1]:C. Bucilu, R. Caruana, and A. Niculescu-Mizil. Model compression. In Proceedings of the 12th ACM SIGKDD international conference on Knowledge discovery and data mining,pages 535–541. ACM, 2006.

[2]:J. Ba and R. Caruana. Do deep nets really need to be deep? In Advances in neural information processing systems, pages 2654–2662, 2014.

[3]:G. Hinton, O. Vinyals, and J. Dean. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531, 2015.

[4]:A. Romero, N. Ballas, S. E. Kahou, A. Chassang, C. Gatta, and Y. Bengio. Fitnets: Hints for thin deep nets. arXiv preprint arXiv:1412.6550, 2014.

[5]:Deep Model Compression-Distilling Knowledge from Noisy Teachers

[6]:Dumitru Erhan, Yoshua Bengio, Aaron Courville, Pierre-Antoine Manzagol, Pascal Vincent, and Samy Bengio. Why does unsupervised pre-training help deep learning? The Journal of Machine Learning Research, 11:625–660, 2010.

---