[GENERAL SUPPORT]: Getting best predicted point of a botorch model

[RoeyYadgar]

Question

Hi,

I'm trying get the best predicted point (the point which minimizes the posterior mean instead of the best point observed so far) of a fitted model.
What I currently do is perform the optimization loop with the following GenerationStep:

GenerationStep(model=Models.GPEI,num_trials=-1)

And when i want to get the best predicted point I create a new model with PosteriorMean Acqusition class and generate a trial with it:

gp = Models.BOTORCH_MODULAR(experiment=exp,data=exp.fetch_data(),botorch_acqf_class=PosteriorMean,optimization_config = exp.optimization_config)
best_prediction = gp.gen(1)

This indeed works but it needs to re-fit the data again (even though I already fitted that same data in the optimization step), is there a better way to do this which doesn't require fitting the data again?

Thanks :)

Please provide any relevant code snippet if applicable.

No response

Code of Conduct

• I agree to follow this Ax's Code of Conduct

[bernardbeckerman]

Hi @RoeyYadgar, I'll follow up internally to see the best way to answer this. One quick follow-up is, when are you looking to obtain the best posterior mean point? E.g., is this after the experiment concludes and the last trial has been evaluated, or are you trying to obtain the best posterior mean point after each trial completes (and if so, why is that)? I'm asking because if it's the former, and no additional trials are being generated, you would only be training the model once (to produce the posterior mean minimizer) and there wouldn't be wasted compute.

[RoeyYadgar]

I'm indeed trying to do this after each trial. I do that to evaluate and analyze the performance of a simulated optimization problem.

I'm asking because if it's the former, and no additional trials are being generated, you would only be training the model once (to produce the posterior mean minimizer) and there wouldn't be wasted compute.

I'm not sure I understand what you mean by this, Why is it that in the former option I'd only be training it once?
I assumed that when I create a new model and pass it the sampled data it fits the GP to this data, and then minimizes the acqusition function. If I now create a new model with a different acquistion function (but the exact same data), I'd like to use the already-trained GP instead of having to train it again (I was thinking that I can achieve this by changing the acqusition function of an existing model, or creating a new model but passing it the underlying trained GP, but I couldn't figure out if the API allows me to do that)

[bernardbeckerman]

I do that to evaluate and analyze the performance of a simulated optimization problem.

I'd be interested to hear more about how your performance analysis incorporates the posterior mean minimizer.

Why is it that in the former option I'd only be training it once?

If you were only performing best-posterior-mean acquisition at the end of the optimization, you ostensibly wouldn't have performed a GPEI acquisition since you would use that to generate the next point to sample and there is no next point to sample (since optimization is complete).

cc @esantorella on how to use a new acquisition function on the already-trained surrogate model of a given modelbridge.

[esantorella]

Hi @RoeyYadgar , you might be able to get what you're looking for with generation_step._fitted_model.model_best_point(); see here. That gives the best point of those tried so far.

If you want the best point out of the whole search space, not just the points evaluated so far, replacing the acquisition function with a BoTorch PosteriorMean acquisition function is indeed an option. That would look something like this:

from ax.models.torch.botorch_modular.acquisition import Acquisition
from botorch.acquisition.analytic import PosteriorMean
from ax.models.torch_base import TorchOptConfig

generation_step = ...
model = generation_step._fitted_model.model

acq = Acquisition(
    botorch_acqf_class=PosteriorMean,
    surrogates=model.surrogates
    search_space_digest=model.search_space_digest,
    torch_opt_config=TorchOptConfig(...)
)
best_point, value, _ = acq.optimize(n=1, search_space_digest=model.search_space_digest)

If this doesn't work, could you provide a reproducible example?

I recommend working with the Service API if you can. If you're using that, then ax_client.get_best_parameters() should give what you're looking for.

[RoeyYadgar]

Hi,

I'd be interested to hear more about how your performance analysis incorporates the posterior mean minimizer.

I'm using this as a measure for "how well the optimization process is doing" as a function of the number of samples, and compare the predicted posterior mean minimizer (and the posterior variance) to the actual minimal value of the function. nothing too fancy :)

If you want the best point out of the whole search space, not just the points evaluated so far, replacing the acquisition function with a BoTorch PosteriorMean acquisition function is indeed an option.

This works well, thanks!
However if I understand correctly this gives me the best point without the ModelBridge "inverse transform", am I correct?
I took a look at the code of ax.modelbridge.torch.TorchModelBridge and ax.models.torch.botorch_modular.models.BoTorchModel and noticed they do have a best_point method. where ax.models.torch.botorch_modular.models.BoTorchModel.best_point calls ax.models.torch.botorch_modular.surrogate.Surrogate.best_in_sample_point and there is ax.models.torch.botorch_modular.surrogate.Surrogate.best_out_of_sample_point implemented which essentially performs what you described If I understand correctly?

