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