DARTS 与 ProxylessNAS 论文笔记

最近在研究神经网络架构搜索（NAS），精读了 DARTS 和 ProxylessNAS 两篇论文，在此整理一下笔记。两篇论文的代码都已开源：

DARTS：可微分架构搜索

核心思路

DARTS 的基本想法是：用 softmax 将离散的操作选择松弛成连续的加权混合，再对结构参数 α 求微分，利用梯度下降来优化网络结构。这样就把原本不可微的离散搜索问题转化成了连续优化问题。

操作集

DARTS 的候选操作共 8 种：

  • 3×3 separable convolution
  • 5×5 separable convolution
  • 3×3 dilated separable convolution
  • 5×5 dilated separable convolution
  • 3×3 max pooling
  • 3×3 average pooling
  • identity（直连）
  • zero（不连接）

Cell 的结构

DARTS 同样只搜索 cell，然后将 cell 堆叠成整个网络。

DARTS 概览四子图：(a) 边上操作未知，(b) 在每条边上放置候选操作的混合做连续松弛，(c) 联合优化混合概率与网络权重，(d) 由学到的概率导出最终结构
图 1：DARTS 概览。(a) 边上的操作初始未知；(b) 在每条边上放置候选操作的混合，对搜索空间做连续松弛；(c) 通过求解双层优化，联合优化混合概率与网络权重；(d) 从学到的混合概率中导出最终架构。

一个 cell 是多个节点组成的有向无环图，有 2 个输入节点和 1 个输出节点。输出节点是所有中间节点输出的 concat。论文中每个 cell 取 7 个节点，去掉输入输出节点后共有 4 个中间节点。

搜索阶段，cell 是一张全排列的有向图——任意两节点之间都有一条「混合操作」边，混合操作是所有候选操作输出的加权和，权重由 α 经 softmax 归一化得到。这样一来，所有操作的 feature map 和参数都需要同时保留在内存中，内存开销较大。ProxylessNAS 针对这一点做了改进，它不用混合操作，而是直接根据概率采样一个具体的操作。

搜索完成后：

  • 对每个中间节点，根据指向它的所有边的「强度」（定义为该边上权重最大的操作所对应的权重值），保留强度最大的 2 条边。
  • 对每条保留边，argmax 取权重最大的单一操作，得到最终的离散结构。

双层优化

梯度迭代是权重 w 和结构参数 α 交叉更新的：w 在训练集上更新，α 在验证集上更新。论文把 CIFAR 的训练集平均分成两份，分别充当这里的 train 和 val。

结构参数的梯度计算公式如下：

