Fully Bayesian Modular Botorch Model (SAASBO in MBM)

Dear developers,

I am deeply appreciating the great work you’ve done on BOTorch and Ax.

As far as I understand, Fully Bayesian BO is for now restricted to the qNEI acquisition and the SingleTaskGP surrogate. What would be the best approach to access the entire BOTorch modelling kit using Fully Bayesian estimation? Are you plannning anything in this direction?

I thought it would be great to have a Fully Bayesian option within botorch_modular, something along the lines of:

class ax.modelbridge.registry.Models(value)
...
FULLYBAYESIAN_MODULAR = 'FullyBayesianBotorchModel'

With FullyBayesianBotorchModel defined as:

from ax.models.torch.botorch_modular.model import BotorchModel
....
class FullyBayesianBotorchModel(FullyBayesianBotorchModelMixin, BotorchModel):

which would then allow fully flexible use of the BOTorch components. Of course I’d be happy to help out on this, Thank again for the great work!

Paolo

[lena-kashtelyan]

This functionality is now on latest! To set up the generation strategy:

from botorch.models import SaasFullyBayesianSingleTaskGP
from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
from ax.modelbridge.registry import Models

GenerationStrategy(
    steps=[
        # Quasi-random initialization step
        GenerationStep(
            model=Models.SOBOL,
            num_trials=5,  # How many trials should be produced from this generation step
        ),
        # Bayesian optimization step using the custom acquisition function
        GenerationStep(
            model=Models.BOTORCH_MODULAR,
            num_trials=-1,  # No limitation on how many trials should be produced from this step
            # For `BOTORCH_MODULAR`, we pass in kwargs to specify what surrogate or acquisition function to use.
            model_kwargs={
                "surrogate": Surrogate(SaasFullyBayesianSingleTaskGP),
            },
        ),
    ]
)

This will also be included in the next stable release of Ax.

[lena-kashtelyan]

This is now on latest stable: 0.2.5.1 : )