DARTS 与 ProxylessNAS 论文笔记

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

DARTS：https://github.com/khanrc/pt.darts
ProxylessNAS：https://github.com/MIT-HAN-LAB/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 堆叠成整个网络。

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

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

搜索完成后：

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

双层优化

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

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

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

此外也可以尝试用 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 上重新训练：

ProxylessNAS：直接架构搜索

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

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

解决内存问题

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

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

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

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

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

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

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