神经架构搜索 (Neural Architecture Search, NAS) 是深度学习领域最令人振奋的前沿之一。试想,有一种算法可以自动为你设计出最先进的神经网络——并且完美契合你的特定任务。NAS 的愿景是用一种有原则的、自动化的搜索过程,取代繁琐且依赖直觉的人工网络设计。

然而,多年来,这一愿景伴随着巨大的代价。早期的 NAS 方法需要数万个 GPU 小时才能发现一个架构——成本高得令人望而却步,大多数研究人员和工程师都无法承担。为了让 NAS 可行,社区提出了一种巧妙的变通方法: 不直接在庞大的目标任务 (如 ImageNet) 上搜索,而是在更小、更容易管理的代理任务上进行搜索——例如用 CIFAR-10 代替 ImageNet、缩短训练周期,或者只搜索一个可复用的模块而非整张网络。

但这引出了一个关键问题:** 为小型代理任务优化的架构,对真实目标任务来说也是最优的吗?** 答案越来越倾向于*“不一定”*。在 CIFAR-10 上效果惊艳的设计,面对 ImageNet 的复杂性时可能并不理想。更重要的是,当你考虑到特定手机或 GPU 上的推理延迟等现实约束时,在代理任务上得到的架构可能远非最佳。

这正是《ProxylessNAS: 直接在目标任务和硬件上进行神经架构搜索》这篇论文要解决的问题。研究人员提出了一种彻底消除代理的方法,可以直接在像 ImageNet 这样的大规模任务上进行搜索,同时针对延迟等硬件特定指标优化。正如我们将看到的,他们的方法在仅需极少计算资源的情况下,取得了业界最前沿的成果。

旧的基于代理的方法与直接的 ProxylessNAS 方法之间的比较。代理方法有一个中间步骤,而 ProxylessNAS 直接在目标任务和硬件上学习。

图 1: 基于代理的 NAS 在将架构迁移到目标任务之前,依赖一个中间的代理任务。ProxylessNAS 移除了代理,直接在目标数据集和硬件上学习架构。


可微 NAS 的问题: 内存消耗大户

随着像 DARTS 这样的可微 NAS方法的引入,NAS 的效率实现了飞跃。其思想非常巧妙: 与其训练数千个独立的候选网络,不如创建一个巨大的、过参数化的网络,它包含了搜索空间中所有可能的架构

想象一下网络中的一层。在常规设计中,你会选择它是 3×3 卷积、5×5 卷积、池化层还是其他操作。在可微 NAS 中,你把所有这些选项都并行包含进来。这种“混合操作”会输出每个候选操作结果的加权和:

\[ m_{\mathcal{O}}^{\text{DARTS}}(x) = \sum_{i=1}^{N} \frac{\exp(\alpha_i)}{\sum_j \exp(\alpha_j)}\, o_i(x) \]

这里,\(o_i(x)\) 是第 \(i\) 个操作 (如 3×3 卷积) 的输出,权重 \(\alpha_i\) 是可学习的架构参数,经 softmax 归一化。搜索过程变为标准优化问题: 使用梯度下降同时训练网络权重和架构参数。

这将搜索成本从数万个 GPU 小时降到短短几天。但它也引入了一个致命瓶颈:** GPU 内存**。计算加权和需要计算并存储每个候选操作的输出特征图。如果每层有 7 个候选操作,内存占用约为普通网络的 7 倍。在 ImageNet 这样的大规模任务上,这种内存爆炸不可行——迫使研究人员回到代理任务。


ProxylessNAS 的核心思想: 路径二值化

ProxylessNAS 用一个看似简单的技巧解决了内存问题:** 路径二值化**。网络不再计算所有候选路径的加权和,而是在每次运行时只采样一条路径来激活。

架构参数 \(\alpha_i\) 仍然通过 softmax 转换为概率 \(p_i\)。但这些概率不再用作连续权重,而是用来采样一条激活路径,从而生成一个二元门控向量 \(g\):

\[ g = \mathrm{binarize}(p_1, \cdots, p_N) = \begin{cases} [1, 0, \cdots, 0] & \text{with prob.\ } p_1 \\ \vdots & \\ [0, 0, \cdots, 1] & \text{with prob.\ } p_N \end{cases} \]

通过二元门控,混合操作的输出变为:

\[ m_{\mathcal{O}}^{\text{Binary}}(x) = \sum_{i=1}^{N} g_i\, o_i(x) \]

训练过程中,网络在每个模块仅执行一个操作——将内存占用降到普通紧凑网络的水平。凭借这一点,即可在 ImageNet 等大型数据集上实现直接无代理搜索。


训练不可直接反传的架构: 学习二值化结构

采样离散路径带来了新挑战: 如何在这种随机过程中训练架构参数 \(\alpha_i\),因为无法直接对随机选择进行反向传播?

ProxylessNAS 通过交替的两步训练过程化解了这一问题:

ProxylessNAS 的两步训练过程: 一步更新采样架构的权重参数,另一步更新架构参数。

