基于对抗学习和知识蒸馏的神经网络压缩算法

期刊:计算机工程与应用

时间:2021

研究院:战略支援部队信息工程大学

关键词:知识蒸馏 ;对抗学习 ;互学习 ;模型压缩 ;人脸识别

方法简介

(1)改进经典的知识蒸馏损失,使其不仅能够学习 教师网络输出的正确软目标,而且能够从中间层获取丰 富的隐含知识

(2)引入对抗学习中的判别器,鉴别教师网络和学 生网络特征图的差异,进一步缩小大模型和容量有限的 小模型最优解空间之间的差异。

(3)在训练过程中采用互学习的策略,使教师网络 和学生网络学习对方的特征图,提升泛化能力。

为了使学生网络能够探索自己的最优 解空间,加入判别损失这一更加宏观的标准,使学生网 络在训练过程中具有更多的自主性

由于学生网络 具有较小的容量,很难使其完全精确地模仿教师网络的 软目标,增加对抗损失,使学生网络能够能快地收敛于 教师网络的最优解空间。由于判别器过早达到平衡会 使学生网络无法从教师网络学习到有效的梯度,引入对判别器的正则化,避免判别器支配后续的训练过程。

具体方法

基于分类概率的知识蒸馏优化

因此改进传统知识蒸馏的方法,忽略教师网络错误的预测分布,只把正确的预测分布传递给学生网络,具 体目标函数如式(3)所示
互蒸馏综述神经网络 知识蒸馏 神经网络压缩_机器学习

对抗学习辅助下的特征图迁移

  • 思想来源

由于学生网络的容量小,它 可能无法精确地再现某一特定的输出模态,并且实际中 学生网络与教师网络具有不同的结构,没有必要精确地 模拟一个教师网络的输出来达到良好的表现。因此本 研究提出一个面向教师网络和学生网络的对抗学习机 制。对抗训练策略缓解了人工设计损失函数的困难

  • 对抗损失

互蒸馏综述神经网络 知识蒸馏 神经网络压缩_最小化_02

在训练过程中,判别器的目的是最小化对抗损失, 确保正确区分两个不同的分布;学生网络的目的则是使 判别器无法区分其与教师网络的差异,以此构成对抗训 练 互蒸馏综述神经网络 知识蒸馏 神经网络压缩_互蒸馏综述神经网络_03

相较分类概率,高维特征图能够保留更多的特征, 采用高维的特征图作为判别器的鉴定对象能够使判别 器具有更强的鉴别能力,指导学生网络的更新,最小化 与教师网络的差异

学生网络和教师网络的深度互学习

采用学生网络和教师网络特征图间的 Jensen-Shannon 散度作为互学习的目标函数,如式(5)所 示。相较于 KL 散度,JS 散度是对称的,解决了 KL 散度 非对称的问题。

互蒸馏综述神经网络 知识蒸馏 神经网络压缩_互蒸馏综述神经网络_04

  • 损失目标函数
  • 互蒸馏综述神经网络 知识蒸馏 神经网络压缩_人工智能_05

  • 判别器

模型中的重要组成部分,必须在简单性和 网络容量之间取得平衡。判别器由三个全连接层组成 (128fc-256fc-128fc),中间激活为非线性激活单元 ReLU。 输出则是由 Sigmoid 函数给出的二元预测,判定输入的 特征图来自哪个网络。

实验对比

  • 本文方法
  • 消融实验