Custom means for different dimensions...

[ghutchis]

I'm intending to port some code to BoTorch using GPyTorch - and I'd like to supply priors for different dimensions. For now, let's assume a positive quadratic would suffice e.g., (x - c)**2, in which I can register a learnable parameter c for the likely position of the minima. (I'm a chemist, and one can usually guess some parameters like bond lengths or angles.)

That works great for one dimension - I defined my QuadraticMean similar to LinearMean.

Instead, I want to give a separate custom mean for each dimension (e.g., maybe I want to specify a CosineMean or a pre-trained neural network for some dimensions).

For kernels, there are examples using AdditiveStructureKernel(GridInterpolationKernel(..., num_dims=1))

Anything similar for stacking means?

Thanks in advance!

[Balandat]

What exactly do you mean by "a separate custom mean for each dimension"? The mean function (in the standard, non-multi-task setting) is a scalar function. Do you want to just combine those functions together in some pre-defined way, say additively?

I don't think there is anything pre-canned, but you should be able to just define a custom mean that picks out the right dimensions?

class ComboMean(Mean):

    def __init__(self, input_size, batch_shape=torch.Size(), pre_trained_network, lin_weights):
        super().__init__()
        self.pre_trained_network = pre_trained_network
        self.register_parameter("lin_weights", parameter=torch.nn.Parameter(torch.randn(*batch_shape, input_size, 1)))

    def forward(self, x):
        mean1 = self.pre_trained_network(x[..., :k])   # NN pass on the first k dimensions
        mean2 =  x[..., k:].matmul(self.lin_weights).squeeze(-1)  # linear function on the rest
        return mean1 + mean2

You can also get fancier by allowing to pass other mean functions into the arguments of the constructor of such a combo mean if that's the kind of flexibility you'd like.

[ghutchis]

Okay, let me restate. I want to hint to the GP an initial belief (e.g., there is a minimum of x1 near 1.08, a minimum of x2 at 1.54, a minimum of x3 at 104, etc.) since data points are expensive and I do not have many initial values.

In the simplest case, I'd be okay with a quadratic mean of the form weight*(x - center)**2 + bias if I can supply a list of the centers and/or weights, e.g.

self.mean_module = CustomQuadraticMean()
self.mean_module.initialize(center = [ 1.08, 1.54, 104 ])

But yes, I think I can reframe the ComboMean() example to handle my use case.

[Balandat]

I see. So you don't really have a prior on the function value, just that you expect an optimum at a particular point?

There is some recent work on this particular problem (not sure how well that works, the empirical results focus a lot on the relatively large sample regime and the differences are kind of hard to interpret [log difference of 10-18?]).

Another approach one could take here is to sample the initial points from a biased distribution where you put more mass around where you think the optimum is (e.g. some MVN centered at the location of the optimum, or a mixture of Gaussians if there are possibly multiple).

[ghutchis]

Sometimes we have a good prior on the function value and expected optimum. The trick is that different variables (e.g., bond lengths, angles, ..) have different length scales and priors (e.g., quadratic for some, periodic for others). For example a 1D cut might be:

So the ideal is to supply a tuned prior for each variable (thus the question). Failing that, a quadratic isn't a bad approximation near the minima. Well.. it's at least worth some tests to see if a quadratic is better than a flat prior since it can suggest the minimum.

I also considered sampling from a biased distribution. It seems like an easy way to suggest an optimum, particularly if acquisition is biased towards exploitation. If it's a good point, the acquisition function should give up on other regions quickly.

Benchmarking both strategies will be useful - just needed an example of coding that style of mean. Thanks!