ProxylessNAS：直接在目标任务和硬件上进行神经架构搜索

Cai H, Zhu L, Han S. Proxylessnas: Direct neural architecture search on target task and hardware[J]. arXiv preprint arXiv:1812.00332, 2018.

第一章 引言与研究背景

神经架构搜索（NAS）在自动化神经网络架构设计方面取得了显著成功，特别是在图像识别和语言建模等深度学习任务中。然而，传统NAS算法的计算成本极其高昂，需要在目标任务上训练数千个模型，单次实验就需要$10^4$ GPU小时。这种计算密集性使得直接将NAS应用于大规模任务（如ImageNet）变得极其昂贵甚至不可能。

作为折中方案，Zoph等人提出在代理任务上搜索构建块的策略——包括训练更少的epoch、从较小的数据集（如CIFAR-10）开始，或使用更少的块进行学习。然后将表现最好的块堆叠并迁移到大规模目标任务。这种范式被后续的NAS算法广泛采用。然而，这些在代理任务上优化的块无法保证在目标任务上也是最优的，特别是当考虑延迟等硬件指标时。更重要的是，为了实现可迁移性，这些方法需要只搜索少数几个架构模式然后重复堆叠相同的模式，这限制了块的多样性并损害了性能。

图1描述：该图展示了ProxylessNAS与传统基于代理的方法的对比。左侧显示传统方法需要先在代理任务上学习，然后迁移到目标任务和硬件；右侧显示ProxylessNAS直接在目标任务和硬件上优化架构。图中还包含了三种不同NAS方法的GPU小时数和GPU内存对比：普通训练的NAS需要元控制器和代理；DARTS和One-shot不需要元控制器但需要代理；ProxylessNAS既不需要元控制器也不需要代理。

第二章 方法论详解

2.1 超参数化网络的数学构建

将神经网络表示为有向无环图$\mathcal{N}(e_1,\cdots ,e_n)$，其中$e_i$代表图中的某条边。设O={oi}i=1NO = \{o_i\}_{i=1}^NO={oi​}i=1N​为包含$N$个候选原始操作的集合，包括各种卷积、池化、恒等映射和零操作。为了构建包含搜索空间中任何架构的超参数化网络，不将每条边设置为确定的原始操作，而是设置为具有$N$条并行路径的混合操作$m_O$。

在One-Shot方法中，给定输入$x$，混合操作$m_O$的输出定义为：
$$m_O^{\text{One-Shot}}(x)=\sum _{i=1}^N{o}_i(x)$$

在DARTS中，输出是加权和：
$$m_O^{\text{DARTS}}(x)=\sum _{i=1}^N{p}_i{o}_i(x)=\sum _{i=1}^N\frac{\exp ({\alpha }_i)}{{\sum }_{j=1}^N\exp ({\alpha }_j)}{o}_i(x)$$

其中权重通过对$N$个实值架构参数{αi}i=1N\{\alpha_i\}_{i=1}^N{αi​}i=1N​应用softmax计算得到。

2.2 路径二值化的创新机制

图2描述：该图详细展示了学习权重参数和二值化架构参数的过程。图分为两部分：左侧显示更新权重参数时，架构参数被冻结，二进制门控制哪条路径激活（例如CONV 3x3路径激活，其他路径不在内存中）；右侧显示更新架构参数时，权重参数被冻结，不同的路径被激活用于梯度计算。图中用不同颜色区分了在内存中的特征图和不在内存中的特征图。

为了减少内存占用，ProxylessNAS在训练超参数化网络时只保留一条路径。引入$N$个实值架构参数{αi}i=1N\{\alpha_i\}_{i=1}^N{αi​}i=1N​，然后将实值路径权重转换为二进制门：

