CVPR2021 | AttentiveNAS：通过注意力采样改善神经架构搜索

本文是 Facebook 发表在CVPR 2021的NAS论文。在 OFA 和 BigNAS 在训练超网络的采样技巧上，进一步提出了基于帕累托感知的采样策略，提高了超网训练的效率，并取得了 SOTA 的结果。

• 文章题目：AttentiveNAS: Improving Neural Architecture Search via Attentive Sampling.
• 文章链接：https://arxiv.org/pdf/2011.09011.pdf
• 文章代码: https://github.com/facebookresearch/AttentiveNAS

摘要

神经结构搜索(NAS)在设计最先进的(SOTA)模型方面表现出了巨大的潜力，既准确又快速。近年来，BigNAS 等两阶段 NAS 将模型训练和搜索过程解耦，取得了良好的搜索效率。两阶段 NA S在训练过程中需要对搜索空间进行采样，这直接影响最终搜索模型的准确性。尽管均匀抽样的广泛应用是为了简化，但它不考虑模型性能的帕累托前沿，而帕累托前沿是搜索过程中的主要关注点，因此错过了进一步提高模型精度的机会。在这项工作中，我们建议关注于采样网络，以提高性能的帕累托。在训练过程中，本文还提出了有效识别帕累托网络的算法。无需额外的再训练或后处理，就可以通过广泛的 FLOPs 同时获得大量的网络。本文发现的模型家族 AttentiveNAS 模型在 ImageNet 上的准确率最高，从77.3%到80.7%，优于包括 BigNAS、Once-for-All networks 和 FBNetV3 在内的 SOTA 模型。并且本文还实现了 ImageNet 的精度为80.1%，只需491 MFLOPs。

动机

深度神经网络(DNNs)近年来取得了显著的经验成果。然而，网络规模和计算成本的快速增长给将DNNs引入边缘设备带来了很大的挑战。设计精确而高效的网络是一个重要而具有挑战性的问题。

神经结构搜索(NAS)为高效的神经网络设计提供了一个强大的自动化工具。NAS需要优化模型架构和模型参数，从而产生了一个具有挑战性的嵌套优化问题。传统的NAS算法利用进化算法和强化学习，并且由于MFLOPs需要成千上万的模型，可能会非常昂贵。最近的NAS技术进步将参数训练和架构优化解耦为两个独立的阶段：

• 第一阶段通过权重共享对搜索空间中所有候选网络的参数进行优化，使所有网络在训练结束时同时达到优越的性能。
• 第二阶段利用典型的搜索算法，如进化算法，在各种资源约束下找到性能最佳的模型。

这种NAS范式提供了最先进的经验结果，具有很高的搜索效率。

两阶段NAS的成功很大程度上依赖于第一阶段的候选网络训练。为了使每个候选网络都有更好的性能，训练的过程是从搜索空间中抽取候选网络，然后通过一步随机梯度下降(SGD)对每个样本进行优化。关键是要弄清楚在每个SGD步骤中采样哪个网络。现有的方法通常使用均匀的抽样策略，对概率相等的所有网络进行抽样。尽管有很好的结果被证明，但是均匀采样策略使得训练阶段和搜索阶段是独立的。更具体地说，搜索阶段关注的是位于准确性和推理效率的帕累托前沿的网络集合，而训练阶段并不是为了改进帕累托前沿而进行的，对每个网络候选对象都同等重要。这种方法错过了在训练阶段提高网络在帕累托上的准确性的机会。

在这项工作中，建议通过更多地关注那些更有可能产生更好的帕累托前沿的模型来改进基准均匀抽样。具体回答以下两个问题:

• 在训练过程中，应该采样哪些候选网络集?
• 应该如何有效地抽样这些候选网络，而不增加太多的训练开销?

为了回答第一个问题，本文探索了两种不同的抽样策略。第一种策略，称为“最优策略”（BestUp），在传统的帕累托最优 NAS 之后，研究一种帕累托最佳前沿感知采样策略，将更多的训练预算用于改进当前的帕累托最佳前沿。第二种策略被称为“最差策略（(WorstUp），专注于改善候选网络，从而产生最差情况下的性能权衡。本文将这些候选网络称为帕累托最差集。这种采样策略类似于难例挖掘，将 Pareto-worst 模型视为难训练示例。突破帕累托最坏网络的极限，有助于更新权享网络中最优参数，使所有参数都得到充分训练。

第二个问题也不是无关紧要的，因为在帕累托最佳和最坏的前面确定网络不是简单的。本文提出了两种方法来利用1)训练损失和2)预训练预测器预测的准确性作为准确性比较的代理。总的贡献可总结如下:

