【科普】联邦知识蒸馏概述与思考

2022-03-30 11:11:10 浏览数 (1)

随着深度学习与大数据的进一步发展,效果好的模型往往有着较大的规模和复杂的结构,往往计算效率与资源使用方面开销很大,无法部署到一些边缘设备、移动终端或者嵌入式设备上。因此,如何在保证模型性能的前提下减少模型的参数量以及加快模型前向传播效率,这是一个重要的问题,总的来说不同的模型压缩与加速技术具体可以分为以下四类:

1)参数剪枝与共享(Parameter pruning and sharing):参数剪枝和共享用于降低网络复杂度和解决过拟合问题。网络模型越大,参数也越多,但是可能很大一部分参数是冗余的,所以我们可以将那些对输出结果影响不大的参数减去,这样就可以使得模型运行速度更快、模型文件更小。参数共享可以参考卷积神经网络每个过滤器的权重是固定的,可以看作模板,只关注数据的某一种特性,同时可以减少参数与计算量。具体如下图1所示。

2)低秩分解和稀疏性(Low-rank factorization and sparsity):从线性代数的角度上来说,低秩矩阵每行或每列都可以用其他的行或列线性表出,其包含大量的冗余信息。低秩分解的目的在于去除冗余,并且减少权值参数,进一步来说权重向量往往分布在一些低秩子空间,所以我们可以用少量参数重建权重向量。在机器学习角度可以采用两个K*1的卷积核替换掉一个K*K的卷积核,这将减少大部分参数。具体如下图2所示。

3)传输/紧凑卷积滤波器(Transferred/compact convolutional filters):由于卷积神经网络的滤波器通过一个个卷积核进行特征提取,而较大的卷积核可以通过小型卷积核的拼凑进行替换,即一个5*5的卷积核可以通过两个3*3的卷积核代替进行特征提取。通过增加卷积核之间的约束关系来优化空间和计算,从而实现压缩卷积核、减少计算的目的。

4)知识蒸馏(Knowledge Distillation):知识蒸馏是近年来发展起来的一类模型压缩与加速技术,其主要是利用一个已经训练好的复杂模型(作为教师),将其学习到的决策信息(知识)迁移到另一个轻量级模型(作为学生)中,帮助和指导学生模型的训练。知识蒸馏旨在通过将知识从深度网络转移到小型网络来压缩和改进模型。

图1:参数剪枝与共享

图2:矩阵向量低秩分解

联邦学习(FL)是一种机器学习设定,其中许多客户端(例如,移动设备或整个组织)在中央服务器(例如,服务提供商)的协调下共同训练模型,同时保持训练数据的去中心化及分散性。其核心目标是为了在不暴露数据的情况下分析和学习多个数据拥有者的数据。

进一步来说,联邦学习主要面临四大挑战:通信效率(communication efficiency)、系统异构性(system heterogeneity)、统计异构性(statistical heterogeneity)与隐私安全(privacy security)。同时由于联邦学习特有的客户端-服务器的系统架构,导致模型往往需要尽可能的便捷,同时需要保证传输模型时的实时性(面向多数设备)与快速性(传输效率高)。

自然而然的,人们就将联邦学习与模型压缩加速技术进行结合。通过模型压缩加速技术在联邦学习的训练过程、传输过程或推理过程进行优化,从而尽可能保障模型的传输高效率和实时性。而近期知识蒸馏的技术发展最为迅速,因此涌现出来不少的联邦知识蒸馏相关的论文与想法。

接下来,我们将从知识蒸馏概述和联邦知识蒸馏两大主题探讨相关论文、经典算法和自己的理解思考。

知识蒸馏概述

知识蒸馏将复杂且大的模型作为 Teacher模型,将模型结构较为简单需要优化的模型作为 Student模型,然后用Teacher来辅助Student模型的训练。知识蒸馏希望学生网络在模拟老师的输出的时候尽可能学到类与类之间的相似性和差异,从而增强学生模型的泛化能力。

知识蒸馏的主要关键点在于:知识的类型、蒸馏方法和教师-学生网络的结构。由于蒸馏方法往往大同小异,主要在于怎么设计Loss(但大部分也是从欧式距离、余弦距离、相似性等角度来设计损失),同时教师-学生网络的结构设计涉及到模型架构方面也较复杂,所以从知识类别角度来对知识蒸馏进行概述(个人感觉如何定义一个新的蒸馏的知识是创新最大的地方)。

Logits KD:作为知识蒸馏的开山之作(链接1),Hinton在2014年提出基于logits的知识蒸馏方法,主要思想在于用学生网络的预测logits去学习教师网络的输出logits,从而引导学生网络训练,可以学习到自身预测不出来的类之间的相似性知识。主要方法是通过基于温度参数的softmax函数,对输出logits进行软化,将其看作一种知识从教师端转移到学生端。具体过程如下图3所示。

