ICML 2021 | AlphaNet：基于α-散度的超网络训练方法

©作者 | 韩翔宇

学校 |  南昌大学

研究方向 | 神经网络架构搜索

概述

本文是 ICML 2021 收录的 NAS（神经网络结构搜索）领域的重磅论文，其作者是 AttentiveNas 的原作者，在 AttentiveNas 的基础上，添加了 α- 散度损失函数，在 ImageNet（NAS 方向）排行榜中取得了 SOTA 结果。

论文标题：

AlphaNet: Improved Training of Supernets with Alpha-Divergence

论文链接：

https://arxiv.org/abs/2102.07954

代码链接：

https://github.com/facebookresearch/AlphaNet

解决的问题

尽管权值共享的 NAS 有较好的效果，因为这种方法构建了一个超网，利用子网来训练超网。然而，权重共享 NAS 的成功很大程度上依赖于将超网络的知识提取到子网络。如果广泛使用蒸馏中的散度，例如 KL 散度，可能导致学生网络高估或者低估教师网络的不确定性，导致了子网的效果变差。论文提出了更广义的 α- 散度来改进超网训练。通过自适应选择 α- 散度，避免了高估或者低估教师模型不确定性。改进之后的 AlphaNet 在 ImageNet top-1 的精度达到了 80%，且参数量只有 444M。

简介

传统的 NAS 方法代价非常大，需要数百个网络结构从头训练、验证效果。超网将所有候选体系结构组装成一个权重共享网络，每个网络结构对应一个子网络。通过同时训练子网和超网，子网可以直接从超网中获得的权重，用来重新训练和验证，而不需要单独训练或微调每个结构，因此成本大大降低。为了稳定超网训练和提高子网络的性能，大家广泛使用知识蒸馏的方式。知识蒸馏用超网中最大的子网预测的软标签来监督所有其他子网，提取教师模型的知识来提高子网性能。

一般情况下，知识蒸馏使用 KL 散度衡量师生网络之间的差异。但是如果学生网络对教师网络的某些部分覆盖不完整，对学生网络的惩罚大。因此，学生模型往往高估了教师模型的不确定性，不能准确近似教师模型的正确预测。

为了解决这个问题，论文提出用更广义的 α- 散度代替 KL 散度，具体来讲，通过自适应控制散度度量中的 α，可以同时惩罚对教师网络高估或者低估教师网络不确定性的行为，鼓励学生网络更好地近似教师网络。直接优化 α 散度可能会受到梯度的高方差的影响，论文游提出了一种简单的梯度裁剪技术，稳定训练过程。通过提出的自适应 α- 散度，我们能够训练高质量的 alphanet，在 200 到 800 MFLOPs 范围表现都达到 SOTA 效果。

关于知识蒸馏的相关知识

在权值共享的 NAS 中，知识蒸馏是重要的方法。假设超网有可训练的参数 θ，训练超网的目的是学习 θ，让所有的子网都能同时得到优化，达到较高的准确率。

上图描述了采用知识蒸馏的超网训练过程。在每个训练步骤中，给定一小批数据，对超网和几个子网络采样。当超网使用真实标签训练时，所有抽样的子网络都使用超网预测的软标签进行监督训练。然后对所有采样的网络的梯度进行融合，更新超网参数。在第 t 个 step 时，超参数 θ 被更新为如下（其中 ε 是 step 的数量）：

g 定义如下：其中 LD 是数据集 D 上超网的损失函数（交叉熵）， 为权重系数， 是超网和采样出子网的 KL 散度， 和 表示：输入 x 的超网和子网 s 的输出概率。

知识蒸馏中KL散度的局限性（α-散度解决的问题）

KL 散度广泛用于衡量教师网络和学生网络之间的差异性。但是 KL 散度的缺点在于，当学生网络高估了教师网络的不确定性时，不能充分惩罚学生网络。我们首先列举一下 KL 散度的公式：

▲ KL散度公式（公式2为f散度，即KL散度的一般形式）

我们用 P 代表教师网络的某个输出概率，用 Q 代表学生网络的某个输出概率。KL 散度有个避零性质，当 P=0 时，我们看等式右半部分 log(P/Q)，因为分子为 0，我们知道 log 函数趋向于 0 的时候，值趋向于负无穷。无论分母 Q 怎么变大（也就是是学生网络某个输出概率变大），值始终是负无穷，对应 KL 散度低，也就是对学生网络的 Q 惩罚小，所以学生网络就可以肆无忌惮地预测，即使预测错误很离谱，也不会惩罚。

