资讯 人工智能
此为临时链接,仅用于文章预览,将在时失效

Hinton 新作「在线蒸馏」,提升深度学习分布式训练表现的利器

作者:MrBear 编辑:杨晓凡
2018/04/24 19:11

雷锋网 AI 科技评论按:Hinton 团队最近发表了一篇关于「在线蒸馏」的新论文。论文里讲了什么呢?

我们为什么需要在线蒸馏?

近年来,随着深度学习技术的蓬勃发展,我们日常生活的每个角落都充满了「人工智能」的影子,并由此催生了许多新的业态、以人工智能应用为核心产品的初创公司也如雨后春笋般在科技行业中崭露头角。是的,这也许是人工智能(特别是深度学习)最好的时代。

然而,深度学习作为当今人工智能领域皇冠上最璀璨的明珠,如果要想将其推向工业级的应用,则往往需要可望而不可及的强大算力!而这种算力则是由价格高昂的硬件、庞大而复杂的分布式计算环境、精妙高效的算法作为支撑的。可想而知,除了 Google、Amazon、阿里巴巴、百度等掌握强大计算资源的科技巨头外,想获得这样的计算资源对于小型企业或者个人用户可谓是天方夜谭。实际上,在很多普通的高等院校和科研单位中,拥有像 Titan X 或 GTX 1080 Ti 这样的民用显卡已经是很奢侈的事情。更重要的是,由于基础架构的限制(分布式计算节点的通信、同步、任务调度等问题)、对模型进行优化求解的障碍、集成学习环境下各模型的贡献缺乏好的决策,使得盲目的堆砌计算资源也可能触碰到模型性能的天花板(例如:分布式随机梯度下降(distrtibuted SGD))。

为此,「深度学习泰斗」Geoffrey E.Hinton 近年来在知识蒸馏(或者「暗知识提取」)方面做了一系列工作,试图通过这种从高计算能力要求、同时也具有高表现的模型中提取出一些隐含的知识,并且将其作为先验,通过设计新的网络结构和目标函数将这种知识「传授」给规模较小的深度学习网络,实现对网络模型的压缩,以显著缩小的网络规模和计算需求展现出尽量高的模型表现。最近,为了将这种思想部署在分布式环境中,用以突破如今常常被使用的分布式 SGD(同步和异步形式)的瓶颈,Hinton 团队又发布了名为LARGE SCALE DISTRIBUTED NEURAL NETWORK TRAINING THROUGH ONLINE DISTILLATION」(通过在线蒸馏的神经网络大规模分布式训练)的论文。在笔者看来,Hinton 的这一系列工作进一步降低了深度学习模型应用的门槛,之前的研究对深度学习模型在移动终端上的部署、这篇论文对提高大规模计算集群上的模型表现和计算效率都有深远意义。

Hinton 新作「在线蒸馏」,提升深度学习分布式训练表现的利器

论文地址:https://arxiv.org/abs/1804.03235 

知识蒸馏的前世今生

要想理解「在线蒸馏」的概念,我们有必要回顾一下 Hinton 从 2014 年开始对 dark knowledge extraction(暗知识提取) 和 knowledge distillation(知识蒸馏)的相关工作,甚至更早的 Caruana et.al 所做的模型压缩的工作。

为了提升神经网络模型的性能,Caruana 等人早在 2006 年(那时深度学习还没有大火)就提出了一种压缩大规模复杂网络的方法。由于集成学习在传统的机器学习领域大获成功,许多深度学习研究人员很自然地想要利用集成学习的思想,将大量的模型聚合到一个神经网络中,通过暴力的训练,为不同的模型赋予不同的权值,对这些模型的输出进行加权平均得到最终的结果,以便充分利用它们各自对于不同的任务所具有的优势。然而,这种暴力的训练过程和臃肿的网络结构需要消耗巨大的计算资源、造成额外的能源消耗。Caruana 等人提出了 MUNGE 的数据增强算法,将大规模模型学习到的函数压缩进规模更小、训练更便捷的模型中。 

受此启发,时隔 8 年后,当深度学习迎来春天、人们沉醉于大规模深度学习网络带来的人工智能在计算机视觉、自然语言处理等方面的成功应用的时候,宗师 Hinton 则认为,是时候回过头来看看 Caruana 的文章,思考如何将这些「丑陋的」大规模的模型压缩到小而快的模型中去。

