在现实世界中部署机器学习模型是一项棘手的任务。对于高端云 GPU 来说完美无缺的模型,可能完全不适合智能手机;而适合智能手机的模型,对于微型控制器来说又显得性能过剩。每种设备都有其独特的约束——延迟、内存和功耗——这种多样性催生了神经架构搜索 (Neural Architecture Search, NAS) 的迅速发展,该领域致力于自动设计针对特定硬件的神经网络。
尽管 NAS 已取得令人瞩目的成果,但它通常伴随着一个隐藏的代价: 复杂性与高昂的计算成本。一种现代流行的 一次性 NAS (one-shot NAS) 方法试图通过训练一个巨大的“超网络”来简化这一过程,这个超网络内部包含了无数个较小的“子”架构。寻找最佳架构因此变得像是在这个超网络中找到最优路径一样简单。
然而,问题在于: 直接从超网络中提取的子模型,其性能往往远逊于从头训练的同等模型。它们的共享权重是在迥然不同的架构之间的妥协,很少能对任何一个特定架构做到最优。这导致了 NAS 领域中根深蒂固的“规则”:** 你必须将发现的架构从头开始重新训练,或者应用复杂的后处理才能获得良好性能**。这个最终步骤可能成为巨大的瓶颈,尤其是在需要为各种设备提供模型时。
但如果这个规则是错的呢?如果你可以只训练一个巨大的模型,然后即时切分出任何尺寸的高性能、可直接部署的子模型,而无需任何重新训练呢?
这正是来自 Google Brain 的研究人员提出的 BigNAS 背后的颠覆性理念。他们挑战了传统观念,并展示了如何训练一个单一模型,使其能够作为适配各种计算预算的通用最先进模型的来源。
图 1: BigNAS 的工作流程 (右) 相比以往方法简化了部署。不再需要复杂、多步骤的重新训练或蒸馏流程,BigNAS 通过训练一个单一模型,可即时切分出各种规模的子模型。
一次性 NAS 的问题所在
在标准的一次性 NAS 工作流程中,你会定义一个搜索空间——一组可能的配置,如卷积核尺寸、网络深度和通道宽度。然后构建一个包含所有这些可能性的超网络。在训练过程中,对不同的子架构进行采样,并更新它们共享的权重。
这种方法很优雅: 你只需训练一组权重。但它也极具挑战性。一套参数必须同时适用于微小、浅层的网络以及庞大、深层的网络——这些架构的学习动态完全不同。结果就是一种妥协: 共享的权重对任何特定子模型都不是最优的。切分出的模型准确率通常只是个近似值,因此不得不进行代价高昂的从头再训练。
BigNAS 的目标就是消除这一差距。旨在训练出一个单阶段模型,使得任意切分出的子模型都已接近最佳性能。
核心方法: 训练高质量的单阶段模型
核心挑战在于平衡冲突的需求: 最小的模型需要更高的学习率和更少的正则化才能有效学习,而最大的模型则需要更低的学习率和更多的正则化来避免过拟合。
BigNAS 引入了五项关键技术来调和这些差异。
1. 三明治法则 (The Sandwich Rule) — 界定性能边界
训练过程同等关注搜索空间的两端和中间段。
在每一步中,BigNAS 不仅采样随机架构,还会始终采样:
- 最大模型 — 最大深度、最大宽度、最大卷积核、最大分辨率
- 最小模型 — 各项参数均取最小值
- 随机模型 — 若干介于其间的不同架构
所有这些模型的梯度会聚合起来更新权重。通过明确训练最大与最小模型,该过程确保性能范围的两端同时得到提升。
2. 原位蒸馏 (Inplace Distillation) — 大模型指导小模型
知识蒸馏利用较大的“教师”网络预测的输出概率 (软标签) 来训练较小的“学生”网络。
在 BigNAS 中,得益于三明治法则,最大模型在每个训练步骤中都作为教师。它根据真实标签进行学习,其预测结果用于监督同一批次的所有其他模型。
作者确保教师和学生看到的是同一块图像,但会调整为各自的分辨率。这让监督信号更加一致,并为子模型带来了 +0.3% 的准确率提升。
3. 智能初始化 (Smart Initialization) — 抑制损失爆炸
在最初训练大型单阶段模型时会出现不稳定的损失爆炸。降低学习率能稳定训练,但会导致准确率下降 (约 –1.0% top-1) 。
解决办法很巧妙: 将每个残差块的最后一个 BatchNorm 层输出,用一个可学习的缩放参数 \(\gamma = 0\) 初始化。这意味着残差路径起始时实际上为零,让跳跃连接成为主要信号通道。
网络从一个“更简单”的状态开始,训练过程中可以逐步利用残差路径。这稳定了学习过程,使得可以使用更高的学习率,并将 ImageNet 准确率提高了约 +1.0%。
4. 改进学习率策略 — 兼顾不同模型收敛
小模型收敛较慢,而大模型会过早达到性能峰值并开始过拟合。
图 2: (a) 大模型 (橙色) 在小模型 (蓝色) 完成学习前就已达到峰值。(b) 所提出的学习率策略先指数衰减,再保持一个较小的恒定值。
改进后的策略,即**“指数衰减并以恒定值收尾”**,先按正常方式衰减学习率,但在降至初始值的 5% 时停止,并保持该值不变。
好处:
- 小模型获得额外的训练空间以完成收敛
- 大模型在训练末期权重出现轻微振荡,起到正则化作用,减少过拟合
5. 定向正则化 (Targeted Regularization) — 少即是多
大模型容易过拟合;小模型容易欠拟合。对两者施加相同的正则化 (dropout、权重衰减) 并非最佳选择。
BigNAS 只对最大模型进行正则化。小模型则不加正则化,以最大限度提升其数据拟合能力。
这样小模型的准确率提升了 +0.5%,甚至大模型也获得了 +0.2% 的微弱提升。
批归一化校准 (Batch Norm Calibration)
训练结束后,最后一步是: 为任意选定的子模型重新计算批归一化统计,只需用几百个批次的数据通过该模型。这一过程很快,不需要更新权重,并能确保推理时性能稳定。
寻找最佳架构
在庞大 (超过 \(10^{12}\) 种可能) 的搜索空间中寻找最优子模型需要高效的策略。BigNAS 采用从粗到细 (coarse-to-fine) 的搜索:
- 粗搜索 — 在一个小规模的全局设置网格 (分辨率、深度、宽度、卷积核尺寸) 上扫描,找到潜力“骨架”架构。
- 细搜索 — 对每个潜力骨架,局部调整细节,如按阶段调整宽度或单个卷积核的大小,以锁定预算内的最佳架构。
图 8: 左 — 粗选定位到好的候选 (红点) 。右 — 细搜索在该候选附近优化,找到预算内的最优配置。
由于 BigNAS 子模型可直接部署,这里测得的准确率即为它们的最终准确率。
在 ImageNet 上的惊人成绩
BigNAS 单阶段模型在 ImageNet 上的训练覆盖了从约 200 MFLOPs (类似 MobileNetV3-Small) 到约 1 GFLOPs (类似 EfficientNet-B2) 的架构范围。团队由此切分出四个代表性模型: S、M、L、XL。
图 3: BigNAS 模型 (红线) 在各种计算预算下的 top-1 准确率均高于此前的模型。
亮点:
- BigNASModel-M — 在约 400 MFLOPs 下达到 78.9% top-1 准确率 (比 EfficientNet-B0 高 1.6%)
- BigNASModel-S — 在约 240 MFLOPs 下达到 76.5% (比 MobileNetV3 高 1.3%)
- BigNASModel-XL — 在约 1 GFLOPs 下达到 80.9%,比 ResNet-50 更准确且计算成本低约 4 倍
以上成果均无需重新训练或微调。
消融研究: 方法为何奏效
五项技术均通过消融实验得以验证。
初始化: 没有改进的初始化方法,训练在高学习率下会不稳定;使用后,训练更快且准确率更高。
图 5: 使用 100% 学习率的改进初始化 (红色) 对小模型 (左) 和大模型 (右) 都取得最佳效果。
正则化: 仅对最大模型施加 dropout/权重衰减可提升大小模型的性能。
图 7: 定向正则化提高了小模型 (左) 和大模型 (右) 的准确率。
微调有帮助吗?
如果我们对切分的 BigNAS 模型进行微调会怎样?
结果显示,提升微乎其微,有时甚至会降低准确率。
表 2: 微调通常无法改善 BigNAS 子模型性能,表明它们已接近最优状态。
这证实了 BigNAS 子模型已处于性能巅峰。
结论 — 一次训练,随处部署
BigNAS 提供了一种精简且可扩展的 NAS 方法。通过直接应对大型权重共享超网络独特的训练动态,它消除了耗时且昂贵的搜索后再训练步骤。
最终成果是一个经过专业训练的单一模型,可以视作“一个权重集内的模型动物园”:
- 需要一个用于微控制器的紧凑模型?切分即可。
- 需要一个用于旗舰手机的更大更快的模型?仍从同一父模型切分即可。
这一创新不仅带来了一个家族的最先进模型,还为 NAS 指明了更高效的未来方向: 真正实现一次训练,随处部署。