摘要
Ax 中的多目标优化能够高效地探索神经网络架构搜索中的权衡（例如模型性能与模型大小或延迟之间的权衡）。这种方法已在 Meta 成功应用于各种产品，如设备端 AI。在这篇博文中，我们提供了一个端到端教程，让您可以亲自尝试。
引言
神经网络在规模和复杂性上持续增长。开发最先进的架构通常是一个繁琐且耗时的过程，既需要领域专业知识，也需要大量的工程工作。为了克服这些挑战，人们提出了多种神经网络架构搜索 (NAS) 方法，以无需人工干预的方式自动设计性能良好的架构。
尽管随机搜索和网格搜索等朴素方法样本效率很低，但它们在超参数优化和 NAS 中仍然很受欢迎（一项在 NeurIPS 2019 和 ICLR 2020 上进行的研究发现，80% 的 NeurIPS 论文和 88% 的 ICLR 论文使用手动调优、随机搜索或网格搜索来调整其机器学习模型超参数）。但由于模型训练通常耗时且可能需要大量计算资源，因此最小化评估配置的数量至关重要。
Ax 是一个通用的黑盒优化工具，它允许用户使用贝叶斯优化等最先进的算法以样本高效的方式探索大型搜索空间。在 Meta，Ax 被用于各种领域，包括超参数调优、NAS、通过大规模 A/B 测试确定最佳产品设置、基础设施优化以及设计尖端 AR/VR 硬件。
在许多 NAS 应用中，多个感兴趣的指标之间存在天然的权衡。例如，在设备端部署模型时，我们可能希望最大化模型性能（例如，准确率），同时最小化功耗、推理延迟或模型大小等竞争指标，以满足部署约束。在许多情况下，通过接受模型性能的轻微下降，我们能够大幅减少计算需求或预测延迟（在某些情况下，我们甚至能够同时提高准确率并减少延迟！）。有效探索此类权衡的原则性方法是可持续 AI 的关键推动因素。
在 Meta，我们成功地在 Ax 中使用了多目标贝叶斯 NAS 来探索此类权衡。我们的方法学正被常规用于优化 AR/VR 设备端 ML 模型。除了 NAS 应用，我们还开发了 MORBO，这是一种高维多目标优化方法，可用于优化增强现实 (AR) 的光学系统。
使用 Ax 实现全自动多目标 NAS
Ax 的调度器允许以闭环方式异步运行实验，通过持续将试验部署到外部系统、轮询结果、利用获取的数据生成更多试验，并重复该过程直到满足停止条件。无需人工干预或监督。调度器的功能包括：
-
可定制并行度、容错能力和许多其他设置；
-
大量最先进的优化算法可供选择；
-
保存进行中的实验（到 SQL 数据库或 json）并从存储中恢复实验；
-
易于扩展到新的后端，以便远程运行试验评估。
以下来自Ax 调度器教程的插图总结了调度器如何与用于运行试验评估的任何外部系统交互：
要使用调度器运行自动化 NAS，主要需要做以下几件事：
-
定义一个 Runner（运行器），它负责将特定架构的模型发送到我们选择的平台（如 Kubernetes，或者可能只是本地机器上的 Docker 镜像）进行训练。在下面的教程中，我们使用 TorchX 来处理训练作业的部署。
-
定义一个 Metric（度量），它负责从训练作业中获取目标指标（如准确率、模型大小、延迟）。在我们的教程中，我们使用 Tensorboard 记录数据，因此可以使用 Ax 捆绑的 Tensorboard 指标。
教程
在我们的教程中，我们将展示如何使用 Ax 对流行的 MNIST 数据集上的简单神经网络模型运行多目标 NAS。虽然底层方法学可用于更复杂的模型和更大的数据集，但我们选择了一个可以在笔记本电脑上轻松端到端运行且耗时少于一小时的教程。在我们的示例中，我们将调整两个隐藏层的宽度、学习率、dropout 概率、批量大小和训练 epoch 数。目标是使用多目标贝叶斯优化来权衡性能（验证集上的准确率）和模型大小（模型参数数量）。
本教程使用了以下 PyTorch 库：
-
PyTorch Lightning（用于指定模型和训练循环）
-
TorchX（用于远程/异步运行训练作业）
-
BoTorch（为 Ax 算法提供支持的贝叶斯优化库）
完整的可运行示例作为PyTorch 教程提供。
结果
在教程中进行的 NAS 优化最终结果可以在下面的权衡图中看到。图中每个点对应一个试验结果，颜色表示其迭代次数，星号表示我们对目标施加的阈值定义的参考点。我们看到，我们的方法能够成功探索验证准确率和参数数量之间的权衡，找到了验证准确率高的大模型，也找到了验证准确率低的小模型。根据性能要求和模型大小约束，决策者现在可以选择要使用的模型或进行进一步分析。
可视化
Ax 提供了多种可视化功能，可以用来分析和理解实验结果。在这里，我们将重点介绍模拟未知目标的高斯过程模型的性能，这些模型用于帮助我们更快地发现有前景的配置。Ax 使通过留一交叉验证更容易理解这些模型的准确性以及它们在未见数据上的表现。在下面的图中，我们看到模型拟合看起来相当不错——预测值与实际结果非常接近，并且预测的 95% 置信区间很好地覆盖了实际结果。此外，我们观察到模型大小 (num_params) 指标比验证准确率 (val_acc) 指标更容易建模。
要点总结
-
我们展示了如何使用 Ax 运行全自动多目标神经网络架构搜索。
-
使用 Ax 调度器，我们能够以完全异步的方式自动运行优化——这可以在本地完成（如教程所示），也可以通过将试验远程部署到集群来完成（只需更改 TorchX 调度器配置）。
-
Ax 中提供的最先进的多目标贝叶斯优化算法使我们能够有效地探索验证准确率和模型大小之间的权衡。
高级功能
Ax 还具有本教程中未讨论的许多其他高级功能。其中包括：
提前停止
在评估新的候选配置时，神经网络训练作业运行时通常会提供部分学习曲线。我们可以利用部分曲线中包含的信息来识别表现不佳的试验并提前停止，以便为更有前景的候选者释放计算资源。虽然在上面的教程中没有演示，但 Ax 原生支持提前停止功能。
高维搜索空间
在我们的教程中，我们使用了带有标准高斯过程的贝叶斯优化以保持较低的运行时间。然而，这些模型通常只能扩展到大约 10-20 个可调参数。我们新的 SAASBO 方法（论文、Ax 教程、BoTorch 教程）样本效率很高，可以调优数百个参数。只需将 use_saasbo=True 传递给 choose_generation_strategy 即可轻松启用 SAASBO。
致谢
我们感谢 TorchX 团队（特别是 Kiuk Chung 和 Tristan Rice）在将 TorchX 与 Ax 集成方面的帮助，以及 Meta 的自适应实验团队对 Ax 和 BoTorch 的贡献。
