跳转到主要内容
博客

使用 Ax 进行高效的多目标神经架构搜索

作者: 2022年11月22日2024年11月15日暂无评论

简而言之

Ax 中的多目标优化能够在神经架构搜索中高效探索权衡(例如模型性能与模型大小或延迟之间的权衡)。此方法已在 Meta 成功应用于各种产品,例如设备端人工智能。在这篇文章中,我们提供了一个端到端教程,让您可以亲自尝试。

引言

神经网络在规模和复杂性方面持续增长。开发最先进的架构通常是一个繁琐且耗时的过程,需要领域专业知识和大量的工程投入。为了克服这些挑战,已经提出了几种神经架构搜索 (NAS) 方法,以自动设计性能良好的架构,而无需人工干预。

尽管效率低下,但像随机搜索和网格搜索这样的简单方法在超参数优化和 NAS 中仍然很流行(NeurIPS 2019 和 ICLR 2020 进行的一项研究发现,80% 的 NeurIPS 论文和 88% 的 ICLR 论文使用手动调优、随机搜索或网格搜索来调整其机器学习模型超参数)。但由于模型通常训练耗时且可能需要大量计算资源,因此最大限度地减少评估的配置数量非常重要。

Ax 是一个通用的黑盒优化工具,允许用户使用贝叶斯优化等最先进算法,以样本高效的方式探索大型搜索空间。在 Meta,Ax 被用于各种领域,包括超参数调优、NAS、通过大规模 A/B 测试识别最佳产品设置、基础设施优化以及设计尖端 AR/VR 硬件。

在许多 NAS 应用中,多个感兴趣的指标之间存在自然的权衡。例如,在设备上部署模型时,我们可能希望最大限度地提高模型性能(例如准确性),同时最大限度地减少相互竞争的指标,例如功耗、推理延迟或模型大小,以满足部署限制。在许多情况下,我们通过接受模型性能的微小下降,能够大幅减少计算要求或预测延迟(在某些情况下,我们能够同时提高准确性并减少延迟!)。高效探索这种权衡的原则性方法是可持续人工智能的关键推动力。

在 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 Research 博客,2021年7月。