首页 > 技术文章 > 模型压缩,反其道而行之?

dagis 2020-12-23 04:45 原文

矛盾存在于一切事物中,并且贯穿于事物发展过程的始终,即矛盾无处不在,矛盾无时不有。机器学习领域也无例外。
 
近几年大型网络不断刷新记录,看看这些SOTA(state of the art)模型规模:
BERT-Base: 110m
GPT-3: 175b
VGG16/19: 138m
ResNet: 23m
 
谷歌gmail的自动补齐功能,为了控制预测时延,后台的预测模型是跑在TPU上的。
https://ai.googleblog.com/2018/05/smart-compose-using-neural-networks-to.html
Roblox有一篇blog分享了,如何在cpu上用BERT做归类,实现一天1b的请求,其中包含了bert distill,模型量化,缓存等等工程上的优化,非常值得推荐。https://blog.roblox.com/2020/05/scaled-bert-serve-1-billion-daily-requests-cpus/
 
不是所有的公司都像谷歌一样是大户,动辄上TPU。当然很多公司是由于系统架构的限制,无法简单的引入GPU或TPU。更多的是机器学习在处理器能力有限的嵌入式设备上的普及。至2018底,已有22b的IoT设备接入网络,预测至2030年底,将有50b的设备接入。所以如何在有限资源的设备上跑时尚机器学习模型自然成了热点之一。其中如何精简压缩模型在大型网络时代也暗流涌动。
 
论文<Distilling the Knowledge in a Neural Network>的开头一段讲的非常好。许多昆虫的幼虫形式,通过优化,可从环境中提取能量和养分,而成虫则完全不同,其成虫形式针对不同的旅行和繁殖需求进行了优化。机器学习模型也一样,在训练的时候,可能考虑的是模型预测准确度,而在预测的时候,考虑更多的是如何做到更快更节省资源。
 
而在评判一个模型性能的时候,准确度也已经慢慢不再是唯一的判断标准。除了正在热议的道德问题以外,模型是不是低碳环保也正在讨论作为评判标准的一部分。如果这个被通过,那GPT-3训练一次4.6m美刀对总体评分肯定伤害不少。
 
嗯嗯,那我们赶紧聊聊怎么压缩模型吧。
 
压缩模型可以分为两大类,一是模型修剪,二是知识蒸馏。
 
模型修剪比较好理解,从深度上来说,可以减少神经网络的层数,从广度上来说,可以减少隐藏层的节点数。这个方法必须重新训练模型,当然无法解决我们所说的降低训练成本的问题。当然修剪后的模型对于后期预测速度的提高,对于大部分工程上的案例是有实际意义的。
修剪模型另外一种方式是模型量化,即降低精度,比如从32位降到8位。这个方法不需要重新训练模型。但同修剪节点方法一样,不具有普遍性。可能在某些特定的情况下能够保持性能,但很多情况下,不可避免的牺牲了模型的性能。
 
如今更加popular的做法是知识蒸馏。当然知识蒸馏也无法逃离重新训练模型。
 
第一次接触知识蒸馏还是team在讨论用BERT的时候,一个同事接了这个BERT Distill的活,研究如何从以BERT vanilla 为基础的DistillBert作为我们网络模型。(注: vanilla在软件业指的是软件的default配置,BERT有很多版本,我们这里用了default的)。知识蒸馏可以理解成,一个老师在上课,老师可能会无重点的把所有知识都讲一遍,而学生呢,则会把自己认为的重点都记下来。模型蒸馏中,有一个或多个的原始模型作为老师模型,一个目标模型作为学生模型。你可能会想到meta-learning。是的,meta-learning可以说是知识蒸馏的一个例子。meta-learning网络,能够引导学生模型目标到合适的网络。
 
那实际上促使我们做知识蒸馏的intuition是什么呢?
 
当一个model训练出来后,一般我们认为网络参数就是我们的产物。但其实不然,真正的价值产物是如何把输入矢量转换成输出矢量。我们知道,能实现同等的转换的网络并不是唯一的,我们完全可以寻找一个更加有效的网络能够做同样的转化。因此,如果我们能找到一个精简的网络可以达到同样的转换,那可以说我们这个精简网络就是原来网络的蒸馏。下图就是一个简化的框图,表示了知识蒸馏的原理。(图片来自网络)
 
 
 
从上图上我们可以知道,寻找合适的loss函数进行训练学生模型是个很重要的环节。在归类问题当中,有一个误区,就是用最后的类属来作为校准的对象。但是实际上原有网络,在得到最后类属的之前有很多的信息含量,比说在在用softmax得到最后类属前的概率值。举例来说,多项归类问题,如果一个猫没有被归类成猫,那比起汽车来说,它可能更倾向于被归类成狗。但如果用最后的类属作为训练校准对象,这个信息将被丢失,那学生模型在性能方面将无法达到原始模型的效果。为了更好的蒸馏出原有网络的知识,在计算这个loss的时候,引入了温度参数。对于不同类型的网络蒸馏,可以调整温度参数。
 
以上提到的知识蒸馏只是个引子。知识蒸馏还有很多的方法。天下没有免费的午餐,和大部分其他的研究方向一样,每种方法都有其优缺点。但大家都面临的同样问题,理论上对其有效性的解释和证明。
 
最后推荐两篇论文。
文中提到的<Distilling the Knowledge in a Neural Network>,不仅在知识蒸馏方面解释的非常到位,而且对于机器学习方法的总结也是极其点睛到位。实在是高手之作。
<Knowledge Distillation: A Survey> 总结了模型压缩和知识蒸馏的方方面面。如果要作此方面的研究话,这就像一个catalog一样,可以省了不少搜索功夫。
 
其它参考文章:
https://towardsdatascience.com/distilling-knowledge-in-neural-network-d8991faa2cdc
https://towardsdatascience.com/how-to-compress-a-neural-network-427e8dddcc34
https://medium.com/neuralmachine/knowledge-distillation-dc241d7c2322
 
更多更新作者原创文章,关注作者微信公众号:

推荐阅读