• 提出了一种新的策略，AttentiveNAS，通过对帕累托最佳（Pareto-best）或最差（Pareto-worst）前沿的网络进行注意抽样来改进现有的两阶段NAS。比较了 BestUp 和 WorstUp 两种不同的采样策略。
• 提出了两种方法来有效地指导抽样到帕累托最佳或最差前沿。
• 考虑到搜索 AttentiveNAS 模型家族的 FLOPs 限制，本文实现了 SOTA 的 ImageNet 精度。例如，AttentiveNASA0 比 MobileNetV3 的准确率高2.1%，而相比 FBNetV3, AttentiveNAS-A2 的准确率高 0.8%。

Two-stage NAS

传统 NAS 问题建模

$$\begin{array}{l}\underset{\alpha \in \mathcal{A}}{\mathrm{minL}}\left(W_{\alpha }^{\ast };{\mathcal{D}}^{val}\right)\\ \text{ s.t. }{W}_{\alpha }^{\ast }=\underset{W_{\alpha }}{\mathrm{argL}}\left(W_{\alpha };{\mathcal{D}}^{trn}\right)\\ \mathrm{FLOPs}(\alpha )<\tau \end{array}(1)$$

典型的NAS目标，如等式（1）将搜索范围限制为仅小型子网，从而产生了一个具有挑战性的优化问题，无法利用过参数化的优势。 另外，在等式（1）中定义了NAS的优化仅限于一个单一资源约束。 在各种资源限制下优化 DNN 通常需要进行多次独立搜索。

为了缓解上述缺点，近来，一系列 NAS 技术提出将约束优化问题（1）分解为两个单独的阶段：1）无约束预训练：通过权重联合优化搜索空间中指定的所有可能候选DNN 共享时不考虑任何资源限制； 2）资源受限的搜索：在给定的资源约束下识别性能最佳的子网。 在这个方向上的最新工作包括BigNAS，SPOS，FairNAS，OFA 和 HAT。

第一阶段：无约束的预训练

无约束的预训练阶段的目标是学习权重共享网络的参数。 通常将其设计为解决以下优化问题：
$$\underset{W}{\min }{\mathbb{E}}_{\alpha \in \mathcal{A}}\left[\mathcal{L}\left(W_{\alpha };{\mathcal{D}}^{trn}\right)\right]+\gamma \mathcal{R}(W)(2)$$
其中 W 代表网络中的共享权重。 $W_{\alpha }$ 是由架构 $\alpha$ 指定的 W 的子网，而 R（W）是正则项。 BigNAS 中提出的R（W）的一个例子表示如下：
$$\mathcal{R}(W)=\mathcal{L}\left(w_{{\alpha }_s};{\mathcal{D}}^{trn}\right)+\mathcal{L}\left(w_{{\alpha }_l};{\mathcal{D}}^{trn}\right)+\eta \Vert W{\Vert }_2^2(3)$$
其中 ${\alpha }_s$ 和 ${\alpha }_l$ 分别代表搜索空间 A 中的最小和最大候选子网。 $\eta$ 是重量衰减系数。 这在 BigNAS 中也称为三明治训练规则。

实际上，期望项用等式（2）表示通常用 n 个均匀采样的体系结构近似，并通过SGD进行求解，如下图。 请注意，较小和较大的 DNN 都在等式（2）中进行了联合优化。 这种表述允许通过权重共享和知识蒸馏将知识从较大的网络转移到较小的网络，从而提高整体性能。

第二阶段：资源约束的搜索