$$g=\text{binarize}(p_1,\cdots ,p_N)=\{\begin{array}{ll}[1,0,\cdots ,0]& \text{以概率 }{p}_1\\ [0,1,\cdots ,0]& \text{以概率 }{p}_2\\ ⋮& ⋮\\ [0,0,\cdots ,1]& \text{以概率 }{p}_N\end{array}$$

基于二进制门$g$，混合操作的输出为：
$$m_O^{\text{Binary}}(x)=\sum _{i=1}^N{g}_i{o}_i(x)$$

由于gi∈{0,1}g_i \in \{0, 1\}gi​∈{0,1}且${\sum }_{i=1}^N{g}_i=1$，在任何时刻只有一条路径被激活。

2.3 二值化架构参数的梯度推导

架构参数不直接参与前向计算图，因此不能使用标准梯度下降更新。借鉴BinaryConnect的思想，使用$\partial L/\partial g_i$来近似$\partial L/\partial p_i$：

$$\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}$$

对于softmax函数，有：
$$p_j=\frac{\exp ({\alpha }_j)}{{\sum }_{k=1}^N\exp ({\alpha }_k)}$$

其对${\alpha }_i$的导数为：
$$\frac{\partial p_j}{\partial {\alpha }_i}=p_j({\delta }_{ij}-p_i)$$

其中${\delta }_{ij}$是Kronecker delta函数（当$i=j$时为1，否则为0）。

因此：
$$\frac{\partial L}{\partial {\alpha }_i}=\sum _{j=1}^N\frac{\partial L}{\partial g_j}{p}_j({\delta }_{ij}-p_i)$$

2.4 两路径采样的内存优化策略

直接计算$\partial L/\partial g_j$需要计算并存储$o_j(x)$，这仍然需要$N$倍的GPU内存。ProxylessNAS的关键创新是将从$N$个候选中选择一条路径的任务分解为多个二元选择任务。直觉是：如果某条路径在特定位置是最佳选择，那么当它与任何其他路径单独比较时，它应该是更好的选择。

在架构参数的每次更新步骤中：

1. 根据多项分布$(p_1,\cdots ,p_N)$采样两条路径
2. 屏蔽所有其他路径，候选数量临时从$N$减少到2
3. 相应地重置路径权重{pi}\{p_i\}{pi​}和二进制门{gi}\{g_i\}{gi​}
4. 使用梯度公式更新这两条采样路径的架构参数
5. 重新缩放更新后的架构参数值，保持未采样路径的权重不变

第三章 处理硬件指标的算法设计

3.1 延迟建模与可微化

图3描述：该图展示了如何通过引入延迟正则化损失使延迟可微。左侧显示了网络结构，包含一系列可学习块，每个块都有多个候选操作（CONV 3x3、CONV 5x5、Identity、POOL 3x3等）。右侧显示了延迟计算公式：每个块的期望延迟是各操作延迟的加权和，整个网络的期望延迟是所有块期望延迟的总和，最终损失函数包含交叉熵损失、权重衰减和延迟正则化项。

将网络延迟建模为神经网络维度的连续函数。对于具有候选集{oj}\{o_j\}{oj​}的混合操作，每个$o_j$关联路径权重$p_j$，表示选择$o_j$的概率。可学习块（即混合操作）的期望延迟为：

$$E[{\text{latency}}_i]=\sum _j{p}_j^i\times F(o_j^i)$$

其中$E[{\text{latency}}_i]$是第$i$个可学习块的期望延迟，$F(\cdot )$表示延迟预测模型，$F(o_j^i)$是$o_j^i$的预测延迟。

$E[{\text{latency}}_i]$对架构参数的梯度为：
$$\frac{\partial E[{\text{latency}}_i]}{\partial p_j^i}=F(o_j^i)$$

对于包含一系列混合操作的整个网络，由于这些操作在推理时顺序执行，网络的期望延迟表示为：
$$E[\text{latency}]=\sum _{i}E[{\text{latency}}_i]$$

最终的损失函数为：
$$\text{Loss}={\text{Loss}}_{\text{CE}}+{\lambda }_1\Vert w{\Vert }_2^2+{\lambda }_{2}E[\text{latency}]$$

其中${\text{Loss}}_{\text{CE}}$表示交叉熵损失，${\lambda }_1\Vert w{\Vert }_2^2$是权重衰减项，${\lambda }_2>0$控制准确性和延迟之间的权衡。

