tl;dr
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 设备端机器学习模型。除了 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 的贡献。
参考文献
D. Eriksson, P. Chuang, S. Daulton, M. Balandat。使用贝叶斯多目标神经网络架构搜索优化模型准确性和延迟。Meta 研究博客，2021 年 7 月。