Knowledge Distillation with Generative Adversarial Networks

Abstract

  知识蒸馏(Knowledge distillation)旨在训练一种轻量分类器,能够在资源受限的多标签学习中提供精准的推理。分类器不是直接使用特征-标签对,而是由教师进行训练。比如,一个使用大量资源训练的高能力模型,这种分类器的精确度由于很难从教师那里学习到真实数据分布从而达不到理想水平。另一种方法是在类似生成对抗网络中,分类器对抗判别器,以保证分类器在博弈均衡状态下学习到真实的数据分布。但这种双人博弈方法因为高方差梯度更新,往往需要很长的时间才能达到博弈均衡状态。为了解决这些问题,作者提出了一种取名为KDGAN的包括分类器、教师、判别器的三人博弈游戏。分类器和教师通过蒸馏损失相互学习,并且通过对抗损失来对抗判别器。通过同时优化蒸馏损失和对抗损失,分类器或者教师将会在平衡状态下学得的真实数据分布。作者用一个具体的分布来近似分类器(或教师)学习到的离散分布。从具体分布出发,生成连续样本,得到低方差梯度更新,加快了训练速度。实际数据集的大量实验验证了KDGAN在精度和训练速度上的优越性。
  论文地址

Introduction

  在机器学习中,通常在模型的训练阶段比在模型部署阶段(推理阶段)使用更多的资源(比如输入特征或计算资源),这些多出来的部分被称为特权条款。图一展示了一个图像标记推荐的示例应用程序,在训练阶段比在推理阶段使用了更多的输入特征(即特权信息):训练阶段不仅输入的有图片,还有图片标题和注释,但推理阶段输入的只有图片本身。在一名智能手机用户上传一张图片后,他将给图片一个标签,但在手机上输入标签不方便及思考标签需要时间,如图1b所示,基于图片的标签推荐是非常有用的。另一个实例应用程序是手机人脸解锁。我们往往将人脸识别模型部署在手机上,这样合法用户就可以不需要远程服务或网络连接就可以解锁手机。训练阶段往往在功能强大的服务器上完成,其计算资源要比在手机上完成的推理阶段多得多。因此,一个关键的问题是如何利用特权条款,如,资源只能在训练阶段使用,训练出来的模型具有优秀的推理能力。

