Derivative Enabled GPs with qKG

[michad1111]

Question about qKnowledgeGradient in botorch

Hello,

I am encountering an issue while using qKnowledgeGradient in botorch, as discussed in #1626 . I've built a mockup example based on the provided example by @Balandat (see #1626 (comment)), implementing a test function sin(10*x1) * cos(5*x2).

While it works with LogEI, I'm facing an issue with flattening training labels while optimizing the acquisition function when using qKG. To my understanding in case of MC acquisition functions I need to use an MCAcquisitionObjective which in my case should probably be an LinearMCObjective. I suspect the problem lies in the usage of ScalarizedPosteriorTransform or LinearMCObjective.

Here is my implementation:

weights = torch.zeros(1 + ndim)
weights[0] = 1.0
posterior_trans = ScalarizedPosteriorTransform(weights)
objective_mc = LinearMCObjective(weights)

acq_func = qKnowledgeGradient(
                    model,
                    num_fantasies=NUM_FANTASIES,
                    current_value=max_pmean,
                    maximize=True,
                    posterior_transform=posterior_trans,
                    objective=objective_mc,
                )

However, I get the following error:

File ~/miniconda3/envs/botorch-test/lib/python3.11/site-packages/gpytorch/models/exact_prediction_strategies.py:51, in DefaultPredictionStrategy.__init__(self, train_inputs, train_prior_dist, train_labels, likelihood, root, inv_root)
     47     train_labels = train_labels.reshape(
     48         *train_labels.shape[: -len(self.train_shape)], self._train_shape.numel()
     49     )
     50 except RuntimeError:
---> 51     raise RuntimeError(
     52         "Flattening the training labels failed. The most common cause of this error is "
     53         + "that the shapes of the prior mean and the training labels are mismatched. "
     54         + "The shape of the train targets is {0}, ".format(train_labels.shape)
     55         + "while the reported shape of the mean is {0}.".format(train_prior_dist.mean.shape)
     56     )
     58 self.train_inputs = train_inputs
     59 self.train_prior_dist = train_prior_dist

RuntimeError: Flattening the training labels failed. The most common cause of this error is that the shapes of the prior mean and the training labels are mismatched. The shape of the train targets is torch.Size([32, 32, 18]), while the reported shape of the mean is torch.Size([32, 32, 32, 6, 3]).

I believe I might have misunderstood the usage of LinearMCObjective.

Could someone kindly guide me on how to resolve this problem? Thank you in advance for your help!

Best regards.

My Notebook:
demo_deriv_enabled_BO_qKG.ipynb.txt

[saitcakmak]

Hi @michad1111. The issue seems to be the GPWithDerivatives model does not support fantasizing. If I try to fantasize using model.fantasize(train_x, SobolQMCNormalSampler(torch.Size([8]))), I get similar errors due to shapes of tensors not matching. Unfortunately, this is not something we can handle using the MCObjective or PosteriorTransform classes. My guess is that the condition_on_observation method of the model (in particular the get_fantasy_strategy method) would have to be rewritten to support fantasizing from the derivative GP.

[michad1111]

Thank you very much for the quick reply @saitcakmak . I will take a look into this and try to find a way to make this work.