Changing the prior mean on an already trained model

[kentwar]

I'm trying to perform Expected Improvement (EI) over a single GP but in discretized niches of the objective space. The prior mean over the different areas of the objective space will be different and this will induce errors in the EI calculation. I want to learn the hyper-parameters from the whole space but then replace the prior mean for each niche independently. Is there a simple way to do this without spending time retraining?

I was thinking I could train a single GP and then replace the mean (I'm using a botorch model here but it's obviously built on gpytorch)

model = SingleTaskGP(xtrain, ytrain)
mll = ExactMarginalLogLikelihood(model.likelihood, model)
fit_gpytorch_model(mll); 
model.mean_module = gpytorch.means.ConstantMean(prior=gpytorch.priors.NormalPrior(NicheMean,0.1)) 

but this doesn't work without retraining.

As I will need to do this many times, is there a quick way to swap the prior mean value?

I appreciate there are potential problems with this approach, but for what I want to do it is at least an interesting first step.

[kentwar]

I managed to accomplish what I was after by the following

niche_mean = -10
model.mean_module.constant = torch.nn.Parameter(torch.tensor([niche_mean]).float())
state_dict = full_model.state_dict()
niche_model = SingleTaskGP(xtrain, ytrain)

# Load parameters without standard shape checking.
niche_model.load_strict_shapes(False)
niche_model.load_state_dict(state_dict)

ie - initialising a new GP with the learned hypers but with the amended mean.

[gpleiss]

What exactly do you want the niche mean to be? Right now in your example you are modifying the prior on the mean function, rather than the mean itself.

[gpleiss]

Can you post a runnable code example?

[kentwar]

import numpy as np
import matplotlib.pyplot as plt
import gpytorch
import torch
from botorch.models import SingleTaskGP
from botorch.fit import fit_gpytorch_model
from gpytorch.mlls import ExactMarginalLogLikelihood


def CreateModel(xtrain, ytrain):
    '''
    Creates and trains a GpyTorch GP model.
    '''
    model = SingleTaskGP(xtrain, ytrain)
    mll = ExactMarginalLogLikelihood(model.likelihood, model)
    fit_gpytorch_model(mll);

    return(model)

fig = plt.figure(figsize=(15,10))

#### INITIAL SETUP
dim = 1 ; init_points = 3
ObjFun = lambda x : (x+0.2)**2
lb = -1
ub = 1

#True function samples
x = np.linspace(-1,1,100)
y = ObjFun(x)
plt.plot(x,y, label ='True')
## Random sampling for initial points
xtrain = torch.rand(init_points,dim)*(ub-lb)+lb
yvals = ObjFun(xtrain)
ytrain = torch.tensor(yvals).reshape(init_points,1)

## Create and Fit GP model
model = CreateModel(xtrain, ytrain)

## Plot GP Posterior
x = np.linspace(-1,1,100)
y = [model.posterior(torch.tensor([xi]).float()).mean for xi in x]
plt.plot(x,y , label = 'original GP')
plt.scatter(xtrain,ytrain)


# Make adjustement to Prior Mean value in posterior model
model.mean_module.constant = torch.nn.Parameter(torch.tensor([-0.2]).float())

## Plot new GP
x = np.linspace(-1,1,100)
y = [model.posterior(torch.tensor([xi]).float()).mean for xi in x]
plt.plot(x,y , label = 'GP with constant')
plt.scatter(xtrain,ytrain)

# Make adjustement to Posterior and rebuild from model state
state_dict = model.state_dict()
model = SingleTaskGP(xtrain, ytrain)
# Load parameters without standard shape checking.
model.load_strict_shapes(False)
model.load_state_dict(state_dict)

## Plot new GP
x = np.linspace(-1,1,100)
y = [model.posterior(torch.tensor([xi]).float()).mean for xi in x]
plt.plot(x,y , label = 'retrained GP with constant')
plt.scatter(xtrain,ytrain)
plt.legend()

You can see the green line (the one that doesn't pass through the points) is just the first GP that's been dragged down.

The one that performs as I'd expect is the retrained GP (red) that fits through the points and is dragged down towards the new prior mean value.

[kentwar]

Just as an update, the method i called 'retraining' works well and is not really retraining at all, so the wallclock cost is negligible.

[gpleiss]

The problem is that the model caches some terms of the GP posterior. When you update the mean, the caches aren't cleared, and so you compute the wrong thing.

This is an issue we are aware of (#1610). For now, you can force a cache reset by doing the following.

model.mean_module.constant = torch.nn.Parameter(torch.tensor([-0.2]).float())
# Clear cache
model.train()
model.eval()

[kentwar]

Great.

For anyone reading this thread, the speed up from this method is considerable (2 to 3 times faster than my previous method).