Use ModelListGP for composite functions

[ToennisStef]

Hello,
I am trying to use the ModelListGP as the multi-output probability model for composite functions. I am trying to essentially replicate the Composite Bayesian Optimization with the High Order Gaussian Process tutorial by defining a new objective and replacing the HOGP with the ModelListGP.

The problem the i wanted to tackle is the following:
I want to model the optimize the composite function $h(x)$:

$$ \min h(x)=g(f(x)) \quad \text{ with } \quad f(x) = \left\{ \begin{aligned} f_1(x_{1})\\ f_2(x_{2},x_{3}) \end{aligned} \right. \ \quad \text{ and } \quad g(x)=f_1(x)+f_2(x)$$

$x=(x_{1}, x_{2},x_{3})^{T}$

with $f_{1}(x_{1})$ beeing the Forrester function:
$f_1(x_{1}) =(6x_{1}-2)^{2}\sin(12x_{1}-4)$

and $f_{2}(x_{2},x_{3})$ beeing the Six-hump camel function:
$f_2(x_{2}, x_{3}) =(4-2.1x_{2}^{2}+\frac{x_{2}^{4}}{3})x_{2}^{2}+x_{2}x_{3}+(-4+4x_{3}^{2})x_{3}^{2}$

I modeled $f_{1}(x_{1})$ and $f_{2}(x_{2},x_{3})$ as seperate SingelTaskGP models and combined them with the ModelListGP. The different GP models are not correlated and have different input dimensions. Reading from the ModelListGP description the different input dimensions should not be a problem since the model "supports different-shaped training inputs for each of its sub-models". The final transformation $g(x)$ is a GenericMCObjective. I am now having problems with the input dimensions of the combined model as well as the input dimensions of the model.posterior(SomeRandomData) and also with the input dimension of the acquisition function.

This is the code:
With a miscellaneous setup code:

torch.manual_seed(0)
device = (
    torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:4")
)
dtype = torch.double

print("Using ", device)

models_used = (
    "rnd",
    "ei",
    "ei_MLGP_cf",
)

Test function definitions as well as bounds and objective definition code:

def torchForresterFunction(X1):
    r"""
    Implementation of Forrester optimization test function returning a torch tensor
    """
    Y = torch.pow(6*X1-2,2)*torch.sin(12*X1-4)
    return Y

def torchSixHumpCamelFunction(X2,X3):
    r"""
    Implementation of the Six-Hump camel test function returning a troch tensor
    """
    term1 = (4-2.1*(X2**(2))+((X2**(4))/(3)))*X2**(2)
    term2 = X2*X3
    term3 = (-4+4*(X3**(2)))*(X3**(2))
    return term1+term2+term3

def FFunction(X):
    
    term1 = torch.stack([torchForresterFunction(x) for x in X[:,0]])
    term2 = torch.stack([torchSixHumpCamelFunction(*x) for x in X[:,1:3]]) 
    term3 = torch.stack((term1,term2)).T
    return term3

def Obj_funcAddition(samples):
    sum = torch.sum(samples,1)
    return sum

def gen_rand_points(bounds, num_samples):
    points_nlzd = torch.rand(num_samples, bounds.shape[-1]).to(bounds)
    return bounds[0] + (bounds[1] - bounds[0]) * points_nlzd

def optimize_ei(qEI, bounds, **options):
    cands_nlzd, _ = optimize_acqf(qEI, bounds, **options)
    return cands_nlzd

def prepare_data(device=device, dtype=dtype):
    bounds= torch.tensor(
        [[0, -3, -2], [1, 3, 2]],
        device=device,
        dtype=dtype
    )
    
    objective = GenericMCObjective(Obj_funcAddition)

    return objective, bounds

The BO Loop code:

#Number of BO iterations
n_batches = 1

with gpt_settings.cholesky_jitter(1e-4):
    objective, bounds = prepare_data(device=device,dtype=dtype)
    
    train_X_init = gen_rand_points(bounds,n_init)
    train_Y_init = FFunction(train_X_init)
    
    train_x = {k: train_X_init.clone() for k in models_used}
    train_y = {k: train_Y_init.clone() for k in models_used}

    #run BO loop
    for i in range(n_batches):
        tic = time.monotonic()
    
        #get best observation,log status
        best_f = {k: objective(v).max().detach() for k, v in train_y.items()}
    
        print(
                f"Iteration {i+1:>2}/{n_batches}, best obs.: ",
                ", ".join([f"{k}: {v:.3f}" for k, v in best_f.items()])
            )
        # generate random candidates
        cands = {}
        cands["rnd"] = gen_rand_points(bounds, 1)
        
        optimize_acqf_kwargs = {
            "q": 1,
            "num_restarts": 10,
            "raw_samples": 512,
        }
        sampler = IIDNormalSampler(sample_shape=torch.Size([128]))
        
        train_Y_ei = objective(train_y["ei"]).unsqueeze(-1)
        model_ei = SingleTaskGP(
            train_x["ei"],
            train_Y_ei,
            input_transform=Normalize(train_x["ei"].shape[-1]),
            outcome_transform=Standardize(train_Y_ei.shape[-1]),
        )
        mll = ExactMarginalLogLikelihood(model_ei.likelihood,model_ei)
        fit_gpytorch_mll_torch(mll,step_limit=1000, optimizer=partial(Adam,lr=0.01))

        #generate qEI candidate (single output modeling)
        qEI = qExpectedImprovement(model_ei,best_f=best_f["ei"],sampler=sampler)
        cands["ei"] = optimize_ei(qEI, bounds, **optimize_acqf_kwargs)
        
        #Setting up the ModelListGP
        model1 = SingleTaskGP(
                train_X=train_x["ei_MLGP_cf"][:,0:1],
                train_Y=train_y["ei_MLGP_cf"][:,0:1],
                input_transform=Normalize(train_x["ei_MLGP_cf"][:,0:1].shape[-1]),
                outcome_transform=Standardize(train_y["ei_MLGP_cf"][:,0:1].shape[-1]),
                )
        model2 = SingleTaskGP(
                train_X=train_x["ei_hogp_cf"][:,1:3],
                train_Y=train_y["ei_hogp_cf"][:,1:2],
                input_transform=Normalize(train_x["ei_MLGP_cf"][:,1:3].shape[-1]),
                outcome_transform=Standardize(train_y["ei_MLGP_cf"][:,1:2].shape[-1]),
                )
        model = ModelListGP(model1,model2)
    
        mll_model = SumMarginalLogLikelihood(model.likelihood,model)
        with _fast_solves(True):
            fit_gpytorch_mll_torch(mll_model,step_limit=1000,optimizer=partial(Adam,lr=0.01))

        #generate qEI candidate (multi-output modeling)
        
        qEI_MLGP_cf = qExpectedImprovement(
            model,
            best_f=best_f["ei_MLGP_cf"],
            sampler=sampler,
            objective=objective,
        )
        cands["ei_MLGP_cf"] = optimize_ei(qEI_MLGP_cf,bounds,**optimize_acqf_kwargs)

        # make observations and update data
        for k, Xold in train_x.items():
            Xnew = cands[k]
            if Xnew.shape[0] > 0:
                train_x[k] = torch.cat([Xold,Xnew])
                train_y[k] = torch.cat([train_y[k], FFunction(Xnew)])
                
        print(f"Wall time: {time.monotonic() - tic:1f}")
    objective_dict = {k: objective(train_y[k]) for k in train_y}

the code is currently failing at this step:

cands["ei_hogp_cf"] = optimize_ei(qEI_hogp_cf,bounds,**optimize_acqf_kwargs)

Since i didn't get it to work i was playing around with the model itself:
I defined a tensor x_input which should serve as input for the acquisition function qEI_MLGP_cf and the model.posterior.

x_input= torch.tensor([0.5, 0.0, 0.0]).reshape([3, 1, 1])
posterior = model.posterior(x_input)
qEI_hogp_cf(x_input)

The model.posterior at least accepted the dimensions of x_input but the output was not what i expected.
qEI_hogp_cf(x_input) threw another error.
I wanted the entries of the tensor to function as the $x_{1}, x_2 , x_3$ values and expected that one entry of the tensor would be used as input for model1 in the sense of the function description above $f_1(x_1)$ and the two other entries would be used as the input for the model2 similarly as $f_2(x_2,x_3)$ . But instead model.posterior(x_input), produced a tensor with the shape: torch.Size([3, 1, 2]). It seemd to have broadcasted each entry of x_input to the input dimension of model1 and model2. To clarify: the resulting posterior.mean was

tensor([[[-4.7261, 30.1441]],

        [[ 2.8009,  6.3355]],

        [[ 2.8009,  6.3355]]], dtype=torch.float64,
       grad_fn=<TransposeBackward0>)

I was looking at the batching documentation but did not get smarter with regards to evaluating the posterior of a ModelListGP model and acquisiton functions using the ModelListGP. I also couldn't find a tutorial or discussion covering this issue.

What is the correct way to evaluate the posterior here? That each entry gets used as the correct model and that i end up with only one two dimensional output? And following this up how can i then set up the ModelListGP for composite function modeling?

I wanted to use the ModelListGP to not be restricted to the individual model requiring the same ammount of training data.

I am thankful for every help!

[saitcakmak]

Hi @ToennisStef. ModelListGP is generally not designed to work with sub-models with different input features out of the box. When you evaluate a ModelListGP with test_X, each sub-model gets evaluated using the same test_X and the outputs are concatenated. In your case, test_X is batch x q x 3, so each sub-model has to support inputs of this shape.

One way to do this is to use the FilterFeatures input transform. This will filter only the desired columns from the original input for each model, and evaluate the model using the resulting input of lower dimension. A setup like this should work:

        model1 = SingleTaskGP(
                train_X=train_x["ei_MLGP_cf"][:,0:1],  # this is n_1 x 1
                train_Y=train_y["ei_MLGP_cf"][:,0:1],
                input_transform=ChainedInputTransform( 
                    # Disabling transform_on_train since the train_X is already filtered to x_0 only
                    filter=FilterFeatures(feature_indices=torch.tensor([0], dtype=torch.long), transform_on_train=False),
                    normalize=Normalize(train_x["ei_MLGP_cf"][:,0:1].shape[-1]),
                ),
                outcome_transform=Standardize(train_y["ei_MLGP_cf"][:,0:1].shape[-1]),
                )
        model2 = SingleTaskGP(
                train_X=train_x["ei_hogp_cf"][:,1:3],
                train_Y=train_y["ei_hogp_cf"][:,1:2],
                input_transform=ChainedInputTransform( 
                    filter=FilterFeatures(feature_indices=torch.tensor([1, 2], dtype=torch.long), transform_on_train=False),
                    normalize=Normalize(train_x["ei_MLGP_cf"][:,1:3].shape[-1]),
                ),
                outcome_transform=Standardize(train_y["ei_MLGP_cf"][:,1:2].shape[-1]),
                )
        model = ModelListGP(model1,model2)

With this setup, evaluating model.posterior(test_X) with test_X of shape batch x q x 3 will be roughly equivalent to evaluating model1.posterior(test_X[..., 0:1]) and model2.posterior(test_X[..., 1:3]) and concatenating the outputs.

Also to clarify, since your ModelListGP has two sub-models of each with one outcome dimension, it will produce 2 outcomes. So, model.posterior(test_X).rsample(sample_shape) will produce an output of shape sample_shape x batch x q x 2, where output[..., 0] will be the output of model1 and output[..., 1] will be the output of model2.

[ToennisStef]

Thanks a lot! That did the trick! I needed to change to objective a little bit but now its working! I have still to validate the results but the optimization loop returns at least plausable results. Also thanks for clarifying the output dimensions of the model. That is very helpful

[Balandat]

"supports different-shaped training inputs for each of its sub-models"

To clarify this part: This is meant to say that the models can support different numbers of training data points (but not different feature spaces). I guess this is something we can improve the clarity of.

[ToennisStef]

Ah ok thanks for clarifying!