图3:Logits KD框架图(链接2

Hints KD:该论文(链接3)是第一个考虑到模型中间隐藏层的知识蒸馏方法,其主要是将模型隐藏层的特征看作是一种知识,然后学生网络通过学习教师网络的隐藏层特征知识,可以提升学生模型自己的性能,同时可以和logits KD方法在一起结合使用。

其主要示意图如下图4所示。损失函数主要使用了MSE Loss,如果教师网络和学生网络的隐藏层输出特征不同大小,那么将其变换成相同大小再进行损失计算。

图4:Hints KD框架图与损失函数(链接3)

Attention KD:该论文(链接4)将神经网络的注意力作为知识进行蒸馏,并定义了基于激活图与基于梯度的注意力分布图,设计了注意力蒸馏的方法。大量实验结果表明AT具有不错的效果。

论文将注意力也视为一种可以在教师与学生模型之间传递的知识,然后通过设计损失函数完成注意力传递,本质上来说学生模型学习到了教师模型针对输入数据权重更高的地方,即输入数据对模型的影响程度。具体如下图5所示。

图5:Attention KD框架图与损失函数(链接4

Cooperative KD:通常来说,大小模型的效果对比十分显著,意味着logits的差别是很大的,因而使小模型去逼近大模型的logits无疑是比较难的任务。该论文(链接5)通过结合课程学习思想,思考模型从易到难训练可能会有更好的效果,而从易到难的方法简单来说就是大小模型同时在下游任务进行微调,并在每个周期epoch中让大模型引导小模型进行学习。

关键之处在于:teacher模型和student模型共同训练、teacher和student模型loss损失的计算。同时教师模型不再固定参数,而是随着学生模型反馈而进行轻微改动。主要过程如下图6所示。

图6:Cooperative KD框架示意图(链接5

Similarity KD:在该论文(链接6)中提出了一种新的知识蒸馏损失形式,称为相似性知识,灵感来自于相似的输入会倾向于在训练的网络中引起相似的激活模式,与以前的蒸馏方法相比,学生不需要模仿教师的表示空间,而是需要在其自己的表示空间中保持与教师网络成对的相似性。

保持相似性的知识蒸馏指导学生网络的训练,使在训练的教师网络中产生相似激活的输入也在学生网络中产生相似激活。具体如下图7所示,当输入相同的数据时,学生网络学习教师网络的神经元激活。

更进一步来说,如果两个输入在教师网络中产生高度相似的激活,那么引导学生网络,这也会导致两个输入在学生中产生高度相似的激活;相反地,如果两个输入在教师中产生不同的激活,我们就希望这些输入在学生中也产生不同的激活。

图7:Similarity KD框架与损失函数示意图(链接6

Relation KD:该论文中(链接7)作者引入了一种新的方法,称为关系知识蒸馏(RKD),它可以转移数据示例的相互关系,对于RKD的具体实现,作者提出了距离和角度的蒸馏损失来惩罚关系中的结构差异。

作者主要将关系结构看作是一种知识,然后又通过欧式距离与余弦距离作为损失函数来传递知识,从而使得RKD训练学生模型形成与教师相同的关系结构。具体如下图8所示。

图8:RKD示意图与框架结构(链接7

知识蒸馏概述部分总结

通过若干知识蒸馏论文的解读我们可以发现,相比较而言,蒸馏知识的创新比较重要,一种知识的定义主要在于两方面:

1)在模型学习训练过程中,什么构成了知识;

2)如何将知识从教师网络转移到另一个学生模型中。

进一步而言,我们可以考虑从哪方面定义知识,例如模型的预测标签、模型网络的中间层、模型权重等等方面,或者综合多个角度进行结合,也可以考虑基于任务的知识蒸馏(例如目标检测、语义分割等任务),从而设计出基于任务的知识/蒸馏算法或者师生模型架构。

在传统联邦学习,传输的数据量和模型比例正相关,而且会随数据的Non-IID分布影响并减少准确度,所以有一些文章思考如何结合知识蒸馏与联邦学习从而带来效率与性能方面的提升。接下来,我们将结合一些文章解析基于联邦学习的知识蒸馏算法与思想。

参考链接

[1] Hinton, Geoffrey, Oriol Vinyals, and Jeff Dean. "Distilling the knowledge in a neural network." arXiv preprint arXiv:1503.02531 2.7 (2015).

[2] https://towardsdatascience.com/paper-summary-distilling-the-knowledge-in-a-neural-network-dc8efd9813cc.

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

[4] Zagoruyko, Sergey, and Nikos Komodakis. "Paying more attention to attention: Improving the performance of convolutional neural networks via attention transfer." arXiv preprint arXiv:1612.03928 (2016).

[5] Jin, X., Peng, B., Wu, Y., Liu, Y., Liu, J., Liang, D., ... & Hu, X. (2019). Knowledge distillation via route constrained optimization. In Proceedings of the IEEE/CVF International Conference on Computer Vision (pp. 1345-1354).

[6] Tung, F., & Mori, G. (2019). Similarity-preserving knowledge distillation. In Proceedings of the IEEE/CVF International Conference on Computer Vision (pp. 1365-1374).

[7] Park, W., Kim, D., Lu, Y., & Cho, M. (2019). Relational knowledge distillation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 3967-3976).

END

0 人点赞