[论文解读] Distilling Model Knowledge

George Papamakarios|arXiv (Cornell University)|Oct 8, 2015

Time Series Analysis and Forecasting参考文献 9被引用 9

一句话总结

本论文提出了一种通用的知识蒸馏框架，通过匹配输出和导数，使紧凑模型模仿复杂模型。该方法在模型压缩、贝叶斯预测分布蒸馏以及不可解生成模型蒸馏方面均取得了最先进性能，引入了导数匹配和在线蒸馏等新方法。

ABSTRACT

Top-performing machine learning systems, such as deep neural networks, large ensembles and complex probabilistic graphical models, can be expensive to store, slow to evaluate and hard to integrate into larger systems. Ideally, we would like to replace such cumbersome models with simpler models that perform equally well. In this thesis, we study knowledge distillation, the idea of extracting the knowledge contained in a complex model and injecting it into a more convenient model. We present a general framework for knowledge distillation, whereby a convenient model of our choosing learns how to mimic a complex model, by observing the latter's behaviour and being penalized whenever it fails to reproduce it. We develop our framework within the context of three distinct machine learning applications: (a) model compression, where we compress large discriminative models, such as ensembles of neural networks, into models of much smaller size; (b) compact predictive distributions for Bayesian inference, where we distil large bags of MCMC samples into compact predictive distributions in closed form; (c) intractable generative models, where we distil unnormalizable models such as RBMs into tractable models such as NADEs. We contribute to the state of the art with novel techniques and ideas. In model compression, we describe and implement derivative matching, which allows for better distillation when data is scarce. In compact predictive distributions, we introduce online distillation, which allows for significant savings in memory. Finally, in intractable generative models, we show how to use distilled models to robustly estimate intractable quantities of the original model, such as its intractable partition function.

研究动机与目标

• 解决部署大型复杂机器学习模型所面临的实际挑战，包括推理速度慢和存储成本高。
• 开发一种统一的知识蒸馏框架，使简单模型能够复现复杂模型的行为。
• 在三个关键应用中提升模型效率：模型压缩、紧凑贝叶斯预测分布，以及不可解生成模型的蒸馏。
• 通过新型蒸馏技术，在数据稀缺条件下实现更优性能并减少内存使用。
• 利用蒸馏模型对不可解量（如分区函数）进行稳健估计。

提出的方法

• 通过最小化惩罚输出值和导数值偏差的损失函数，训练紧凑的学生模型以模仿复杂的教师模型。
• 在模型压缩中使用导数匹配，以在训练数据有限时提升泛化能力。
• 在贝叶斯推断中实现在线蒸馏，通过维护充分统计量的运行平均值来减少内存使用。
• 通过最小化KL散度，利用随机梯度训练将不可归一化模型（如RBMs）蒸馏为可处理模型（如NADEs）。
• 利用R{backprop}高效计算学生模型和教师模型中的二阶导数。
• 使用小批量和在线EM算法训练蒸馏过程中的混合模型，并采用自适应学习率调度以确保收敛。

实验结果

研究问题

• RQ1 知识蒸馏能否被推广为一种统一框架，适用于多种机器学习应用？
• RQ2 在数据稀缺时，导数匹配如何提升蒸馏性能？
• RQ3 在线蒸馏是否能显著降低贝叶斯预测分布学习中的内存使用？
• RQ4 蒸馏模型能否稳健估计复杂模型（如RBMs）的不可解量（如分区函数）？
• RQ5 在模型压缩、贝叶斯推断和不可解生成建模中，蒸馏的相对性能提升如何？

主要发现

• 导数匹配显著提升了模型压缩中的蒸馏性能，尤其在数据稀缺时，通过同时对齐函数值和梯度实现。
• 在线蒸馏通过维护充分统计量的运行平均值而非存储完整数据批次，显著降低了内存消耗。
• 蒸馏后的NADE模型成功以高精度估计了原始RBM的不可解分区函数。
• 该框架能够有效将大型神经网络集成蒸馏为紧凑的单模型等效结构，且预测性能相当。
• 结合自适应学习率的随机梯度训练确保了在线蒸馏中稳定且收敛的优化过程。
• 所提出的蒸馏框架在所有三项应用中均取得了最先进性能：模型压缩、贝叶斯推断和不可解生成建模。