由此,Hinton 提出了「dark knowledge」的概念。在他看来,这种被称为「暗知识」的东西才是深度学习本质上学到的知识(或许这也是寻求深度学习可解释性的一种途径)。Dark knowledge,顾名思义,就是隐藏在深度学习表面上所展现出来的网络结构、节点之间的连接权重、网络的输出这些看得到的数据之下的知识。如果能够找到一种途径,使得我们能够获取这种知识,并且将其包装成一种先验概率,迁移到更小的模型中去,是否能够提升小模型的性能呢?事实上,笔者认为,这也可以看作是在迁移学习的框架下,将大规模网络视作信息充分的 source domain,将小规模网络视作需要大规模网络补充信息的 target domain,而我们提取到的 dark knowledge 则是两个任务之间的 common knowledge。

然而,该从哪里下手,获取这种知识呢?Hinton 敏锐地观察到:我们在绝大多数的预测任务的深度学习网络中,都会使用 softmax layer 为大量的标签分配概率分布。然而这种处理方式存在一个负作用:与正确标签相比,模型为所有的误标签都分配了很小的概率;然而实际上对于不同的错误标签,其被分配的概率仍然可能存在数个量级的悬殊差距。例如:在图片分类问题中,我们要将图片分成猫、狗、老虎三类。在一次训练中,我们给三类分配的概率分别为 [0.0010, 0.0001, 0.9989],从而最终得到 [0,0,1] 的 one-hot 编码作为分类结果(即 hard-target),我们认为图片所代表的是一只老虎。然而,softmax 函数输出的概率往往包含着类别之间潜在的相关性。在这个例子中,我们可以看到,图片可能是猫的概率比图片是狗的概率更接近图片是老虎的概率,这说明猫和老虎之间存在的内在联系更加强大。类似地,Hinton 也举例说:在识别一辆宝马汽车的图片时,分类器将该图片识别为清洁车的概率是很小的,然而这种概率比起将其识别为胡萝卜的可能是会大出很多。由于在宏观上由于这些概率都很小,这一部分的知识很容易在训练过程中淹没,这无疑是浪费了重要的可以用于将大规模网络的知识迁移到小规模网络中去的宝贵先验概率。

为了充分利用这种类类别之间的相关性,我们需要通过某种方式去改变概率分布,使其更加平缓。而 Hinton 仅仅对我们经常使用的 softmax 函数进行了一点点修改,就达到了这一目标,他究竟是怎么做的呢?

事实上,如下面的公式所示,Hinton 向 softmax 函数添加了一点「佐料」——参数「T, 温度」(如今 T 已经成为了许多深度学习模型的标配,例如在生成文本的 RNN 中提高 T 可以增加生成文本的多样性):


Hinton 新作「在线蒸馏」,提升深度学习分布式训练表现的利器

其中,z 为每一个类别输入的 logit。式中 T=1 时,退化成传统的 softmax;T无穷大时,结果趋近于 1/C,即所有类别上的概率趋近于相等。T>1 时,我们就能获得 soft target label。通过提高 T,softmax层的映射曲线更加平缓,因而实例的概率映射将更为集中,便使得目标更加地「soft」。

有了这个 distillation 的内核,Hinton 按照以下的步骤对大规模网络进行「蒸馏」:

Hinton 新作「在线蒸馏」,提升深度学习分布式训练表现的利器

