[RFC] Support for non-MC acquisition objectives for MC acquisition functions

[saitcakmak]

What

I want to add support for non-MC AcquisitionObjectives to MCAcquisitionFunctions, i.e., add support for a posterior_objective that transforms the posterior before sampling. The existing objectives for MCAcquisitionFunctions operate on the samples of the posterior rather than the posterior itself.

Why

The posterior sampling from q-points jointly has an O(q^3) cost (using Cholesky). Using risk measures, i.e., InputPerturbation or AppendFeatures, multiplies the size of the q-batch with n_w, leading to a joint posterior over q * n_w points and a sampling cost of O(q^3 n_w^3). If using Expectation as the risk measure, we have an option of transforming the q * n_w posterior into a posterior over q expectations before sampling (I think at a cost of O(q^2 n_w^2)), reducing the cost of sampling to O(q^3) (This transform is similar to ScalarizedObjective but operates over the q-batch dimension). This is currently not supported by MCAcquisitionFunctions.

Proposal

I can think of two ways of implementing this. In both versions, we need to introduce new attribute posterior_objective (possibly with a better name).

i) add posterior = self.posterior_objective(posterior) if self.posterior_objective is not None else posterior line after the posterior call in MCAcquisitionFunctions.

ii) Move the _get_posterior method of AnalyticAcquisitionFunction to AcquisitionFunction and use the posterior_objective in here, essentially unifying the objectives operating on the posterior across both analytic and MC acquisition functions. This would also require replacing the model.posterior(...) calls in MCAcquisitionFunctions with a _get_posterior call.

cc @sdaulton, @Balandat

[sdaulton]

This makes sense to me. Curious to hear @Balandat's thoughts. No strong preference between 1) and 2). Perhaps 2) is a bit cleaner since the posterior_objective behavior/code is unified across analytic and MC acquisition functions.

[Balandat]

Are there other concrete transformations / objectives you would want to compute on a posterior that way? If we can't think of any then I'm wondering whether we should jump to implement this as generic functionality (over e.g. defining a simple subclassed model that overwrites the posterior call and applies this transformation there).

If we do want to go the generic route, I also don't have particularly strong feelings about which way to go. Though it seems that if this is the default method to use in all acquisition functions, maybe it should be more prominent than a "private" method with underscore? In that case I wonder if this is something that should happen directly in the posterior call itself (e.g. by allowing that to take an objective)?

[saitcakmak]

Are there other concrete transformations / objectives you would want to compute on a posterior that way?

Not really. In theory, you could apply any other transform that way as well, just returning a TransformedPosterior object. Since we have the MCAcquisitionObjective class that supports anything that would be transforming the samples, this doesn't seem like a real use case.

In that case I wonder if this is something that should happen directly in the posterior call itself (e.g. by allowing that to take an objective)?

I like this design. We could simply replace this https://github.com/pytorch/botorch/blob/1da6044afb9219b9c95d85d6f0fc2d91d2729777/botorch/models/gpytorch.py#L360
with return objective(posterior).

[Balandat]

Yeah that would work (there are a few other models where we'd need to do this. Then if we make transform an optional arg to all posterior calls we could just check for that being not None and apply it as you suggested.

I wonder if this design could also easily be used to implement ScalarizedObjective. Seems like this could potentially make it a lot easier to apply a both a posterior transform and an objective.

One possible way to do this would be for all acquisition functions to take in a posterior_transform arg (that would replace the objective arg for the analytic ones, and all MCAcquisitionFunction constructors to take in an MCAcquisitionObjective objective arg. Inside the posterior calls it would then be applied after a potential untransform_posterior from the outcome transform was applied. In fact, in an ideal world one could just make a PosteriorTransform class that implements the basic interface, and then subclass OutcomeTransform from that since outcome transforms add strictly more capabilities (transforming/untransforming the observations).

I guess I'm getting carried away here ... maybe we can start with a simple proof of concept that doesn't make substantial changes all over the place...

[saitcakmak]

I wonder if this design could also easily be used to implement ScalarizedObjective.

I don't see why not. This would simply be absorbing in the existing functionality in _get_posterior (which is to return objective(posterior) if objective exists).

I can put up a limited implementation with a posterior_transform arg that is passed into the posterior. We can go from there.