在阶段1进行预训练之后，所有候选DNN都经过了充分优化。 下一步是搜索可产生最佳性能和资源折衷的DNN，如下所示：
{ α i ∗ } = arg ⁡ min ⁡ α i ∈ A L ( W α i ∗ ; D v a l ) ,  s.t.  FLOPs ⁡ ( α i ) < τ i , ∀ i ( 4 ) \begin{array}{l} \left\{\alpha_{i}^{*}\right\}=\underset{\alpha_{i} \in \mathcal{A}}{\arg \min } \mathcal{L}\left(W_{\alpha_{i}}^{*} ; \mathcal{D}^{v a l}\right), \\ \text { s.t. } \operatorname{FLOPs}\left(\alpha_{i}\right)<\tau_{i}, \quad \forall i \end{array} \quad (4) {αi∗​}=αi​∈Aargmin​L(Wαi​∗​;Dval), s.t. FLOPs(αi​)<τi​,∀i​(4)
这里 $W^{\ast }$ 是在第1阶段学习的最佳权重共享参数。由于不需要重新训练或微调，因此该阶段的总体搜索成本通常很低。 此外，等式（4）自然地支持广泛的部署约束，而无需进一步修改，从而为机器学习从业人员提供了更加灵活的NAS框架。

Attentive Sampling NAS

NAS 的目标是在不同的计算约束下找到具有最佳准确性的网络体系结构。 尽管优化公式（2）中$\alpha \in \mathcal{A}$的平均损失似乎是很自然的选择，它并不是为提高任务性能和 DNN 资源使用之间的权衡而量身定制的。实际上，人们常常对形成最佳权衡的帕累托最优 DNN 感兴趣，如下图所示：

在等式（2）中调整无约束的预训练目标。 在等式（4）中寻求更好的解决方案。 文献中尚未针对两阶段的 NAS 进行探索。 直观地，一个简单的想法是将更多的训练预算投入可能构成帕累托最优集合的模型上，并用更多的数据和迭代来训练那些模型。 在实践中，增加训练预算已被证明是改善DNN性能的有效技术。

然而，改进表现最差的模型也可能很重要。突破帕累托最差集的性能极限可能会导致一个更好的优化的权重共享图，这样所有可训练的组件(如通道和层)，在最终的性能贡献中发挥其最大潜力。此外，改进 Pareto-worst 体系结构的原理与难例挖掘类似，将Pareto-worst子网络视为难数据例。它可以带来更多的信息梯度，并在体系结构空间中进行更好的探索，从而产生更好的性能。

感兴趣的子网络

Pareto-best architecture set

给定一个优化状态 W（权重共享图的参数），如果没有其他架构 $a^{\prime }\in \mathcal{A}$ 比 $\alpha$ 在获得更少或相似的计算开销的同时却可以获得更好的性能，则子网是 Pareto-best 最佳的。

Pareto-worst architecture set

同样，如果架构 $\alpha$ 在准确性上总是被其他具有相同或更大FLOP的架构所超越，则将其定义为帕累托最差的架构。

帕累托感知预训练

在公式(2)，所有候选 DNN 的优化概率相等。用帕累托感知目标重新表述(2)，使优化集中在帕累托最佳或最差集合上。这里，首先把预期损失改写成如下表示:
$$\underset{W}{\min }{\mathbb{E}}_{\pi (\tau )}{\mathbb{E}}_{\pi (\alpha \mid \tau )}\left[\mathcal{L}\left(W_{\alpha };{\mathcal{D}}^{trn}\right)\right](5)$$

容易看出，将等式（5）简化为等式（2）通过将 $\pi (\tau )$ 设置为搜索空间 $\mathcal{A}$ 和 $\pi (\alpha \mid \tau )$ 指定的 FLOP 的先验分布，并将其作为条件 FLOP 的体系结构的均匀分布。为简单起见，在这里删除了正则项 R（W）。

通过将 $\pi (\alpha \mid \tau )$ 设置为始终吸引 Pareto 最佳或最差体系结构的感知采样分布，可以进行 Pareto感知采样。 此优化目标制定如下：
$$\underset{W}{\min }{\mathbb{E}}_{\pi (\tau )}\sum _{\pi (\alpha \mid \tau )}\left[\gamma (\alpha )\mathcal{L}\left(W_{\alpha };{\mathcal{D}}^{trn}\right)\right](6)$$

当且仅当 $\alpha$ 是Pareto最佳子网（或Pareto最差子网）时， $\gamma (\alpha )$ 才定义为1，否则为0。

