通过Ax进行高效优化:一个用于自适应实验的开放平台
我们发布了 Ax 1.0,这是一个利用机器学习自动引导复杂、资源密集型实验的开源平台。Ax在Meta内部被大规模用于改进AI模型、调整生产基础设施以及加速机器学习和硬件设计领域的进展。我们的配套论文《Ax:一个自适应实验平台》解释了Ax的架构、方法论,以及它与其他先进的黑盒优化库相比如何。
研究人员如何有效地理解和优化那些拥有海量可能配置的AI模型或系统?这在现代AI开发与部署等以复杂、相互作用的系统为特征的领域中,是一个尤其普遍的挑战。在这些场景下进行优化需要实验,而当评估单一配置极其耗费资源和/或时间时,效率便至关重要。
自适应实验通过主动提出用于顺序评估的新配置,并利用从先前评估中获得的见解,为解决这个问题提供了方案。
今年,我们发布了Ax的1.0版本,这是一个开源的自适应实验平台,它利用机器学习来引导和自动化实验过程。Ax采用贝叶斯优化,使研究人员和开发人员能够进行高效实验,找到最优配置以优化他们的系统和流程。
配合这次重要发布,我们发表了一篇题为《Ax:一个自适应实验平台》的论文,该论文探讨了Ax的核心架构,更深入地解释了驱动优化的方法论,并将Ax的性能与其他黑盒优化库进行了比较。
Ax已在Meta内部成功应用于各个学科领域,包括:
- 传统的机器学习任务,如超参数优化和架构搜索。
- 应对生成式AI中的关键挑战,包括为训练AI模型寻找最优的数据混合比例。
- 在生产环境中调整基础设施参数或编译器标志。
- 在物理工程任务中优化设计参数,例如设计AR/VR设备。
通过使用Ax,开发人员可以运用最先进的方法论来进行复杂实验,最终更深入地理解并优化其底层系统。
如何开始使用Ax
要开始使用Ax来高效调整复杂系统中的参数,请通过 pip install ax-platform 安装库的最新版本,并访问Ax网站获取快速入门指南、教程以及关于Ax底层所用方法的深度解析。
Ax用于现实世界的实验
自适应实验非常有用,但运行起来可能具有挑战性。这些实验不仅需要使用复杂的机器学习方法来驱动优化,还需要专门的基础设施来管理实验状态、自动化编排、提供有用的分析和诊断等。此外,任何给定实验的目标通常比单纯改进单一指标更为复杂。在实践中,实验通常是在多个目标指标与多个约束条件和防护措施之间进行的谨慎平衡。
我们构建Ax,是为了让用户能够使用最先进的技术轻松配置和运行这些动态实验,并为研究人员提供一个稳健、成熟的平台,以便将尖端方法直接集成到生产系统中。
Ax用于理解系统
除了高效地寻找最优配置,Ax还是一个理解被优化底层系统的强大工具。Ax提供了一套分析工具(图表、表格等),帮助用户理解优化过程随时间如何推进,通过帕累托前沿展示不同指标之间的权衡,可视化一个或两个参数在整个输入空间上的影响,并解释每个输入参数对结果的贡献程度(通过敏感性分析)。
这些工具让实验者不仅能获得一个可以部署到生产环境的最优配置,还能更深入地理解他们的系统,从而为未来的决策提供信息。
Ax如何工作
默认情况下,Ax使用贝叶斯优化,这是一种有效的自适应实验方法,擅长在探索(了解新配置的性能)和利用(改进先前观察到的良好配置)之间取得平衡。Ax依靠BoTorch来实现其贝叶斯优化组件。
贝叶斯优化是解决全局优化问题的一种迭代方法,它不假设函数f的任何形式信息。在实践中,这意味着通过评估一些候选配置(即尝试一些配置并测量其效果),利用这些数据构建一个代理模型,使用该代理模型识别下一个最有希望评估的配置,并重复此过程,直到找到最优解或实验预算耗尽。
在典型设置下,Ax在贝叶斯优化循环中使用高斯过程作为代理模型。GP是一种灵活的模型,可以在量化不确定性的同时进行预测,并且在数据点非常少时尤其有效。然后,Ax使用一个称为期望改进的采集函数家族中的函数,通过计算任何新配置相对于先前评估的最佳配置的期望改进值,来建议下一个要评估的候选配置。
以下动画展示了这个循环过程,其中GP建模的目标指标在上方以蓝色绘制,EI在下方以黑色绘制;EI的最高值指示了下一个要评估的x值。一旦新的x值被评估,GP会使用新的数据点重新拟合,然后我们计算下一个EI值。
这个一维示例可以扩展到许多输入和输出维度,使Ax能够优化具有许多(可能数百个)可调参数和结果的问题。事实上,更高维度的设置,即覆盖搜索空间的成本呈指数级增长,是这种基于代理的方法相较于其他方法真正大放异彩的地方。
您可以在Ax网站的“贝叶斯优化介绍”页面上阅读更多关于贝叶斯优化的内容。
我们在Meta如何使用Ax
Ax已在Meta大规模部署,用于解决公司最具挑战性的优化问题。Meta有数千名开发人员使用Ax来完成诸如AI模型的超参数优化和架构搜索、在线推荐和排名模型的参数调优、基础设施优化以及AR/VR硬件设计的仿真优化等任务。
这些实验优化微妙的目标并利用复杂的算法。例如,我们曾使用多目标优化来同时提高机器学习模型的准确性并最小化其资源使用。当研究人员需要缩小自然语言模型以适应第一代Ray-Ban Stories时,他们使用Ax来寻找在大小和性能之间达到最优权衡的模型。此外,Meta工程师使用约束优化技术来调整推荐系统,在避免其他指标退化的同时优化关键指标。
最近,Ax被用于设计新型快速固化、低碳的混凝土混合物,并已应用于我们其中一个数据中心建筑工地。这些新混合物在推进我们到2030年实现净零排放的目标中发挥着重要作用。
我们在各个领域都看到这样的问题:系统的最终质量受到参数的影响,而这些参数之间的相互作用在没有实验的情况下难以分析,且实验本身成本高昂。Ax通过采用数据驱动的方法在实验进行过程中进行调整来解决这些挑战,使我们能够高效且有效地解决这些问题。
Ax的未来
阅读论文
《Ax:一个自适应实验平台》
要了解更多关于Meta开源的信息,请访问我们的网站,订阅我们的YouTube频道,或在Facebook、Threads、X、Bluesky和LinkedIn上关注我们。
致谢
Ax由Meta的自适应实验团队创建:Sebastian Ament, Eytan Bakshy, Max Balandat, Bernie Beckerman, Sait Cakmak, Cesar Cardoso, Ethan Che, Sam Daulton, David Eriksson, Mia Garrard, Matthew Grange, Carl Hvarfner, Paschal Igusti, Lena Kashtelyan, Cristian Lara, Ben Letham, Andy Lin, Jerry Lin, Jihao Andreas Lin, Samuel Müller, Miles Olson, Eric Onofrey, Shruti Patel, Elizabeth Santorella, Sunny Shen, Louis Tiao, 和 Kaiwen Wu。