图 2: 权重参数与二值化架构参数交替更新,在保持内存占用低的同时,高效学习两者。

  1. 更新权重参数
    冻结架构参数。对每批训练数据,为每个混合操作模块采样一条激活路径 (组成该批次的紧凑网络) 。在训练集上用标准梯度下降更新网络权重。

  2. 更新架构参数
    冻结网络权重。用验证集,通过一种受 BinaryConnect 启发的梯度方法更新架构参数。核心梯度估计为:

    \[ \frac{\partial L}{\partial \alpha_i} \approx \sum_{j=1}^N \frac{\partial L}{\partial g_j} \, p_j \, (\delta_{ij} - p_i) \]

    为降低内存占用,每次更新只计算两条随机采样路径的梯度——直接比较并相应更新它们的参数。

通过交替执行这两步,ProxylessNAS 高效学得网络权重及其离散架构。


让硬件优化成为一等公民

ProxylessNAS 的一大亮点是可以直接针对硬件指标 (如延迟) 进行优化,而延迟通常是不可微的。研究人员提出了两种方法:

1. 延迟正则化损失

为每个目标硬件训练一个延迟预测模型。对于包含候选操作 \(o_j\) 和概率 \(p_j^i\) 的模块,其期望延迟为:

\[ \mathbb{E}[\text{latency}_i] = \sum_j p_j^i \times F(o_j^i) \]

其中 \(F(o_j^i)\) 是目标硬件上的预测延迟。因为总延迟只是各模块延迟的总和,该期望可微,可添加到损失函数中:

期望延迟根据不同操作的概率计算,并添加到损失函数中。

图 3: 将延迟建模为可微的正则项,引导搜索更快的架构。

\[ Loss = Loss_{CE} + \lambda_1 \|w\|_2^2 + \lambda_2\, \mathbb{E}[\text{latency}] \]

其中 \(\lambda_2\) 用于平衡精度与延迟。

2. 基于 REINFORCE 的替代方案

采用 REINFORCE 算法采样架构,评估其奖励 (精度和延迟结合) ,并据此更新概率。这无需延迟可微,灵活性更高。


结果: ProxylessNAS 的表现

CIFAR-10 上的顶尖成绩

在 CIFAR-10 上,基于梯度的 ProxylessNAS (Proxyless-G) 取得了 2.08% 测试错误率,参数量仅 5.7M——超过 AmoebaNet-B 的 2.13% 错误率,且参数量少 6 倍以上。

各模型在 CIFAR-10 上测试错误率与参数量对比。ProxylessNAS 以极少参数量取得最低错误率。

表 1: ProxylessNAS 在同类 NAS 模型中,以显著更少的参数量取得最佳性能。

移动设备场景下的 ImageNet

在 ImageNet 上,当移动端延迟限制低于 80 ms 时,ProxylessNAS 达到 74.6% top-1 精度,延迟仅 78 ms——超越 MobileNetV2 (72.0%) 和 MnasNet (74.0%),搜索耗时仅 200 GPU 小时,而 MnasNet 需 40,000 GPU 小时。

移动延迟限制下的 ImageNet 性能对比。ProxylessNAS 精度更高,搜索成本更低。

表 2: 面向移动设备的 ProxylessNAS 模型,在显著降低搜索成本的同时,超越人工设计和已有 NAS 架构。

Top-1 精度与移动延迟关系图。ProxylessNAS 整体曲线高于 MobileNetV2,表明其精度-延迟权衡更优。

图 4: 在不同延迟条件下,ProxylessNAS 始终优于 MobileNetV2,以更低延迟达到相同精度。

为了在无设备集群的前提下实现这一点,研究团队训练了精确且轻量的延迟预测模型——预测值与真实测量高度吻合。

预测延迟与真实测量延迟的强相关性散点图。

图 5: 预测延迟与真实世界测量高度相关,使设备感知 NAS 无需昂贵设备部署。


一刀切不可行: 为不同硬件定制模型

ProxylessNAS 针对三类平台搜索了专用架构: 手机、CPU 和 GPU。

不同硬件优化的模型在各平台上的性能对比。最佳成绩仅在目标平台上出现。

表 4: 针对特定硬件优化的模型仅在目标平台表现最佳,凸显专用化的必要性。

ProxylessNAS 为 GPU、CPU 和移动平台发现的架构。各模型结构模式明显不同。

图 6: 面向 GPU 的模型倾向浅而宽;面向 CPU 的模型更深更窄;移动设备则在两者间平衡。

架构洞察:

  • GPU 模型: 偏好浅而宽,采用大卷积核 (如 7×7 MBConv) 以充分利用并行性。
  • CPU 模型: 偏好深而窄,采用小卷积核以匹配顺序处理特性。
  • 移动模型: 在资源限制下寻求两者平衡。

结论: 实用 NAS 的新纪元

ProxylessNAS 是实用高效 NAS 的重大突破。路径二值化消除了可微 NAS 的内存瓶颈,首次实现直接在 ImageNet 等大规模数据集上进行架构搜索,成本仅相当于训练一张普通网络。

核心要点:

  1. 无需代理 —— 架构可直接在目标任务上搜索,缩小“代理鸿沟”。
  2. 搜索成本骤降 —— 降至数百 GPU 小时,与训练单个深度网络相当。
  3. 硬件优化必不可少 —— 最优架构取决于运行平台;ProxylessNAS 可自动找到专用设计。

这项工作让神经架构搜索走向普及,从资金雄厚的工业实验室专属工具,变为普通研究者和工程师也能掌握的技术。NAS 的范式由此转向生成多款专用模型——每个模型都为其部署环境量身打造。