为了解决此优化问题，可以使用 n 个 Monte Carlo 样本的 FLOPs { τ o } \left\{\tau_{o}\right\} {τo​} 进行近似于 $\pi (\tau )$ 的期望。然后，对于每个目标FLOP ${\tau }_o$ ，可以用 k 个采样架构 { a 1 , ⋯   , a k } ∼ π ( α ∣ τ o ) \left\{a_{1}, \cdots, a_{k}\right\} \sim \pi\left(\alpha \mid \tau_{o}\right) {a1​,⋯,ak​}∼π(α∣τo​) 近似 $\mathrm{FLOPs}\left({\alpha }_i\right)={\tau }_o,\forall 1\le i\le k$ 如下:

$$\underset{W}{\min }\frac{1}{n}\sum _{{\tau }_o\sim \pi (\tau )}^n\left[\sum _{{\alpha }_i\sim \pi \left(\alpha \mid {\tau }_o\right)}^k\gamma \left({\alpha }_i\right)\mathcal{L}\left(W_{{\alpha }_i};{\mathcal{D}}^{trn}\right)\right](7)$$

令 $P(\alpha )$ 表示参数为 W 的模型 $\alpha$ 的性能估计。 如果目标是专注于Pareto最佳体系结构，则指定 $\gamma \left({\alpha }_i\right)=\mathbb{I}\left(P\left({\alpha }_i\right)>P\left({\alpha }_j\right),\forall j\ne i\right)$，其中 $\mathbb{I}(\cdot )$ 是指示函数。 如果目标是专注于Pareto最差架构，则将 $\gamma \left({\alpha }_i\right)=\mathbb{I}\left(P\left({\alpha }_i\right)<P\left({\alpha }_j\right),\forall j\ne i\right)$。

算法1 提供了基于感知抽样的 NAS 框架的元算法，称为AttentiveNAS。将总是选择性能最好的架构进行训练的采样策略表示为BestUp，总是选择性能最差的架构进行训练的采样策略表示为WorstUp。

性能评估器 $P(\alpha )$ 的选择是将其设置为验证损失，即 $P(\alpha )=-\mathcal{L}\left(W_{\alpha };{\mathcal{D}}^{\text{val }}\right)$。然而，由于验证集可能很大，这在计算上通常很昂贵。本文试验了许多可以有效计算的替代性能指标，包括由预先训练的精度预测器或小批量损失给出的预测精度。这将导致各种感知的架构采样实现，将在下面的实验部分中讨论。

实验结果

搜索空间

基于高效性能估算的感知采样

感知采样方法要求从一组候选抽样网络中选择最佳或最差的子网络。对验证集进行精确的性能评估在计算上是昂贵的。在本部分中，介绍两种有效的子网性能估计算法:

• Minibatch-loss 作为性能估计器:对于每个架构，使用当前小批训练数据上测量的训练损失作为代理性能度量;

• 准确性预测器作为性能评估:在验证集上训练准确性预测器;然后，对于每个体系结构，使用精度预测器给出的预测精度作为其性能评估。

第一种方法直观而直接。对于第二种方法，文献中广泛观察到，通过权重共享学习的不同子网之间的性能等级相关性在不同的运行过程中存在显著差异，导致肯德尔τ值极低。如果这仍然是两阶段NAS的情况，一个预先训练的准确性预测者不能很好地跨不同的设置。因此，首先了解候选子网在不同训练阶段和设置下的性能变化是很重要的。

训练精度预测器的设置

进行如下步骤:1)首先将原始训练数据集分割为90%的训练和10%的测试;2)对子样本训练集进行无约束的预训练。我们将训练限制为30个epoch，因此只引入不到整个两阶段NAS计算时间的10%。一旦训练完成，随机抽样1024个子网，并在次抽样测试数据划分上评估它们的性能;3)将1024对子网络及其精度划分为同等大小的训练和评估子集。使用带有100棵树的随机森林回归器作为精度预测器，并将最大深度设置为每棵树15。

精度预测器有效性的结果

对于所有测试子网，本文测量其预测精度和在次抽样测试数据集上测量的实际精度之间的秩相关(Kendall 's τ)。如下图所示，在评价子集上，预测精度与实际精度之间的Kendall 's τ为0.89，表明秩相关性非常高。

与 SOTA 的 NAS 结果比较

---