Changing prior mean after model is trained

[kentwar]

I'm using a FixedNoiseGP and after training on the whole data set I'd like to perform Expected Improvement on 2 subregions using the global GP. I'm interested in finding the optima for each region, but lets suppose that one region, on average, performs very highly and the the other performs very poorly. If I compare the EI for these regions it is not a fair comparison because the mean of the Global GP drives the overall performance for both regions (under representing the high performing region and vice versa which leads to problems when compating EI values across regions).

Instead, I'd like to train a GP on the whole dataset, to get a global set of hyperparameters and then simply replace the prior mean on the model to represent the individual regions.

I can achieve the desired outcome by training a FixedNoiseGP and then :

model.mean_module.constant = torch.nn.Parameter(torch.tensor(NEW_MEAN).double())
model.train()
model.eval()

Thinking about the predictive step it should be simpler than this, I should be able to replace the mean function with a constant value (as above) and then also update the vector of distances from the mean used in the prediction step, the variance is unaffected by the prior mean, so it shouldn't need to be changed.

After searching, I've found prediction_strategy.mean_cache which holds the values I need to change, but it doesn't seem to be working as I cant seem to set it.

I wonder if anyone might know the easiest way to do this or propose another way to do it that might, even fractionally, reduce the overheads of running the model.train(), model.eval() line. I'm running it thousands of times in my current context.

[saitcakmak]

You could technically delete it with

mem_cache = model.prediction_strategy.__dict__["_memoize_cache"]
mean_cache_key = list(mem_cache.keys())[0]
del mem_cache[mean_cache_key]

but it won't work since mean_cache will be re-generated from other cached values, such as model.prediction_strategy.train_prior_dist. You could delete that as well, but then you'd have to replace it, otherwise you'll get an error. You can give it a try and see if it speeds things up.

In general, I am a bit suspicious of the idea of training the covar_module along with the mean_module then swapping out the mean_module with something else. Wouldn't it be better to set a custom mean_module from the start (e.g., fix it to what you plan to set to after training) and train the covar_module along with it?

[kentwar]

Thanks for your response.

In general, I am a bit suspicious of the idea of training the covar_module along with the mean_module then swapping out the mean_module with something else. Wouldn't it be better to set a custom mean_module from the start (e.g., fix it to what you plan to set to after training) and train the covar_module along with it?

In the specific use case the reason is that I am training the GP to model a single fitness function over a single input space. The variance and lengthscale/ARD learned during training on the whole training set is valid for the whole space.

If I partition my input space into subregions and define my goal as attempting to find the local optima in each sub-region and attempt to use EI to find it three things happen. The first is that the best point used in the EI calculation is different for each region (which is trivial). The other is that my model is now trained with a global mean which is not locally-representative of any of the sub-regions (this causes a lot of problems when comparing acquisition values between regions as the prior mean influences the posterior prediction strongly), the third is that the subregion may contain no observations ( there may be many subregions or sparse initial sampling..) and so I need to use the whole observation set.

I'm a little out of my depth with the code if I'm honest but I thought I had a few options:

I can choose to train a model only with the local points (throwing away a lot of inference power.. bad idea)
I could do as you say and train the mean module for each region and train the covariance module with it (I don't want to do this as the covariance module is invariant if trained on all the points, which is what I want).
I could train the global model then swap out the prior mean.

Assuming it's only the mean module which needs to change I was basing my approach on the assumption that botorch was implementing something like the original rasmussen formulation which would mean a simple swap of the mu0 and (f-mu0) vector.

posterior predicted mu = mu0 + K(x)T K(xx)^-1 (f - mu0)

I have found a hacky way to change the mem_cache by creating a custom prediction strategy that takes a mu and mem_cache as an extra argument and calculates it on the fly, but this ends up being slower than simply changing the constant and then .train() .eval().

quick question.. does calling model.posterior(x) simply call model.call()?

[saitcakmak]

botorch was implementing something like the original rasmussen formulation which would mean a simple swap of the mu0 and (f-mu0) vector.

Due to the way things are implemented in GPyTorch, which utilizes lots of caching to speed things up, swapping parts of the vector is not as easy as the formula would suggest. You need to invalidate the caches and recompute them, which is what happens when you call train() / eval() on the model. This ends up computing all the caches all over again, which leads to the slowdown you observe.

does calling model.posterior(x) simply call model.call()?

Yes and no. It calls model.call() but also applies a bunch of transforms (if given to the model) and reshapes the tensors for multi-output models. The posterior implementation for SingleTaskGP is here.

In the specific use case the reason is that I am training the GP to model a single fitness function over a single input space. The variance and lengthscale/ARD learned during training on the whole training set is valid for the whole space.

So, this would still work if you implemented a custom mean_module and trained the global model with that. I am assuming you want to use some pre-determined constants, say constant_i for subregion_i. A (psuedo) code for the mean_module I suggest would like this:

from gpytorch.means import mean

class CustomMean(mean)
    # may want to implement __init__ to define your constant / subregions 
    def forward(self, input):
        if input in subregion_i:  # could do this in a loop to check for each sub region
            return constant_i  # need to make this into a tensor of appropriate shape

You could either use buffers for constants, in which case they would not be trained, or register them as parameters, in which case the constants would be trained along with the parameters of the covar_module. The benefit of this is it lets you train a global model while keeping local means, without having to swap constants after training.