∇αLval(w′,α)−ξ ∇α,w2Ltrain(w,α) ∇w′Lval(w′,α)\nabla_\alpha \mathcal{L}_{val}(w', \alpha) - \xi\, \nabla^2_{\alpha,w}\mathcal{L}_{train}(w, \alpha)\, \nabla_{w'} \mathcal{L}_{val}(w', \alpha)

其中 w′=w−ξ∇wLtrain(w,α)w' = w - \xi \nabla_w \mathcal{L}_{train}(w,\alpha)。

其中 Hessian 矩阵用一种近似算法来快速计算：

∇α,w2Ltrain(w,α) ∇w′Lval(w′,α)≈∇αLtrain(w+,α)−∇αLtrain(w−,α)2ϵ\nabla^2_{\alpha,w}\mathcal{L}_{train}(w, \alpha)\, \nabla_{w'}\mathcal{L}_{val}(w', \alpha) \approx \frac{\nabla_\alpha \mathcal{L}_{train}(w^{+}, \alpha) - \nabla_\alpha \mathcal{L}_{train}(w^{-}, \alpha)}{2\epsilon}

其中 w±=w±ϵ ∇w′Lval(w′,α)w^{\pm} = w \pm \epsilon\, \nabla_{w'}\mathcal{L}_{val}(w', \alpha)。

此外也可以尝试用 REINFORCE 方法来计算这样的梯度，基于蒙特卡洛采样来估计期望。

Normal Cell 与 Reduction Cell

搜索阶段只搜出一种 cell 结构，但通过 normal cell 和 reduction cell 两种用途加以区分：normal cell 保持特征图尺寸不变，reduction cell 则将 cell 内输入节点的第一个操作改为 stride=2，起下采样作用。在完整网络中，1/3 和 2/3 深度位置放置 reduction cell，其余位置放置 normal cell。

搜索成本与 ImageNet 迁移

一次完整搜索耗时约 1 天（单张 GTX 1080 Ti）。论文为了实验稳定性，用了 4 个不同随机数种子各跑一次，总共 4 天。

在 CIFAR 上搜索完成后，将搜到的结构迁移到 ImageNet 上重新训练：

ArchitectureTest Error (%)Params (M)Search Cost (GPU days)Search Method
DenseNet-BC (Huang et al., 2017)3.4625.6–manual
NASNet-A + cutout (Zoph et al., 2017)2.653.31800RL
NASNet-A + cutout (Zoph et al., 2017)†2.833.13150RL
AmoebaNet-A + cutout (Real et al., 2018)3.34 ± 0.063.23150evolution
AmoebaNet-A + cutout (Real et al., 2018)†3.123.13150evolution
AmoebaNet-B + cutout (Real et al., 2018)2.55 ± 0.052.83150evolution
Hierarchical Evo (Liu et al., 2017b)3.75 ± 0.1215.7300evolution
PNAS (Liu et al., 2017a)3.41 ± 0.093.2225SMBO
ENAS + cutout (Pham et al., 2018b)2.894.60.5RL
Random + cutout3.493.1––
DARTS (first order) + cutout2.942.91.5gradient-based
DARTS (second order) + cutout2.83 ± 0.063.44gradient-based
表 1：在 CIFAR-10 上与当前先进图像分类器的对比。标注 † 的结果是用本文的训练设置重新训练对应架构得到的。
ArchitectureTest Error top-1 (%)Test Error top-5 (%)Params (M)+× (M)Search Cost (GPU days)Search Method
Inception-v1 (Szegedy et al., 2015)30.210.16.61448–manual
MobileNet (Howard et al., 2017)29.410.54.2569–manual
ShuffleNet 2× (v1) (Zhang et al., 2017)29.110.2~5524–manual
ShuffleNet 2× (v2) (Zhang et al., 2017)26.3–~5524–manual
NASNet-A (Zoph et al., 2017)26.08.45.35641800RL
NASNet-B (Zoph et al., 2017)27.28.75.34881800RL
NASNet-C (Zoph et al., 2017)27.59.04.95581800RL
AmoebaNet-A (Real et al., 2018)25.58.05.15553150evolution
AmoebaNet-B (Real et al., 2018)26.08.55.35553150evolution
AmoebaNet-C (Real et al., 2018)24.37.66.45703150evolution
PNAS (Liu et al., 2017a)25.88.15.1588~225SMBO
DARTS (searched on CIFAR-10)26.99.04.95954gradient-based
表 3：在 mobile 设定下与当前先进图像分类器在 ImageNet 上的对比。

ProxylessNAS：直接架构搜索

ProxylessNAS 感觉主要是在 DARTS 的基础上进行改进，同时也做了更加全面的实验。几个重要贡献如下：

  1. 处理了 DARTS 太耗内存的问题
  2. 从 CIFAR 到 ImageNet 并没有使用 transfer，而是直接在 ImageNet 上搜索
  3. 在更新结构参数上提出了两种计算梯度的方法：BinaryConnect 方式和 REINFORCE 方式
  4. 针对移动平台做了多目标搜索的实验，把 latency 也作为优化目标之一，并且将 latency 用一个神经网络来预测，从而使其可微分

解决内存问题

DARTS 使用混合操作，所有操作后的 feature map 都会保存起来，内存消耗很大。ProxylessNAS 的思路是：

  • 更新 w 时：冻结结构参数 α，根据 α 计算每个操作的采样概率 p，然后采样得到一个确定的单路径结构，只更新该路径的参数。
  • 更新 α 时：冻结 w，计算 α 的梯度然后梯度下降。

节点间的连接方式和 DARTS 一样，每个中间节点取两个前置节点作为输入，输出节点是所有中间节点输出的 concat。生成最终结构的方式也和 DARTS 一样，具体操作取权值最大的那个。

但这样还是会有内存问题，因为在计算 α 的梯度时，需要用到如下的二值化权重：

∂L/∂gi\partial L / \partial g_i

每条边的操作由下列公式确定，g 与 o 都保留在计算图中——即使本质上是采样，实现方式还是混合权值，只是将权值二值化了，因此和 DARTS 的内存消耗差不多：

mOBinary(x)=∑i=1Ngi oi(x)={o1(x)with probability p1⋯oN(x)with probability pN.m^{\mathrm{Binary}}_{\mathcal{O}}(x) = \sum_{i=1}^{N} g_i\, o_i(x) = \begin{cases} o_1(x) & \text{with probability } p_1 \\ \quad\cdots & \\ o_N(x) & \text{with probability } p_N. \end{cases}

更有效的方案是：更新 α 时，不把全部 N 个操作放进计算图，而是根据 α 从 N 个操作中做多项采样，采样出 2 个操作，将 N 临时压缩到 2，再用上面的实现方式进行参数更新。这相当于 α 的更新不是针对所有操作一起进行，而是每次以 2 个操作为一组进行更新。

为什么不直接采样 1 个操作？因为采样操作本身不在计算图中——采样 1 个时无法计算 α 的梯度，采样 2 个才有办法形成对比，计算出 α 的梯度方向。更新完这 2 个操作对应的 α 之后，在所有 α 上做一次缩放，让其他 α 对应的权重（α 经 softmax 得到）保持不变。经过多次两两比较的累积，如果一个操作比所有操作都合适，那么在两两比较中它也能持续保留下来。

直接在 ImageNet 上搜索

得益于内存优化，ProxylessNAS 可以直接在 ImageNet 上搜索，不再依赖 CIFAR 作为代理。这里不是把 cell 直接堆叠起来，而是用 MobileNet V2 作为骨干框架，搜索其中每个 block 的操作选择。

两种结构参数梯度估计方法

方法一：BinaryConnect 方式

借鉴 BinaryConnect 中对离散权重求梯度的思路：

∂L∂αi=∑j=1N∂L∂pj∂pj∂αi≈∑j=1N∂L∂gj∂pj∂αi=∑j=1N∂L∂gj∂ ⁣(exp⁡(αj)∑kexp⁡(αk))∂αi=∑j=1N∂L∂gj pj(δij−pi)\frac{\partial L}{\partial \alpha_i} = \sum_{j=1}^{N} \frac{\partial L}{\partial p_j}\frac{\partial p_j}{\partial \alpha_i} \approx \sum_{j=1}^{N} \frac{\partial L}{\partial g_j}\frac{\partial p_j}{\partial \alpha_i} = \sum_{j=1}^{N} \frac{\partial L}{\partial g_j}\frac{\partial\!\left(\frac{\exp(\alpha_j)}{\sum_k \exp(\alpha_k)}\right)}{\partial \alpha_i} = \sum_{j=1}^{N} \frac{\partial L}{\partial g_j}\, p_j(\delta_{ij} - p_i)

方法二：REINFORCE 方式

用蒙特卡洛采样来估计期望梯度：

J(α)=Eg∼α[R(Ng)]=∑ipi R(N(e=oi)),∇αJ(α)=∑iR(N(e=oi)) ∇αpi=∑iR(N(e=oi)) pi ∇αlog⁡(pi)=Eg∼α[R(Ng) ∇αlog⁡(p(g))]≈1M∑i=1MR(Ngi) ∇αlog⁡(p(gi))\begin{aligned} J(\alpha) &= \mathbb{E}_{g\sim\alpha}[R(\mathcal{N}_g)] = \sum_i p_i\, R(\mathcal{N}(e = o_i)), \\ \nabla_\alpha J(\alpha) &= \sum_i R(\mathcal{N}(e = o_i))\, \nabla_\alpha p_i = \sum_i R(\mathcal{N}(e = o_i))\, p_i\, \nabla_\alpha \log(p_i) \\ &= \mathbb{E}_{g\sim\alpha}[R(\mathcal{N}_g)\, \nabla_\alpha \log(p(g))] \approx \frac{1}{M}\sum_{i=1}^{M} R(\mathcal{N}_{g^i})\, \nabla_\alpha \log(p(g^i)) \end{aligned}

上面这个公式中，计算梯度的部分可以直接用 softmax 的解析式计算出来。

将 Latency 纳入可微分优化

在手机上，利用一些特征（操作类型、输入输出 feature map 的尺寸、其他超参数）来训练一个小型神经网络预测 latency。这样 latency 就变成了关于 α 的可微分函数，可以和精度目标一起纳入多目标优化。

多目标搜索的实验结果：

ModelTop-1Top-5Mobile LatencyHardware-awareNo ProxyNo RepeatSearch cost (GPU hours)
MobileNetV1 [16]70.689.5113ms--✗Manual
MobileNetV2 [28]72.091.075ms--✗Manual
NASNet-A [34]74.091.3183ms✗✗✗48,000
AmoebaNet-A [27]74.592.0190ms✗✗✗75,600
MnasNet [29]74.091.876ms✓✗✗40,000
MnasNet (our impl.)74.091.879ms✓✗✗40,000
Proxyless-G (mobile)71.890.383ms✗✓✓200
Proxyless-G + LL74.291.779ms✓✓✓200
Proxyless-R (mobile)74.692.278ms✓✓✓200
表 2：ProxylessNAS 在 mobile latency 约束（≤80ms）下于 ImageNet 上达到 SOTA 精度（%），搜索成本（GPU 小时）仅为其他方法的 1/200。「LL」表示 latency 正则化损失。
ModelTop-1Top-5GPU latency
MobileNetV2 (Sandler et al., 2018)72.091.06.1ms
ShuffleNetV2 (1.5) (Ma et al., 2018)72.6-7.3ms
ResNet-34 (He et al., 2016)73.391.48.0ms
NASNet-A (Zoph et al., 2018)74.091.338.3ms
DARTS (Liu et al., 2018c)73.191.0-
MnasNet (Tan et al., 2018)74.091.86.1ms
Proxyless (GPU)75.192.55.1ms
表 3：ImageNet 精度（%）与 GPU latency（Tesla V100）。

这个实验用了多个优化目标，除了精度之外还有 latency，在精度上直接达到了 SOTA。最终的模型效果：

散点图：横轴为预测 latency（ms）、纵轴为实测 latency（ms），红色散点紧贴 y=x 虚线
图 2：移动端 latency 预测模型的效果，预测值与实测值非常接近 y=x（latency RMSE 为 0.75ms）。
  • 论文笔记
  • NAS
  • 深度学习