Ax平台中SurrogateSpec与ModelConfigs的正确使用方法

背景介绍

在Facebook开源的Ax平台中，随着版本迭代，一些API接口发生了变化。近期版本中，关于贝叶斯优化模型配置的方式有了重要更新，特别是从直接使用botorch_model_class参数转向使用更结构化的surrogate_spec.model_configs配置方式。

旧版配置方式的问题

在旧版Ax中，用户通常直接在GenerationStrategy的GenerationStep中通过Surrogate类来指定模型配置，例如：

GenerationStep(
    model=Models.BOTORCH_MODULAR,
    model_kwargs={
        "surrogate": Surrogate(
            botorch_model_class=SaasFullyBayesianSingleTaskGP,
            mll_options={
                "num_samples": 64,
                "warmup_steps": 128,
            },
        ),
    }
)

这种方式虽然直观，但随着Ax平台功能的扩展，逐渐暴露出扩展性不足的问题。因此，Ax团队决定引入更灵活的配置方案。

新版推荐配置方式

新版本推荐使用SurrogateSpec和ModelConfig来结构化地定义模型配置：

GenerationStep(
    model=Models.BOTORCH_MODULAR,
    model_kwargs={
        "surrogate_spec": SurrogateSpec(
            model_configs=[
                ModelConfig(
                    botorch_model_class=SaasFullyBayesianSingleTaskGP,
                    mll_options={
                        "num_samples": 64,
                        "warmup_steps": 128,
                    },
                )
            ]
        )
    }
)

新配置方式的优势

1. 更好的扩展性：可以轻松添加多个模型配置，支持更复杂的集成模型场景
2. 更清晰的代码结构：将模型配置封装在专门的配置类中，提高代码可读性
3. 未来兼容性：旧方式将被逐步淘汰，新方式是未来版本的标准做法
4. 统一配置接口：与其他Ax配置方式保持一致的风格

实际应用建议

对于正在迁移到新版本的用户，建议：

1. 检查现有代码中所有使用Surrogate类的地方
2. 逐步替换为SurrogateSpec和ModelConfig的组合
3. 利用这个机会重新审视模型配置，可能需要调整一些参数
4. 测试新配置是否产生与旧配置相同的优化效果

总结

Ax平台的这一API变更反映了项目向更结构化、更可扩展方向发展的趋势。虽然需要一定的迁移工作，但新的配置方式将为更复杂的优化场景提供更好的支持。对于新项目，建议直接采用新的配置方式以避免未来的兼容性问题。