I'm thinking of just implementing a ax.models.torch.botorch_modular.models.BoTorchModel subclass which calls
ax.models.torch.botorch_modular.surrogate.Surrogate.best_out_of_sample_point in best_point instead. Is this a good approach for doing this?

I also noticed that ax.models.torch.botorch_modular.surrogate.Surrogate.best_out_of_sample_point raises "Fixed features not yet supported." but I'm not sure I understand why this implementation doesn't support fixed features? it seems to work well if I remove this exception.
https://github.com/facebook/Ax/blob/b0156250749f534a56d0a3d0a685931f9716db13/ax/models/torch/botorch_modular/surrogate.py#L662

Thanks a lot :)

[esantorella]

I'm using this as a measure for "how well the optimization process is doing" as a function of the number of samples, and compare the predicted posterior mean minimizer (and the posterior variance) to the actual minimal value of the function. nothing too fancy :)

Makes sense! The true function value at the model's predicted best point is known as "inference regret" (as you may know) and is useful for benchmarking. I'd also find this valuable. Ax has generic best-point functionality through BestPointMixin, but that only uses "in-sample" prediction on arms that have been tried so far.

If you want the best point out of the whole search space, not just the points evaluated so far, replacing the acquisition function with a BoTorch PosteriorMean acquisition function is indeed an option.

This works well, thanks! However if I understand correctly this gives me the best point without the ModelBridge "inverse transform", am I correct?

Ah yes, good catch.

I'm thinking of just implementing a ax.models.torch.botorch_modular.models.BoTorchModel subclass which calls ax.models.torch.botorch_modular.surrogate.Surrogate.best_out_of_sample_point in best_point instead. Is this a good approach for doing this?

Hmm, I think you might run into the same problem that there is needed transform information on the ModelBridge that isn't attached to the BoTorchModel

I think what you want to do is change the best_point_recommender on the BoTorchModel to recommend_best_out_of_sample_point (the default is recommend_best_observed_point). If you do that, then the model_best_point method of the ModelBridge should give what you want.

I also noticed that ax.models.torch.botorch_modular.surrogate.Surrogate.best_out_of_sample_point raises "Fixed features not yet supported." but I'm not sure I understand why this implementation doesn't support fixed features? it seems to work well if I remove this exception. https://github.com/facebook/Ax/blob/b0156250749f534a56d0a3d0a685931f9716db13/ax/models/torch/botorch_modular/surrogate.py#L662

Thanks a lot :)

Hmm, this error message looks quite old, so it's possible it is not correct anymore. Could you send a reproducible example? If this is in fact working, I'd be more than happy to accept a PR removing the error and adding unit tests demonstrating that it works, or to add this to our backlog to fix.

[RoeyYadgar]

Makes sense! The true function value at the model's predicted best point is known as "inference regret" (as you may know) and is useful for benchmarking. I'd also find this valuable. Ax has generic best-point functionality through BestPointMixin, but that only uses "in-sample" prediction on arms that have been tried so far.

I see, wasn't aware of this. It can be useful as well, thanks!

Hmm, I think you might run into the same problem that there is needed transform information on the ModelBridge that isn't attached to the BoTorchModel

What I ended up doing is create a BoTorchModel subclass and "register" in in the Models enum, that way the ModelBridge wraps it like a regular model when I call Models.BOTORCH_MODULAR_SUBCLASS(...).

I think what you want to do is change the best_point_recommender on the BoTorchModel to recommend_best_out_of_sample_point (the default is recommend_best_observed_point). If you do that, then the model_best_point method of the ModelBridge should give what you want.

I see, this is really helpful, but I got a bit confused. If I understand, this is used in ax.models.torch.botorch.BotorchModel but not in ax.models.torch.botorch_modular.model.BoTorchModel. What's the difference between the two?
Is there a similar way to do it in BoTorchModel ? because in BoTorchModel.best_point it explicitly calls the surrogate.best_in_sample_point which is why I had to create a BoTorchModel subclass https://github.com/facebook/Ax/blob/main/ax/models/torch/botorch_modular/model.py#L474

Also when I try to use it with recommend_best_out_of_sample_point in BotorchModel I also get an exception.

Attaching a reproducible example of both attempts with (BoTorchModel, BotorchModel):

from ax import (
    Experiment,
    Objective,
    OptimizationConfig,
    Models
)
from ax import (
    Experiment,
    Objective,
    OptimizationConfig,
    ParameterType,
    RangeParameter,
    SearchSpace,
)
from ax.metrics.l2norm import L2NormMetric
from ax import Runner
from ax.models.torch.botorch_defaults import recommend_best_out_of_sample_point

#Initalize an experiment with the example from https://ax.dev/tutorials/gpei_hartmann_developer.html
class MyRunner(Runner):
    def run(self, trial):
        trial_metadata = {"name": str(trial.index)}
        return trial_metadata
    
search_space = SearchSpace(
    parameters=[
        RangeParameter(
            name=f"x{i}", parameter_type=ParameterType.FLOAT, lower=0.0, upper=1.0
        )
        for i in range(2)
    ]
)

