Defining a Model class for ModelListGP

[Abrikosoff]

I have been trying to run a Service version of qMultiFidelityHypervolumeKnowledgeGradient; related thread I have already been able to make the input constructor work (I think!), but I keep getting stuck at the requirement for qMultiFidelityHypervolumeKnowledgeGradient to have a ModelListGP as surrogate: doing this in the GS, for example:

generation_strategy = GenerationStrategy(
                        steps=[
                            GenerationStep(
                                model=Models.SOBOL,
                                num_trials=1,  # https://github.com/facebook/Ax/issues/922
                                min_trials_observed=1,
                                max_parallelism=6,
                                model_kwargs={"seed": 9999},
                            ), 
                            GenerationStep(
                            model=Models.BOTORCH_MODULAR,
                            num_trials=-1,
                            model_kwargs={
                                "botorch_acqf_class": qMultiFidelityHypervolumeKnowledgeGradient,
                                # tried any one of these!
                                # "surrogate": Surrogate(SingleTaskGP),
                                # "surrogate": Surrogate(ModelListGP),
                                # "surrogate": ListSurrogate(SingleTaskGP)
                            },
                            model_gen_kwargs={
                                "model_gen_options": {
                                    "acqf_kwargs": {"cost_intercept": cost_intercept,
                                                    "num_fantasies": 2,
                                                    "cost_aware_utility": cost_aware_utility,
                                                    "target_fidelities": target_fidelities,
                                                    project: project,
                                                    },        
                                },
                            },
                        )
                        ]
                    )

gives me the error ValueError: qMultiFidelityHypervolumeKnowledgeGradient requires using a ModelList.

• "surrogate": Surrogate(SingleTaskGP) gives me the same error;
• "surrogate": Surrogate(ModelListGP) gives meTypeError: ModelListGP.__init__() got an unexpected keyword argument 'train_X', which is mysterious as nowhere do I have this arg in my code (but is probably because something else is trying to pass this to a surrogate object which is not correctly defined here by me).
• "surrogate": ListSurrogate(SingleTaskGP) informs me that ListSurrogate is deprecated (got this idea from this issue, which on first glance I would have thought would be the solution to this problem),

so the only thing left (I think) is via a custom model definition (like given here), ala doing something like class SimpleCustomGP(ModelListGP, GPyTorchModel):, but I am not sure how this can be properly defined (one model per MultivariateNormal?). Since I am not sure if this would be the correct way to go, I would like to ask this first before proceeding.

Thanks for help!

[bernardbeckerman]

Thanks for asking this! Let me follow up internally to see who might be best to answer.

[Balandat]

So the issue here is that by default Ax will select a (batched) non-model-list model that is incompatible with the decoupled acquisition function that qMultiFidelityHypervolumeKnowledgeGradient is a subclass of.

There is a way to avoid using batched models by passing in allow_batched_models=False as part of the surrogate specs of the model (a way of specifying model_kwargs that allow finer granular control). This would look something like passing the following (as part of model_kwargs:

"surrogate_specs": {"a": SurrogateSpec(allow_batched_models=False), "b": SurrogateSpec(allow_batched_models=False)},

However, I'm running into some weird issues with this, so I for now just hacked my way around it, see https://github.com/facebook/Ax/issues/2514#issuecomment-2198331809 - that again surfaced a more serious limitation, see that comment.

[saitcakmak]

A small correction on the SurrogateSpec suggestion: You do not want to define multiple surrogate specs. Multiple surrogates is intended only for some use cases that we never got around to implementing. To make sure the surrogate utilizes ModelListGP, you can use the following:

                            model_kwargs={
                                "botorch_acqf_class": qMultiFidelityHypervolumeKnowledgeGradient,
                                "surrogate_specs": {"model_list": SurrogateSpec(allow_batched_models=False)},
                            },

[Abrikosoff]

Let me test this out @saitcakmak , thanks a lot!

[lena-kashtelyan]

@Abrikosoff, did @saitcakmak's solution work out? Please close the issue if you are all set : )