3.2 基于REINFORCE的替代算法

考虑具有二值化参数$\alpha$的网络，更新二值化参数的目标是找到最优二进制门$g$以最大化奖励$R(\cdot )$。根据REINFORCE算法：

$$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))$$

其中$g^i$表示第$i$个采样的二进制门，$p(g^i)$表示根据公式采样$g^i$的概率，${\mathcal{N}}_{g^i}$是根据二进制门$g^i$的紧凑网络。

第四章 实验结果与架构分析

4.1 实验设置

图4和图5描述：图4展示了ProxylessNAS在不同延迟设置下始终优于MobileNetV2的性能对比图。横轴是移动设备延迟（ms），纵轴是Top-1准确率（%）。图中显示ProxylessNAS的模型在各个延迟点都取得了更高的准确率，特别是在相同准确率下，ProxylessNAS比MobileNetV2快1.83倍。图5展示了移动延迟预测模型的准确性，横轴是估计延迟，纵轴是实际测量延迟，点基本沿着y=x线分布，延迟RMSE仅为0.75ms。

在CIFAR-10实验中，使用树形结构架构空间，以PyramidNet为骨干网络。将PyramidNet残差块中的所有$3\times 3$卷积层替换为树形结构单元，每个单元深度为3，每个节点的分支数设置为2（叶节点除外）。

在ImageNet实验中，使用MobileNetV2作为骨干构建架构空间。不是重复相同的移动倒残差卷积（MBConv），而是允许一组具有不同核大小{3,5,7}\{3, 5, 7\}{3,5,7}和扩展比{3,6}\{3, 6\}{3,6}的MBConv层。通过在候选集中添加零操作，允许带残差连接的块被跳过，实现宽度和深度之间的直接权衡。

4.2 专门化架构的硬件差异

图6描述：该图展示了ProxylessNAS为三种不同硬件平台（GPU、CPU、移动设备）搜索得到的高效模型架构。每个子图都显示了完整的网络结构，从输入到输出，包括各层的操作类型（如MB3 5x5表示扩展比为3的5×5 MBConv）和特征图尺寸。

GPU模型特点：

• 网络更浅更宽，特别是在早期阶段
• 偏好大型MBConv操作（如7×7 MBConv6）
• 在高分辨率特征图阶段使用更宽的通道

CPU模型特点：

• 网络更深更窄
• 偏好较小的MBConv操作（多为3×3）
• 延迟池化操作到网络后期

移动模型特点：

• 在深度和宽度之间取得平衡
• 混合使用不同大小的MBConv操作
• 在下采样位置使用较大的核

所有平台的共同模式是在每个阶段的第一个块（负责下采样）中偏好使用较大的MBConv操作，这可能有助于在下采样时保留更多信息。

4.3 性能对比

在CIFAR-10上，ProxylessNAS达到2.08%的测试误差，仅使用5.7M参数，比之前最佳的AmoebaNet-B（2.13%误差，34.9M参数）减少了6倍参数量。

在ImageNet上的结果：

• 移动设备：74.6% top-1准确率，78ms延迟
• GPU：75.1% top-1准确率，5.1ms延迟
• CPU：75.3% top-1准确率，138.7ms延迟

与MobileNetV2相比，ProxylessNAS在移动设备上将top-1准确率提高了2.6%，同时保持相似的延迟。与MnasNet相比，准确率高0.6%，但搜索成本减少了200倍（200 GPU小时 vs 40,000 GPU小时）。

附录 A：梯度计算

A.1 Softmax梯度

对于softmax函数：
$$p_j=\frac{\exp ({\alpha }_j)}{{\sum }_{k=1}^N\exp ({\alpha }_k)}$$

设分母为$Z={\sum }_{k=1}^N\exp ({\alpha }_k)$，则：
$$p_j=\frac{\exp ({\alpha }_j)}{Z}$$

