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 = \{o_i\}_{i=1}^N$为包含$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$个实值架构参数$\{\alpha_i\}_{i=1}^N$应用softmax计算得到。

2.2 路径二值化的创新机制

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

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

$$g = \text{binarize}(p_1, \cdots, p_N) = \begin{cases} [1, 0, \cdots, 0] & \text{以概率 } p_1 \\ [0, 1, \cdots, 0] & \text{以概率 } p_2 \\ \vdots & \vdots \\ [0, 0, \cdots, 1] & \text{以概率 } p_N \end{cases}$$

基于二进制门$g$，混合操作的输出为：

$$m_O^{\text{Binary}}(x) = \sum_{i=1}^{N} g_i o_i(x)$$

由于$g_i \in \{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. 相应地重置路径权重$\{p_i\}$和二进制门$\{g_i\}$
4. 使用梯度公式更新这两条采样路径的架构参数
5. 重新缩放更新后的架构参数值，保持未采样路径的权重不变

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

3.1 延迟建模与可微化

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

将网络延迟建模为神经网络维度的连续函数。对于具有候选集$\{o_j\}$的混合操作，每个$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\|w\|_2^2 + \lambda_2 E[\text{latency}]$$

其中$\text{Loss}_{\text{CE}}$表示交叉熵损失，$\lambda_1\|w\|_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, 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 \neq 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^*$满足：

$$i^* = \arg\max_{i \in [N]} v_i$$

在两两比较中，如果$i^*$确实是最优的，则对于任意$j \neq i^*$：

$$P(i^* \text{ wins against } j) > P(j \text{ wins against } i^*)$$

通过多次两两比较，可以逐渐增强最优路径的权重。设第$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 \phi(x) + b$$

其中$\phi(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^* = \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]$是动量系数。