被如下图所示，针对某个输出，橙色的教师网络主要预测的类别是 3（概率最大的是 3，所以输出为 3），其次是类别 4，预测其他类别的概率小的可以忽略；再看绿色的学生网络，预测的最主要的类别是 4（概率最大的是 4，所以输出为 4，输出的类别就已经和教师网络有偏差了），而且学生网络预测其他类别的概率也比教师网络高，所以学生网络的不确定性比教师网络大，换句话说，学生网络高估了教师网络的不确定性。

反过来想，我们用 Q 代表教师网络的某个输出概率，用 P 代表学生网络的某个输出概率（和上个假设掉换个位置）。这次让 Q=0，这时就要再讲 KL 散度的另一个性质：零强制性。用这个例子简单来简单解释，就是分母 Q 为 0，P 也要为 0。这很好理解，当分母为 0 的时候，分子如果不为 0 就会报除 0 错误，但是当分子分母都为 0 时，P/Q=1， log(P/Q)=log1=0。

说完这个性质以后，我们想一想，如果 Q 为 0，KL 散度的值为 0，我们希望损失越小越好，我们希望损失比 0 更小，也就是负数。因此最小化 KL 散度会鼓励 P 避免趋向 0（因为当 P 为 0 的时候，Q 可能为 0，Q一旦为 0，损失函数就一定为 0，也就不能继续变小了，我们希望损失函数越小越好，当然 P 不愿意趋向 0）。

从论文的角度来说，也就是学生模型会避免低概率模式，也就是学生模型会倾向于较高的概率。如下图所示，学生模型预测类别 2 的概率为 100%，几乎没有不确定性可言，而教师网络除了预测类别最多的 2，其他类别也有概率分布，所以教师网络的不确定性比学生网络大，也就是学生网络低估了教师网络的不确定性。

使用α-散度训练的超网

为了解决上述 KL 散度的不确定性，论文提出了一个灵活的 α- 散度（α∈R 且 α≠0 且 α≠1）：

和 代表每个类的离散分布，共有 个类别。当 时， 的极限就是 ，同样， 是 时 的极限。 散度的关键点在于，可以通过选择不同的 值来集中惩罚不同类型的差异（低估或高估）。

如图所示，当阿尔法为负值的时候，蓝色线代表了学生网络 高估了教师网络 的不确定性的情况，这时 很大。紫色线表示学生网络 低估了教师网络 的不确定性的情况，此时 很小。当 为正数的时候，情况正好相反。

为了同时缓解超网训练时的高估和低估的问题，用一个正的 和一个负的 ，在知识蒸馏损失函数中使用 和 中最大那个，也就是：

总的 KL 散度为：

进一步改进的稳定的α-散度

人们倾向于将 和 设置为较大，以确保学生模型在低估或高估教师模型的不确定性时受到足够的惩罚。但是直接通过增加 绝对值这种方法，会让优化变得困难，我们首先来看一下 散度的梯度：

不难看出，如果 很大，则 也可能变得很大，从而影响了训练的稳定性。为了让训练变得稳定，我们把 的最大值限定为 （如果小于 则不变，最大不能超过 ），重新定义梯度表达式为：

该梯度等价于 p 和 q 的 f 散度（f 散度是 KL 散度的一般形式，如果你认真看了这篇文章，应该会注意到 KL 散度的公式有两个，没错，第二个公式就是 f 散度），进行梯度更新相当于最小化有效散度。通过裁剪重要性权重的值，我们优化的仍然是一个散度度量，但对基于梯度的优化更友好。

具体实现

自适应的 α- 散度的设置：α- 和 α+ 分别控制对过高估计和过低估计的惩罚幅度。并且，β 控制了教师模型和学生模型之间的密度比率范围（也就是子网占了超网多大一部分）。通过实验发现 AlphaNet 的方法在 α−，α+ 和 β 的选择范围内表现相对稳健。实验中选择 α−=−1，α+=1 和 β=5.0 作为默认值。作者在 Silimmable Network 和权值共享的 NAS 两个层面分别做了实验：

8.1 Slimmable Network部分

Slimmable 网络是支持选择多种通道宽度的超网，其搜索空间包含不同宽度的网络和所有其他参数（深度、卷积类型、kernal 大小）都是相同的。Slimmable 允许不同的设备或应用程序根据设备上的资源限制，自适应地调整模型宽度，以实现最佳精度与能效的权衡。

使用 Slimmable MobileNet v1 和 v2 测试，其中 v1 的宽度范围是 [0.25,1]，v2 的宽度是 [0.35,1]。在每次训练迭代中，分别对 channel 宽度最大的最大子网络、channel 宽度最小的最小子网络和两个随机子网络进行采样，累积梯度（这种方式被称为三明治法则）。使用 gt 标签训练超网，使用知识蒸馏训练采样到的子网，设置知识蒸馏中的系数 γ=3，作为抽样获得的子网数量。

