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=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$个实值架构参数${\alphai}{i=1}^N$应用softmax计算得到。
2.2 路径二值化的创新机制
图2描述:该图详细展示了学习权重参数和二值化架构参数的过程。图分为两部分:左侧显示更新权重参数时,架构参数被冻结,二进制门控制哪条路径激活(例如CONV 3x3路径激活,其他路径不在内存中);右侧显示更新架构参数时,权重参数被冻结,不同的路径被激活用于梯度计算。图中用不同颜色区分了在内存中的特征图和不在内存中的特征图。
为了减少内存占用,ProxylessNAS在训练超参数化网络时只保留一条路径。引入$N$个实值架构参数${\alphai}{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)$$
由于$gi \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$个候选中选择一条路径的任务分解为多个二元选择任务。直觉是:如果某条路径在特定位置是最佳选择,那么当它与任何其他路径单独比较时,它应该是更好的选择。
在架构参数的每次更新步骤中:
- 根据多项分布$(p_1, \cdots, p_N)$采样两条路径
- 屏蔽所有其他路径,候选数量临时从$N$减少到2
- 相应地重置路径权重${p_i}$和二进制门${g_i}$
- 使用梯度公式更新这两条采样路径的架构参数
- 重新缩放更新后的架构参数值,保持未采样路径的权重不变
第三章 处理硬件指标的算法设计
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 延迟预测模型的构建
延迟预测使用以下特征:
- 操作类型(one-hot编码)
- 输入特征图尺寸:$(H{in}, W{in}, C_{in})$
- 输出特征图尺寸:$(H{out}, W{out}, C_{out})$
- 卷积核大小$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]$是动量系数。