知识蒸馏示意图(图片来自网络:https://www.zhihu.com/question/50519680,本图作者YJango)

  1. 训练大模型:先用 hard target(类似于 [0,0,1] 的 one-hot 编码)的样本训练。

  2. 计算 soft target:利用训练好的大模型来计算 soft target 。也就是大模型「软化后」再经过 softmax 的输出。

  3. 重新创建一个小的网络,该网络最后有两个 loss,一个是 hard loss,即传统的 softmax loss,使用 one-hot label;另外一个是 soft loss,即 T>1 的 softmax loss,使用我们第二步保存下来的 soft target label。

整体的 loss 如下式:

Hinton 新作「在线蒸馏」,提升深度学习分布式训练表现的利器

其中

Hinton 新作「在线蒸馏」,提升深度学习分布式训练表现的利器

Hinton 新作「在线蒸馏」,提升深度学习分布式训练表现的利器为第二步输出的 soft label。

用「软化」训练集训练小模型。训练小模型时 T 不变仍然较大,训练完之后 T 改为1。


  4. 预测时,将训练好的小模型按常规方式使用。

现在我们可以把 Hinton 的方法和下图所示的最初 knowledge distillation 的由来作个对比。

Hinton 新作「在线蒸馏」,提升深度学习分布式训练表现的利器

可见,本质上这相当于对数据进行了增强(augmentation),加入了类别之间的关联性的先验信息。将大规模网络学习到的这种关系包装在数据中,用这种更强的数据来训练小规模的模型,充分考虑到了类间的距离和类内的方差信息。从而提升了小规模模型的性能,达到了「蒸馏」的效果。与直接使用预训练模型的结构和权重相比,这是一种相对更「高级」的知识迁移的方式。

在线蒸馏?新瓶装旧酒?

时间的车轮一眨眼就驶向了 2018 年,随着移动终端上的深度学习等应用的出现,网络模型压缩成为了一个广受关注的领域,大量的研究者在 Hinton 的启发下,退出了 distillation 的诸多变形,进行了深入的优化。Hinton 则一如既往地尝试开拓更多新的深度学习范式;当然也可能是谷歌的计算资源太过充足所以遇到了别人没有机会遇到的问题(给跪),Hinton 开始尝试在大规模分布式计算环境下使用「在线蒸馏(online distillation)」方法。这是由于目前的分布式 SGD 方法遇到了瓶颈。而本地的「蒸馏」算法也因其对数据管道的计算操作过于复杂而暴露出越来越多的问题。

具体而言,在分布式 SGD 中,由于边际效益递减的规律,增加参与训练的机器数量而获得的计算效率的提升渐渐变小,直到毫无效果。另一方面,他们也想使用集成的深度学习模型提高预测的准确率。为了在不增加测试时间成本的情况下获得与集成学习同等的收益,他们对一个 n 路集成模型进行蒸馏(distill),得到一个单一模型,这包括两个阶段:

Hinton 他们将这种在线的蒸馏方式称为「codistillation」:即分布式环境中的每个节点之间都可以互为 teacher 和 student,并且互相提取内在的知识,用以提升其它节点的模型性能,具体的算法如下:

Hinton 新作「在线蒸馏」,提升深度学习分布式训练表现的利器

如算法 1 中所示,为了突破分布式 SGD 的瓶颈,Hinton 他们利用蒸馏算法提高了模型的训练效率。使用蒸馏方法更新一个网络的参数只需要对其他网络的预测结果,这些网络可以利用其他网络权重的副本进行本地计算。

值得注意的是,即使教师模型和学生模型是同一神经网络的两个实例,蒸馏也有好处,只要它们有足够的区别(比如,不同的初始化、以不同的顺序接收输入样本;可以参见论文第 3 节的经验证据)。这也说明这种「在线蒸馏」的方法是具有很强的普适性的。

Hinton 他们在这个工作中将蒸馏技术和分布式 SGD 相结合,从而使一个分布式 SGD 的工作组内的各个节点能够交换检查点(checkpoint)保存的网络信息,利用这种信息作为蒸馏出来的知识,通过「老师-学生」的训练,加速学生网络的训练。在这个过程中,学生节点和老师结点的角色是互换的,因此,各个网络互相促,进从而实现共同的蒸馏。

其实,就算「在线蒸馏」是新瓶装旧酒,那也是一个非常恰当的,闪闪发光的新瓶子。它利用蒸馏技术,降低了分布式 SGD 的通信开销,成功提高了预测的准确率,提升模型的计算性能!

结语

笔者通过这篇文章和大家一起回顾了知识蒸馏的相关知识,并且粗浅地了解了 Hinton 在这个领域所做的最新工作。我们可以看到,Hinton 作为一代宗师,每次都能够以超过常人的眼光发现研究的新方向,并且提出简单、美、效果卓越的解决方案,这与他对自然、对生物的神经系统、对生活、对所面临的问题的犀利的观察是密不可分的。这激励着所有的人工智能研究者开拓视野、放飞心灵,充满创造力地去探索新的未知领域。雷锋网 AI 科技评论这样的学术媒体也会不断地把最新的学术研究进展介绍给大家。

论文地址:https://arxiv.org/abs/1804.03235 ,雷锋网 AI 科技评论报道

长按图片保存图片,分享给好友或朋友圈

Hinton 新作「在线蒸馏」,提升深度学习分布式训练表现的利器

扫码查看文章

正在生成分享图...

取消
相关文章