省略号
Ax 中的多目标优化使神经架构搜索(例如在模型性能和模型大小或延迟之间)的权衡探索变得高效。这种方法已在 Meta 的多种产品(如设备端 AI)中成功应用。在这篇文章中,我们提供了一个端到端的教程,让您可以亲身体验。
引言
神经网络在规模和复杂性上持续增长。开发最先进的架构通常是一个繁琐且耗时的过程,需要领域专业知识和大量的工程努力。为了克服这些挑战,已经提出了几种神经架构搜索(NAS)方法,旨在自动设计高性能的架构,而无需人工介入。
尽管非常样本效率低,但像随机搜索和网格搜索这样的朴素方法在超参数优化和神经架构搜索(NAS)中仍然很受欢迎(2019 年 NeurIPS 和 2020 年 ICLR 的一项研究发现,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,我们需要做的主要事情包括:
-
定义一个运行器,它负责将具有特定架构的模型发送到我们选择的平台(如 Kubernetes,或者可能只是我们本地机器上的 Docker 镜像)进行训练。在下面的教程中,我们使用 TorchX 来处理训练作业的部署。
-
定义一个指标,该指标负责从训练作业中获取目标指标(如准确率、模型大小、延迟)。在我们的教程中,我们使用 Tensorboard 记录数据,因此可以使用 Ax 捆绑的 Tensorboard 指标。
教程
在我们的教程中,我们展示了如何使用 Ax 在流行的 MNIST 数据集上对一个简单的神经网络模型运行多目标 NAS。虽然底层方法可以用于更复杂的模型和更大的数据集,但我们选择了一个易于在不到一小时内运行完整的教程。在我们的示例中,我们将调整两个隐藏层的宽度、学习率、dropout 概率、批大小和训练轮数。目标是使用多目标贝叶斯优化在性能(验证集上的准确率)和模型大小(模型参数数量)之间进行权衡。
本教程使用了以下 PyTorch 库:
-
PyTorch Lightning(指定模型和训练循环)
-
TorchX(用于远程/异步运行训练作业)
-
BoTorch(为 Ax 的算法提供动力的贝叶斯优化库)
完整的可运行示例可作为 PyTorch 教程获取。
结果
教程中执行的 NAS 优化最终结果可在下面的权衡图中看到。在这里,每个点对应一次试验的结果,颜色代表其迭代次数,星号表示我们强加给目标的阈值定义的参考点。我们看到,我们的方法能够成功地探索验证准确性和参数数量之间的权衡,找到了具有高验证准确性的大型模型以及具有较低验证准确性的小型模型。根据性能要求和模型大小限制,决策者现在可以选择使用哪个模型或进一步分析。
可视化
Ax 提供了一系列可视化工具,使我们能够分析和理解实验结果。在这里,我们将重点关注建模未知目标的 Gaussian 过程模型性能,这些模型用于帮助我们更快地发现有希望的配置。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 的 Adaptive Experimentation 团队对 Ax 和 BoTorch 的贡献。
参考文献列表
D. 艾里克森,P. 庄,S. 多尔顿,M. 巴兰达特。使用贝叶斯多目标神经架构搜索优化模型准确性和延迟。Meta 研究博客,2021 年 7 月。