Alt text

  解决这一问题的典型方法是基于知识蒸馏。如图2左半部所示,知识蒸馏由一个分类器和教师组成。为了应对资源受限下的知识推理,分类器并不使用特权条款;另一方面,教师使用特权条款,例如,使用一个大模型容量或者使用更多的输入特征。一旦训练完成,教师师将为每一个训练实例输出一个被称为软标签的标签分布。接着,教师通过蒸馏损失如 the L2 loss on logits训练分类器预测软标签。这种训练过程通常被称为“提炼”教师的知识到分类器。由于教师通常不能很好地对真实数据分布进行建模,分类器很难从教师那里学习真实数据分布。
  生成对抗网络提供了另一种学习真实的数据分布的方法。受[49]的启发,作者首先提出了一种naive GAN(NaGAN),拥有两名玩家。如图2右半部分所示,NaGAN由一个分类器和一个判别器组成。分类器作为生成器,当给其一个实例,则生成相关的标签;判别器旨在区分真实标签和生成的标签。该分类器通过对抗损失从判别器处学习完美地建模博弈均衡时地真实数据分布。NaGAN的一个限制是达成博弈均衡需要大量的训练实例和大量的epoches,这就限制了其在当收集标注数据代价很高的领域的应用。在这种双人博弈框架中,来自判别器的更新分类器的梯度在对抗训练中经常消失或者爆炸,导致了该方法的训练需要大量时间。在有限训练实例和epochs的条件下,训练分类器学习到真实的数据分布是一个挑战。
  为了解决这个挑战,作者提出了取名为KDGAN的三人博弈框架,使用生成对抗网络来提炼知识。如图2所示,KDGAN由一个分类器,一位教师,一个判别器组成。除了知识蒸馏中的蒸馏损失和上述提及的NaGAN中的对抗损失外,作者还定义了从分类器到教师的蒸馏损失、教师和判别器之间的对抗损失。具体来说,分类器和教师作为生成器,目的是通过生成类似于真实标签的伪标签来欺骗鉴别器。与此同时,分类器和教师试图通过将他们的知识提炼成彼此的知识,从而就生成什么伪标签达成一致。通过将蒸馏和对抗损失表示为极小极大博弈,作者使分类器能够学习均衡时的真实数据分布(see Methods 3.2)。另外,分类器通过蒸馏损失接收从教师处传来的梯度,通过对抗损失接收从判别器处传来的梯度。从教师处传来的梯度通常具有低方差,这就降低了整个梯度的方差,因此加速了对抗训练(see Methods 3.3)。
  作者进一步考虑了降低从判别器处传来的损失的方差,以加速KDGAN的训练。当使用某些广为使用的策略梯度下降方法时,从判别器处获得的梯度可能存在较大的方差。从判别器处获得低方差的梯度是很重要的,因为分类器和教师生成的离散样本关于它们的参数是不可微分的。作者建议使用Gumbel-Max技巧将分类器和教师学习到的离散分布松弛为具体分布。他们使用具体的分布来生成连续样本,以此达到端到端的微分和充分控制梯度的方差(ps:为什么要连续样本,为什么离散样本不可微分,什么具体分布)。在给定连续样本的情况下,作者从判别器得到低方差梯度,以加速KDGAN的训练。总而言之,作者的贡献如下:

  1. 提出了一种用于多标签学习的KDGAN框架,该框架利用仅用于训练的资源来训练适合于资源约束推理的轻量级分类器。

  2. 通过减小梯度的方差来减少收敛所需的训练时间,这是通过KDGAN的设计和Gumbel-Max技巧实现的。

  3. 在图像标签推荐和深度模型压缩两方面进行了广泛的实验。实验验证了KDGAN相对于现有方法的优越性。

  这里简单的回顾一下知识蒸馏和生成对抗网路相关的研究。
  知识蒸馏目的是从一位强大的教师那转移知识到一个轻量的分类器那里。如,Ba and Caruana[7]通过L2 loss匹配logits训练一个浅层的分类器网络来模拟深度的教师网络。Hinton等人[23]通过训练一个分类器来预测教师提供的软标签泛化了这个工作。Sau and Balasubramanian[39]进一步添加随机扰动到软标签上,以此来模拟从多位老师那里学习。Remero等人[36]提出使用教师网络的中间层而不是软标签来训练分类器。不像之前在分类问题上的工作,Chen et al.[10]应用知识蒸馏和提示学习(hint learning)到目标检测问题上。也有一些工作充分利用知识蒸馏在不同域之间转移知识,比如,在高质量和低质量图片之间[41]。Lopez-Paz el al.[29]将使用特权信息的只是蒸馏统一为一个泛化提炼,存在一个使用特权信息作为输入的与训练好的教师。和知识蒸馏相比,该论文提出的KDGAN框架引入了一个判别器以确保分类器能够在均衡状态时学习到真实数据分布。
  GAN首次被提出来生成连续数据17,通过极大极小博弈对抗地训练一个生成器和判别器。因为离散数据使得判别器的梯度很难通过后向传播来更新生成器,GAN最近才被用来生成离散数据。和那些拥有两个玩家的GANs不同,Li et al.提出一种具有三个玩家的GAN,取名为Triple-GAN13。该论文提出的KDGAN同样拥有三个玩家,但不同于Triple-GAN:(1) KDGAN中的两个生成器在给与特征的条件下,学习标签的条件分布。但是Triple-GAN中的两个生成器分别学习给定特征下标签的条件分布,给定标签下特征的条件分布。(2)KDGAN中两个生成器产生的样本都是离散数据,但Triple-GAN中的两个生成器产生的既有离散数据也有连续数据。这些不同导致了不同的目标函数和训练技术。比如,KDGAN可以使用Gumbel-Max trick来生成从两个生成器产生的样本,但是Triple-GAN不能。
  作者探索了将知识蒸馏和生成对抗网络整合在一起。工作[51]提出来相似的想法,其引入一个判别器来训练分类器。但其和该论文的不同之处在于他们的判别器训练分类器学习教师产生的数据分布,但KDGAN的判别器训练分类器学习真实数据分布。
  作者将KDGAN用来解决深度模型压缩和图像标签推荐问题,但也可以将KDGAN应用到解决其他特权条款可用的问题。比如,考虑上下文信息作为特权信息在意图追踪中使用,将用户评论作为特权条款在电影推荐中使用。

Methods

  作者研究从一个使用特权条款(记为$\varrho$)训练的教师来训练一个轻量分类器来满足受限的推理的问题。推理可能受到(1)以有限的计算资源实时完成推理的要求;(2)缺少某些输入特征。在现有工作[29]之后,作者使用多标签学习问题作为他们用于说明提出的方法的目标应用场景。