param_names = [f"x{i}" for i in range(2)]
optimization_config = OptimizationConfig(
    objective=Objective(
        metric=L2NormMetric(name="l2norm", param_names=param_names),
        minimize=True,
    ))

exp = Experiment(
    name="exp_test",
    search_space=search_space,
    optimization_config=optimization_config,
    runner=MyRunner()
)
NUM_SOBOL_TRIALS = 5
sobol = Models.SOBOL(search_space=exp.search_space)
for i in range(NUM_SOBOL_TRIALS):
    generator_run = sobol.gen(n=1)
    trial = exp.new_trial(generator_run=generator_run)
    trial.run()
    trial.mark_completed()
#################################################################################################################

#Best point with Models.GPEI (which uses BotorchModel)
gpei_model = Models.GPEI(experiment=exp,data=exp.fetch_data())
print(f'Best sampled point GPEI : {gpei_model.model_best_point()}')
try:
    gpei_model = Models.GPEI(experiment=exp,data=exp.fetch_data(),best_point_recommender = recommend_best_out_of_sample_point)  
    print(f'Best out of sample point GPEI : {gpei_model.model_best_point()}')
except Exception as e:
    print(e)

#Best point with Models.BOTORCH_MODULAR (which uses BoTorchModel)
botorch_model = Models.BOTORCH_MODULAR(experiment=exp,data=exp.fetch_data())
print(f'Best sampled point BOTORCH MODULAR : {botorch_model.model_best_point()}')

#Create subclass of BoTorchModel to call surrogate.best_out_of_sample_point when using model_best_point
from ax.models.torch.botorch_modular.model import BoTorchModel
from ax.modelbridge.torch import TorchModelBridge
from ax.modelbridge import registry
from aenum import extend_enum
from ax.core import ObservationFeatures

class BotorchModelOOS(BoTorchModel):

    def best_point(self,search_space_digest,torch_opt_config,options=None):
        try:
            return self.surrogate.best_out_of_sample_point(
                search_space_digest=search_space_digest,
                torch_opt_config=torch_opt_config,
            )[0]
        except ValueError:
            return None

#register the model
registry.MODEL_KEY_TO_MODEL_SETUP["BOTORCH_MODULAR_OOS"] = registry.ModelSetup(
        bridge_class=TorchModelBridge,
        model_class=BotorchModelOOS,
        transforms=registry.Cont_X_trans + registry.Y_trans,
        standard_bridge_kwargs=registry.STANDARD_TORCH_BRIDGE_KWARGS,
    )

extend_enum(registry.Models,"BOTORCH_MODULAR_OOS","BOTORCH_MODULAR_OOS")


try:
    botorch_model = Models.BOTORCH_MODULAR_OOS(experiment=exp,data=exp.fetch_data())
    #still giving an error, needs to return candidates[0], acqf_values in https://github.com/facebook/Ax/blob/e459a083f334170ad155911af06cf665a010e549/ax/models/torch/botorch_modular/surrogate.py#L690 ?
    print(f'Best out of sample point BOTORCH_MODULAR : {botorch_model.model_best_point()}')
except Exception as e:
    print(e) 

try:
    #giving None because of exception raised here? https://github.com/facebook/Ax/blob/e459a083f334170ad155911af06cf665a010e549/ax/models/torch/botorch_modular/surrogate.py#L662
    print(f'Best out of sample point BOTORCH_MODULAR with fixed feature: {botorch_model.model_best_point(fixed_features=ObservationFeatures({"x0" : 0}))}')
except Exception as e:
    print(e)

[esantorella]

Thanks for the repro! I'm able to reproduce this and getting the error you reported:

File ~/ax/Ax/ax/models/torch/botorch_modular/surrogate.py:701, in Surrogate.best_out_of_sample_point(self, search_space_digest, torch_opt_config, options)
    686 acqf = Acquisition(  # TODO: For multi-fidelity, might need diff. class.
    687     surrogates={"self": self},
    688     botorch_acqf_class=acqf_class,
   (...)
    691     options=acqf_options,
    692 )
    693 candidates, acqf_values, _ = acqf.optimize(
    694     n=1,
    695     search_space_digest=search_space_digest,
   (...)
    699     fixed_features=torch_opt_config.fixed_features,
    700 )
--> 701 return candidates[0], acqf_values[0]

IndexError: invalid index of a 0-dim tensor. Use `tensor.item()` in Python or `tensor.item<T>()` in C++ to convert a 0-dim tensor to a number

This loos like a bug to me. Under the hood, the BoTorch function optimize_acqf is called; it can produce a scalar acqf_values when n=1, so there's no need to take the [0] element.

And sorry about the BotorchModel (legacy) vs. BoTorchModel (new) confusion. Your solution makes sense to me.

[esantorella]

That bug is fixed by #2652 .

[esantorella]

Closing since the question is answered and bug is fixed, but feel free to reopen if you have more questions, and we're noting that this functionality is important and ought to be supported better.