Meta 的优化平台 Ax 1.0 能够简化大语言模型和系统优化

如今已进入稳定状态的 Ax 是 Meta 推出的一个开源平台，它的目标是帮助研究人员和工程师将机器学习应用于复杂且资源密集型的实验场景。过去几年中，Meta 已广泛使用 Ax 来改进 AI 模型、加速机器学习研究、调优生产基础设施等。

Ax 特别针对那些需要理解和优化 AI 模型或其他具有复杂配置系统的研究人员。Meta 的研究人员指出，在这类场景中，可能的配置组合数量极其庞大，几乎不可能通过线性方式高效地逐一评估。解决方案是采用自适应实验（adaptive experimentation），即系统会按顺序自动评估配置，并利用先前评估结果中的洞察来指导对解空间的探索。

自适应实验非常强大，但实施起来很有挑战性。这类实验不仅需要借助先进的机器学习方法驱动优化过程，还依赖专门的基础设施来管理实验状态、自动化编排流程、提供有效的分析与诊断等功能。

Meta 在内部使用 Ax 解决的问题包括：机器学习中的超参数优化与神经架构搜索、为训练 AI 模型发现最优的数据混合比例、基础设施调优、编译器标记（compiler flags）优化等。

其中一项尤为引人注目的应用是大语言模型（LLM）。Meta 研究人员已提供了全面的介绍，展示了如何利用 Ax 来编写更有效的提示词（prompts）、选择最高效的样例供 AI 遵循等场景。

优化过程中还存在另外一项挑战，那就是研究人员通常希望在满足约束条件和安全边界的前提下，同时优化多个目标指标。Meta 研究人员举了一个例子，他们曾使用 Ax 进行“多目标优化”，即在提升机器学习模型准确率的同时，尽可能降低其资源消耗。

除了优化本身，Meta 研究人员强调，Ax 还是一个深入理解被优化系统的高效工具：

Ax 提供了一整套分析功能（如图表、表格等），帮助用户了解优化进展、通过帕累托前沿（Pareto frontier）观察不同指标间的权衡、可视化一两个参数在整个输入空间中的影响，并通过敏感性分析（sensitivity analysis）解释每个输入参数对最终结果的贡献程度。

Ax 基于 PyTorch 和BoTorch，采用贝叶斯优化（Bayesian optimization）方法迭代测试候选配置。它使用一个代理模型（surrogate model）来识别下一个最值得评估的配置，并不断重复该过程，直到达成目标或耗尽计算预算。典型的代理模型是高斯过程（Gaussian Process），因为它能从极少量数据点中做出带有不确定性量化（uncertainty quantification）的预测而被广泛采用。

Meta 研究人员特别强调了 Ax 中具有表述性的 API，使其能够探索复杂的搜索空间，同时处理多个目标、约束条件和含噪声的观测结果。除此之外，Ax 支持并行评估多个不同的配置，并可在任意时刻中止实验。Ax 的一大优势在于提供了合理的默认设置，让非优化专家也能轻松使用高级的技术。

Ax 并不是唯一的开源黑盒自适应优化平台。其他可选方案包括SMAC、Nevergrad、Optuna、Dragonfly等。Meta 研究人员表示，Ax 提供了更广泛的功能，例如支持对参数和结果施加约束、处理带噪声的测量数据等。此外，Ax 以及许多同类工具均可与Ray Tune、Hydra等编排框架集成使用。

原文链接：

Meta's Optimization Platform Ax 1.0 Streamlines LLM and System Optimization