为了验证自适应 α- 散度的有效性，使用它替换 baseline 中的 KL散度。使用 SGD 优化器训练 360epoch，动量为 0.9，重量衰减为 10−5,dropout 为 0.2。我们使用余弦学习速率衰减，初始学习速率为 0.8，batch_size 为 2048，使用 16 块 GPU。

在 ImageNet 上进行训练，上表是训练得到的最佳精度，可以看到，无论在任何宽度上，自适应的α. 散度效果均优于 KL 散度（KL-KD）和不加 KL 散度（KD）。

8.2 权值共享的NAS部分

8.2.1 简 介

大多数基于权值共享的 NAS 由以下两个阶段组成（基于强化学习的NAS也一样）：

• 阶段 1：使用可微分权值共享或者看做一个黑盒的优化

• 阶段 2：从头开始训练深度神经网络，以获得最佳的准确率和最终的效果

但是这类 NAS 有缺陷：如果需要不同的硬件约束条件，需要重新搜索。而且要求对所有的候选结果从头训练，达到理想的准确率。因此，显著增加了 NAS 的搜索成本。

论文使用的权重共享的 NAS 是基于超网的权重共享 NAS，搜索过程如下：

阶段1：联合优化搜索空间中的超网和所有可能被采样的子网，使所有可搜索网络在训练结束时都能获得较好的性能。

阶段2：然后所有的子网络同时被优化。然后可以使用典型的搜索算法，比如进化算法，来搜索感兴趣的最佳模型。每个子网络的模型权值直接从预训练的超网络继承而来，无需再进行再训练或微调。

与基于 RL 的 NAS 算法和可微 NAS 算法相比，基于超网络的权重共享的优势主要有：

1. 只需要对超网进行一次训练。搜索空间中定义的所有子网络在第1阶段完全优化后即可使用。不需要再训练或微调；

2. 在阶段 1 中，所有不同模型规模的子网络进行联合优化，找到一组帕累托最优模型，这一组模型天生支持各种需要考虑的资源。

注意，权重共享 NAS 的一个主要步骤是同时训练搜索空间中指定的所有子网络收敛。与训练 Slimmable 神经网络类似，也是用基于 KL 散度的知识蒸馏方法，强制所有的子网络使从超网中学习来实现的。

8.2.2 具体训练

为了简单起见，我们使用均匀抽样策略，使用和 Silimable 的三明治法则相同的方式，每次迭代训练四个网络。使用 SGD 和余弦淬火策略衰减学习率，使用 AutoAugment 进行数据增强，设置数据的标签平滑为 0.1。

上表是论文采用的搜索空间，MBConv 是 mobileNet 的倒残差模块。使用 swish 激活函数，channel width 表示块的输出维度。输入的分辨率表示候选的分辨率，为了简化数据加载过程，预选选取固定大小 224*224，然后再使用双三次插值将它们重新缩放到我们的目标分辨率。

8.2.3 验证

我们比较通过知识蒸馏策略得到不同超网，评价采用准确率和 FLOPs 的帕累托最优，主要包括以下三点：

1. 首先从超网中随机抽取 512 个子网，并估计它们在 ImageNet 验证集上的准确性。

2. 对性能最好的 128 个子网应用交叉和随机变异，将交叉规模和变异规模都固定为 128，共产生 256 个新的子网。然后我们评估了这些子网的性能。

3. 重复步骤 20 次步骤 2，得到 5376 个子网。

8.2.4 最终效果

上表可以看出， AlphaNet 相比于普通的 KL 散度和不使用 KL 散度的训练方法，在各个参数量阶段的准确率都得到了提升。

上图分别是不用 KL 散度（棕色），使用 KL 散度（绿色），AlphaNet（红色）的训练收敛曲线，可以看出，使用 AlphaNet 训练在前期收敛速度会慢于普通的 KL 散度，但是在后期收敛速度会超过普通的 KL 散度。

论文将 AlphaNet 和现有的 NAS 做对比，相比于 EfficientNet、MobileV3、MNasNet、BigNas 等模型都有较为明显的提升，取得了 SOTA 的效果。

更多阅读

#投 稿 通 道#

 让你的文字被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。

总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。

📝 稿件基本要求：

• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注 

• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题

• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算

📬 投稿通道：

• 投稿邮箱：hr@paperweekly.site 

• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿

△长按添加PaperWeekly小编

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·