evaluate qKG [Feature Request]

[shalijiang]

Current implementation of qKG only supports evaluating X_actual along with X_fantasies. If I want to plot the qKG surface for a one or two-dimensional function, I would have to optimize w.r.t. X_fantasies for each X_actual I suppose? It would be nice to implement another function for qKG that does that

[Balandat]

If I want to plot the qKG surface for a one or two-dimensional function, I would have to optimize w.r.t. X_fantasies for each X_actual I suppose?

Yes, that's correct. What you could do is make a simple wrapper around qKG along the lines of:

class DummyKG(MCAcquisitionFunction):

    def __init__(self, qKG, X_actual):
        self.qKG = qKG
        self.X_actual = X_actual

    def forward(self, X):
        return self.qKG(torch.cat([self.X_actual, X], dim=-2))


class EvalKG(MCAcquisitionFunction):

    def __init__(self, qKG):
        self.qKG = qKG

    def forward(self, X):
        dKG = DummyKG(self.qKG, X)
        _, val = optimize_acqf(dKG, bounds, q=self.qKG.num_fantasies)
        return val

This is kind of hacky though, let me think a little about what a good interface for this would be.

Also, you do NOT want to optimize EvalKG - that's the whole point of KG being a "one-shot" formulation.

[shalijiang]

Thanks Max!

[saitcakmak]

@Balandat this issue is pretty much what led me to discover the issue #345. My approach to this problem was to optimize a one-shot acquisition function (similar to qKG) by fixing the value I am trying to evaluate using fixed_features of optimize_acqf and optimizing the remaining features, i.e. the solutions to the fantasy models. In the case of qKG this would be to fix X_actual using fixed_features and calling optimize_acqf with fixed_features specified. @shalijiang you can try this as an alternative solution if you're still interested in evaluating KG.

[Balandat]

@saitcakmak the fixed_features in optimize_acqf are for fixing columns in the input, whereas for computing the KG value for a particular X you need to fix rows of the input. So this won't work that way.

Using qKG as in my example above is inefficient though, b/c we would keep re-conditioning the model on the same samples for each evaluation. What we can do is split out the logic after the fantasizing, make a simple function that computes the value based on the fantasized model, and use that both as an element in qKG, and allow to optimize this only over the fantasy parameterizations.

[hmedal]

I am also interested in plotting the qKG surface for a 2D function. For now I will try Max's example above but if you find something better please let me know.

[saitcakmak]

@Balandat, I have a proper implementation of this as part of a nested KG implementation. This is an evaluate_KG method that takes an n x q x d tensor X as input, generates the fantasies and (globally) optimizes the value function for all fantasies for all X.

It requires optimizing a [num_fantasies, n] batch of PosteriorMean, or value_function in general. This could be accommodated either by minor modifications to gen_batch_initial_conditions and optimize_acqf to allow for a batch_shape argument and passing this to draw_sobol_samples; or more efficiently by making a new gen_initial_conditions method that does smart initialization by first optimizing the current value_function, much like what is done in gen_one_shot_kg_initial_conditions. I can clean it up and create a PR if you think this is a useful addition.

[Balandat]

@saitcakmak This is great. The fact that @shalijiang raised this issues suggests that it indeed is a useful addition. Happy to review the PR for this.

[danielrjiang]

Seconded that this would be great! thanks @saitcakmak !