当$i=j$时：
$$\frac{\partial p_j}{\partial {\alpha }_j}=\frac{\partial }{\partial {\alpha }_j}\left(\frac{\exp ({\alpha }_j)}{Z}\right)$$

使用商法则：
$$=\frac{Z\cdot \exp ({\alpha }_j)-\exp ({\alpha }_j)\cdot \exp ({\alpha }_j)}{Z^2}$$
$$=\frac{\exp ({\alpha }_j)}{Z}\cdot \frac{Z-\exp ({\alpha }_j)}{Z}$$
$$=p_j(1-p_j)$$

当$i\ne j$时：
$$\frac{\partial p_j}{\partial {\alpha }_i}=\frac{\partial }{\partial {\alpha }_i}\left(\frac{\exp ({\alpha }_j)}{Z}\right)$$
$$=-\frac{\exp ({\alpha }_j)\cdot \exp ({\alpha }_i)}{Z^2}$$
$$=-p_j{p}_i$$

综合两种情况：
$$\frac{\partial p_j}{\partial {\alpha }_i}=p_j({\delta }_{ij}-p_i)$$

A.2 两路径采样

设路径$i$的真实价值为$v_i$，在完整的$N$路径比较中，最优路径$i^{\ast }$满足：
$$i^{\ast }=\arg \underset{i\in [N]}{\max }{v}_i$$

在两两比较中，如果$i^{\ast }$确实是最优的，则对于任意$j\ne i^{\ast }$：
$$P(i^{\ast }\text{ wins against }j)>P(j\text{ wins against }{i}^{\ast })$$

通过多次两两比较，可以逐渐增强最优路径的权重。设第$t$次更新后路径$i$的权重为$p_i^{(t)}$，更新规则为：

如果路径$i$和$j$被采样，且$i$获胜：
$${\alpha }_i^{(t+1)}={\alpha }_i^{(t)}+\eta \cdot \frac{\partial L}{\partial g_i}$$
$${\alpha }_j^{(t+1)}={\alpha }_j^{(t)}-\eta \cdot \frac{\partial L}{\partial g_j}$$

其中$\eta$是学习率。经过足够多的更新后，最优路径的权重将收敛到最大值。

A.3 延迟预测模型的构建

延迟预测使用以下特征：

1. 操作类型（one-hot编码）
2. 输入特征图尺寸：$(H_{in},W_{in},C_{in})$
3. 输出特征图尺寸：$(H_{out},W_{out},C_{out})$
4. 卷积核大小$k$、步长$s$、扩展比$e$（对于MBConv）

预测模型使用简单的线性回归：
$$\text{latency}=w^T\varphi (x)+b$$

其中$\varphi (x)$是特征向量，$w$和$b$是学习的参数。

对于MBConv操作，延迟可以分解为三个部分：
$${\text{latency}}_{\text{MBConv}}={\text{latency}}_{\text{expand}}+{\text{latency}}_{\text{depthwise}}+{\text{latency}}_{\text{project}}$$

每个部分的延迟与计算量成正比：

• 扩展：$H_{in}\times W_{in}\times C_{in}\times (e\cdot C_{in})$
• 深度卷积：$H_{out}\times W_{out}\times (e\cdot C_{in})\times k^2$
• 投影：$H_{out}\times W_{out}\times (e\cdot C_{in})\times C_{out}$

A.4 REINFORCE梯度估计的方差减少

原始REINFORCE估计具有高方差。使用基线$b$减少方差：
$${\nabla }_{\alpha }J(\alpha )\approx \frac{1}{M}\sum _{i=1}^M(R({\mathcal{N}}_{g^i})-b){\nabla }_{\alpha }\log (p(g^i))$$

最优基线是奖励的期望值：
$$b^{\ast }={\mathbb{E}}_{g\sim \alpha }[R({\mathcal{N}}_g)]$$

实践中使用移动平均估计：
$$b^{(t+1)}=\beta b^{(t)}+(1-\beta )\cdot \frac{1}{M}\sum _{i=1}^{M}R({\mathcal{N}}_{g^i})$$

其中$\beta \in [0,1]$是动量系数。