Alt text
  因为特权条款只在训练阶段可用,该问题的目标是训练一个不使用特权特权条款就具有高效推理的轻量分类器。为了实现这个目标,作者从NaGAN开始。何其它两玩家框架类似,NaGAN要求大量的训练实例和长时间的训练,这很难在实践中使用。为了解决这个问题,作者提出了三玩家框架的KDGAN,该框架可以加速训练同时保持均衡。

NaGAN Formulation

  首先介绍NaGAN,其包括玩极大极小博弈的一个分类器$C$和一个判别器$D$。由于$D$不用于推理,因此可以利用特权条款。在NaGAN中,$C$给定特征$x$生成伪标签$y$,遵循分类分布$P_c(y|x)$;$D$则计算给定特征$x$标签$y$来自于真实数据分布$p_u(y|x)$的概率${p_d}^{\varrho}(x,y)$。遵循IRGAN[49]的价值函数,作者为NaGAN的极大极小博弈定义的价值函数$V(c,d)$如下:
alt text
令$h(x,y)$和$g(x,y)$分别为$C$和$D$的评分函数,定义$P_c(y|x)$和$p_u(y|x)$如下所示:
alt text
评分函数的实现有很多种方法。比如,$h(x,y)$可以是一个多层感知机。对于具体应用下的评分函数,论文在实验部分有详细介绍。这样的双人博弈框架通过交替地更新$C$和$D$来训练。除非达到均衡状态,这时$C$学习到真实的数据分布,否则训练会持续进行。在均衡状态时,在决定给定标签是否由$C$生成时,$D$不能比随机猜测做的更好。
  作者注意到知识蒸馏和NaGAN的优缺点时可以互补的:(1)知识蒸馏通常要求少量的训练实例和较短的训练时长,但不能确保均衡时$P_c(y|x) = p_u(y|x)$。(2)NaGAN可以确保均衡时$P_c(y|x) = p_u(y|x)$,但通常需要大量的训练实例和训练时间。作者旨在提出一种单一的框架,保持这两种方法的优点,避免二者的缺点。

KDGAN Formulation

  作者将KDGAN形式化表示为一个具有分类器$C$,教师$T$,判别器$D$的极大极小博弈。同分类器$C$一样,教师$T$基于分类分布${p_t}^{\varrho}(y|x) = softmax(f(x,y))$来生成伪标签。$T$和$D$都可以使用特权条款。在KDGAN中,$D$目标是最大化正确区分真标签和伪标签的概率,但$C$和$T$的目标是最小化它们生成的伪标签被$D$拒绝的概率。同时,$C$通过模拟$T$学习到的分布从$T$那学习。为了建立一个通用框架,作者还使$T$能够从$C$学习,因为实际上,通过与学生互动也可以提高教师的能力。这样的相互学习帮助$C$和$T$减少它们生成不同的伪标签的概率。形式化的,作者为KDGAN中的极大极小博弈定义的价值函数$U(c,t,d)$如下:
alt text
$\alpha \in (0,1)$,$\beta \in (0,+\infty)$,$\gamma \in (0,+\infty)$都是超参数。上式中期望项是对抗损失,${\mathcal{L}_{DS}}^c$和${\mathcal{L}_{DS}}^t$是蒸馏损失。蒸馏损失可以有几种定义的方式,比如$L2$损失或KL散度。${\mathcal{L}_{DS}}^c$和${\mathcal{L}_{DS}}^t$分别用来训练分类器和教师。
  下面是理论分析部分,证明KDGAN的分类器能够学习到真实的数据分布。公式(4)相当于一个变形。公式(5)两个期望的最大值是那两个分布的JS散度这一步我没推出来(利用JS散度的定义)。该部分证明了KDGAN当且并仅在$P_c(y|x)={p_t}^{\varrho}(y|x)= p_u(y|x)$到达均衡状态,这是分类器学得真实的数据分布。
alt text

KDGAN Training

  在这一部分,作者细述通过减少需要的epoches数来加速KDGAN训练的技术。正如在之前的研究[8,46]中讨论过那样,训练速度和梯度的方差有很大的关系。和NaGAN相比,KDGAN框架被设计的可以减少梯度的方差。这是因为高方差的随机变量可以被低方差的随机变量减少,正如之前讨论的,$T$能提供比$D$更小方差的梯度。为了减少来自$D$的梯度的方差,并获得对方差的充分控制,作者进一步提出从松弛离散样本得到的连续空间那获得梯度,也就是伪标签,通过重新参数化技巧在分类器(或教师)和判别器之间传播成连续样本。
  首先,我们展示KDGAN如何减少梯度的方差。正如上述所讨论的那样,在NaGAN中,$C$仅从$D$处接收梯度$\triangledown _cV$,但在KDGAN中,它从$D$和$T$那都接收到梯度$\triangledown _cU$,如下所示:
alt text
和之前的研究一致,作者观察到$\triangledown _c{\mathcal{L}_{DS}}^c$往往比$\triangledown _c{\mathcal{L}_{AD}}^n$具有跟小的方差,随后的实验分析也证明了这一点。因此,这轻松地展示了KDGAN中$C$的梯度的方差比NaGAN中的更小,如下所示:
alt text

  接着,作者进一步使用重参数化技巧Gumbel-Max trick来减少梯度的方差。Gumbel-Max trick是一种从离散分布取样的方法,它的形式可以允许我们定义一种可微分的,离散分布的近似取样。Gumbel-Max trick的本质是使用其参数的可微函数加上Gumbel分布的随机变量来重参数化生成离散样本。为了在从分类分布$p_c(y|x)$使用Gumbel-Max trick来生成离散样本,一种具体分布[25,31]被使用。作者使用一个具体的分布$q_c(y|x)$来生成连续样本并且使用连续样本来计算分类器关于对抗损失的梯度$\triangledown _c{\mathcal{L}_{AD}}^n$,如下所示:
alt text
使用温度参数$\tau$在训练阶段控制梯度的方差。当有一个高温度参数时(ps:这里温度不知道是是什么意思),来自具体分布的样本很平滑,导致梯度的方差比较小。注意,具体分布的缺点在于,在高温下,它变得近似原始分类分布,准确度较低,这导致偏差梯度估计。在实验部分会讨论如何调温度参数。
  除了改善$C$的训练之外,同样的方法也可以应用于改善$T$的训练。KDGAN的总题训练逻辑如算法1所示。三个部分都可以先单独预训练,然后通过小批量随机梯度下降联合训练。
alt text

Conclusion

  作者提出了一种取名为KDGAN的框架应用于带特权条款的多标签学习,该框架联合了知识蒸馏和生成对抗网络。KDGAN包括分类器、教师、判别器的三人博弈游戏。分类器和教师通过蒸馏损失相互学习,并且通过对抗损失被对抗地训练来对抗判别器。作者证明了分类器将会在平衡状态下学得真实数据分布。作者使用具体分布在对抗训练阶段控制梯度的方差并且获得低方差的梯度来加速训练。作者展示了KDGAN在图像标签推荐和深度模型压缩这两个重要的应用上取得了当前最佳表现。对未来的工作,作者将探索用于确定模型超参数的自适应方法,以实现更好的训练动态。

Reference

[7] J. Ba and R. Caruana. Do deep nets really need to be deep? In NeurIPS, 2014.
[8] L. Bottou, F. E. Curtis, and J. Nocedal. Optimization methods for large-scale machine learning.
arXiv preprint arXiv:1606.04838, 2016.
13 L. Chongxuan, T. Xu, J. Zhu, and B. Zhang. Triple generative adversarial nets. In NeurIPS,
2017.
17 I. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, and
Y. Bengio. Generative adversarial nets. In NeurIPS, 2014.
[23] G. Hinton, O. Vinyals, and J. Dean. Distilling the knowledge in a neural network. In NeurIPS
workshop, 2014.
[29] D. Lopez-Paz, L. Bottou, B. Schölkopf, and V. Vapnik. Unifying distillation and privileged
information. In ICLR, 2016.
[36] 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.
[46] G. Tucker, A. Mnih, C. J. Maddison, J. Lawson, and J. Sohl-Dickstein. Rebar: Low-variance,
unbiased gradient estimates for discrete latent variable models. In NeurIPS, 2017.
[49] J.Wang, L. Yu, W. Zhang, Y. Gong, Y. Xu, B.Wang, P. Zhang, and D. Zhang. Irgan: A minimax
game for unifying generative and discriminative information retrieval models. In SIGIR, 2017.
[51] X.Wang, J. Qi, K. Ramamohanarao, Y. Sun, B. Li, and R. Zhang. A joint optimization approach
for personalized recommendation diversification. In PAKDD, 2018.