[Feature Request] Implement most likely heteroskedastic GP and write a tutorial notebook

[mc-robinson]

🚀 Feature Request

I would like to submit a potential example notebook, based on my following work here:
https://colab.research.google.com/drive/1dOUHQzl3aQ8hz6QUtwRrXlQBGqZadQgG
(hopefully it is easy for others to run in Colab!)

The notebook details the use of a heteroskedastic GP in BoTorch. However, before finalizing the notebook and cleaning it up a bit, there are a few things that I believe could be changed to help improve the entire training process (also described in notebook comments). I apologize if this reads as many features requests/issues in one -- but I am happy to contribute pull requests and please let me know what you all think / if I am totally wrong about something. I can also open separate feature requests for each statement if you agree on the need.

Listed in roughly decreasing order of perceived importance.

1. I would love to see the implementation of a mll.model.predictive_posterior function. Since the mll.model.posterior(X_test).mvn.confidence_region() function currently only gives the posterior for the function mean, this can be confusing to those who think they are getting the predictive posterior (as is shown in using confidence region with Gpytorch tutorials, for example). Currently, for simple models the predictive posterior is easy to get through mll.likelihood(mll.model(X_test)).confidence_region(). However, this does not work for the HeteroskedasticSingleTaskGP model. For this HeteroskedasticSingleTaskGP model, I could only obtain the prediction intervals using the following code:

mll.eval()
with torch.no_grad():
    posterior = mll.model.posterior(X_test)

    predictive_noise = torch.exp(mll.model.likelihood.noise_covar.noise_model.posterior(X_test).mean)
    # get standard deviation
    predictive_noise_std = torch.sqrt(predictive_noise).squeeze(-1)

   # fit posterior confidence
    lower, upper = posterior.mvn.confidence_region()
    # get posterior predictive confidence by adding noise from noise model
    lower_predictive = lower - 2*predictive_noise_std
    upper_predictive = upper + 2*predictive_noise_std

which is quite burdensome. (Note that the noise model of HeteroskedasticSingleTaskGP currently returns the log of the variance, thus the exponentiation. I believe this may be an error in the noise model.)

In summary, I believe a predictive_posterior function is often what people are after and would be a nice complement to posterior. Furthermore, it would be nice to standardize the methods of obtaining prediction intervals across tasks since the same method does not currently work for all models.

2. I do believe there should probably be a HeteroskedasticGP class that merely accepts train_x, train_y, and does not require the input of train_yvar. This can get rather complicated, but a simple solution for now would be to initially create a SingleTaskGP that is used to estimate train_Yvar once trained. The observed variances could then be fed into the current implementation of HeteroskedasticSingleTaskGP along with train_x, train_y, as is done in the notebook. This is perhaps the more controversial of the feature requests.

Please let me know if you have any other comments on the notebook, and thanks for the library!

[jacobrgardner]

The notebook looks great! a few comments:

1. In general, GPyTorch example notebooks should be independent of BoTorch's convenience methods for wrapping them and fitting them. In GPyTorch, there is no posterior method. The pattern to get predictions is to call model(test_x) to get p(f*|D, x*) and likelihood(model(test_x)) to get p(y*|D, x*). In your case, this just requires a small change because your likelihood also requires test_x as an input. The following code works for me in your notebook to get the predictive posterior after you have the trained GPyTorch model:

# test on the training points 
# call if X_test just for ease of use in future
X_test = torch.linspace(0,1,100)

mll.eval()
with torch.no_grad():
    posterior_f = mll.model(X_test)   # Call the GPyTorch model directly to get p(f*|D, x*)
    test_pred = mll.likelihood(mll.model(X_test), X_test)   # Call the likelihood on the result to get p(y*|D, x*)
    # above works -- we just needed to pass X_test to the likelihood as well because it depends on the input.

Then you should be able to call test_pred.confidence_region() or test_pred.variance as you'd like. This isn't as well documented as it should be because, as you've probably noticed, we are sorely lacking a heteroscedastic example notebook to document such things :-).

2. I think this should be possible to build as a likelihood, if slightly more complicated. Basically what you want is a GP that gets trained through its predictive mean while using something like softplus(predictive_mean) as the observation noise for the outer GP, right?

[jacobrgardner]

Whoops! It seems like I'm lost :-). I've been responding to all the issues opened this weekend that when I got another email about an example notebook for another GP model, I just automatically responded assumed we were on the GPyTorch repo.

Apologies! Obviously an example notebook on the botorch repo should definitely be free to use BoTorch features :-). My comments about getting the predictive posterior are still useful, hopefully.

[Balandat]

@mc-robinson, thanks for the comprehensive issue. We have designed the posterior API in a way to also support getting the posterior predictive, by using the observation_noise kwarg. So you can get the posterior predictive by calling gp.posterior(X_test, observation_noise=True). You will see that for the basic gpytorch model wrapper here, the implementation is exactly what @jacobrgardner suggested above.

However, since the HeteroskedasticGP subclasses SingleTaskGP, the actual code path for your case is through BatchedMultiOutputGPyTorchModel (which for performance reasons does some magic for representing multi-output models with independently modeled outputs as a batched model under the hood). Looking through the posterior code of that, I'm realizing that we appear to have missed handling the observation_noise kwarg in there. This should be easy to add though, I'll put up a PR for this tomorrow.

Note that the noise model of HeteroskedasticSingleTaskGP currently returns the log of the variance, thus the exponentiation. I believe this may be an error in the noise model.

Indeed, the actual noise returned by the likelihood should not be on the log scale. The reason the model fits the inner GP on the log noise is that this way we don't have to worry about zero (or negative) noise predictions. The correct way to handle is to pass in a custom Constraint object with torch.exp transform as the noise_constraint to the noise HeteroskedasticNoise constructor. (Alternatively, the model could also just be fit on the inverse softplus transform of the noise). I'll add this as another PR to do.

I do believe there should probably be a HeteroskedasticGP class that merely accepts train_x, train_y

Generally, a heteroskedastic GP for the case where you don't have noise observations would be a reasonable thing to implement. I haven't thought all too much about this so far though. Let me try to rephrase your suggestion: You suggest first training a model that infers the noise level, then, based on the mean prediction of that model, estimate a heteroskedastic noise function (using the data you've already used to fit the model)? This sounds like it could easily go wrong... maybe some kind of latent-variable model in which the noise levels for the inner GP are themselves model parameters that you optimize over would be more appropriate?

Edit: Just realized this is all in the colab notebook - awesome job! I'll give it a closer look tomorrow.

@jacobrgardner, haha, feel free to keep helping out with issues over here, I don't mind...

[mc-robinson]

Thanks for the detailed responses! And @jacobrgardner, I know it was a mistake haha, but I am also happy to (eventually) try to do a detailed notebook on heteroskedastic GPs in GpyTorch. It might be helpful to show how one might build one from the ground up. I already sort of have a draft ready -- are there any examples currently showing how to embed one model in another such as what is done in the heteroskedastic model in BoTorch?

@Balandat Sorry I missed the observation noise hyperparameter! I will definitely add that to the notebook when I clean it up. Also, I'll reiterate that my way of creating a heteroskedasticGP is basically the simplest/dumbest way I know, but there are lots of schemes that seem to get quite complicated. If anyone here has an idea for something better that isn't too nasty to implement, I would love to hear it.

The iterative scheme found in http://people.csail.mit.edu/kersting/papers/kersting07icml_mlHetGP.pdf is relatively close to what we are already doing. And might not be too hard to implement? I know that there are some issues noted in other papers with convergence and such, but might be worth trying. If a HeteroskedasticGP is built out that doesn't require a train_yvar input, the iteration should be relatively easy to write.

[Balandat]

The "Most Likely Heteroscedastic Gaussian Process Regression:" seems to be pretty much exactly what you're doing, just adding an EM-style fitting process where you iteratively re-estimate the mean and the observation variances from a heteroskedastic model, starting from a homoskedastic mean estimate. This should be just a few lines of code to implement. The only question is what exactly it means for this process to "converge" (the paper doesn't say). It's probably reasonable to just look at the change in mean predictions and noise variance estimates between iterations.

If a HeteroskedasticGP is built out that doesn't require a train_yvar input, the iteration should be relatively easy to write.

We could add a MostLikelyHeteroskedasticGP model that internally just uses the HeteroskedasticGP the same way you are.

[Balandat]

The issue with the observation_noise kwarg not being honored in for BatchedMultiOutputGPyTorchModel models is addressed in #182

[mc-robinson]

Hi guys, thanks for the work you've already put into this. I'll try to update the example notebook soon, and I'm happy to try and help build out the heteroskedastic models. Just let me know if there is a best place to start.

As for your question @Balandat

The only question is what exactly it means for this process to "converge" (the paper doesn't say). It's probably reasonable to just look at the change in mean predictions and noise variance estimates between iterations.

I too am a bit confused as to what they used. I assumed they were just looking at the value of the likelihood, but I think your method would also work well. I have also found the following critique in http://www.comonsens.org/documents/conferences/205_vhgpr_icml.pdf

the algorithm is not guaranteed to converge and may instead oscillate. Furthermore, it may require many iterations (each one requiring to train two standard GPs) before stabilizing. In a related, more recent work (Quadrianto et al., 2009), these issues are addressed by choosing g so as to maximize a penalized likelihood, equivalent up to a constant to p(g|D), thus introducing the MAPHGP approximation.

[Balandat]

I think your method would also work well

I think this would work for small models, but we will quickly run into scalability issues if we have a latent noise variable for every datapoint. You could think of many other ways to address that issue though, like parameterize the noise function for these latent variables using some kind of spline interpolation...

For now I think the EM-style algorithm would be a good thing to try first. Essentially, the model you end up getting is exactly a Heteroskedastic GP with the specific inferred noise values, so I don't think it's necessary to build a full new model.

Maybe for now it's enough to have a helper function of the form

construct_most_likely_heteroskedastic_GP(single_task_gp, options):
    [...]  # perform iterative fitting...
    return het_GP

that takes in a SingleTaskGP and internally performs the fitting, including the construction fo the Heteroskedastic GP, and then returns that fitted model. Happy to accept a PR :)

[mc-robinson]

Hi all,

Sorry for the delay on this one -- but I am finally attaching a slightly updated notebook:
https://colab.research.google.com/drive/1osVrSIfnrm7WWXgsYgP1rR6FDukHNjQF

In it is a first attempt at a MostLikelyHeteroskedasticGP class. As you can see in the notebook, I'm having trouble reaching "convergence" even though the solutions seem to be alright given enough iterations. The original paper, http://people.csail.mit.edu/kersting/papers/kersting07icml_mlHetGP.pdf , mentions that they were usually able to obtain convergence. However, this could be due to their method of estimating noise, which is different than mine (I'm not sure I totally get the reasoning behind their method). Please let me know if you have any suggestions. Hopefully, we can get a pull request started soon.

@Balandat , I also switched over to using the observation_noise parameter of posterior. However, for the heteroskedastic model, I am getting a slightly different result from when I did it by hand. Can you take a look at the differences and let me know if that is just my stupidity?

[Balandat]

Thanks @mc-robinson, I will take a look later today.

[Balandat]

However, for the heteroskedastic model, I am getting a slightly different result from when I did it by hand. Can you take a look at the differences and let me know if that is just my stupidity?

Looks looks like you're doing upper = mean + 2 * sqrt(var_f) + 2 * sqrt(var_obs) here. What you want to do instead is upper = mean + 2 * sqrt(var_f + var_obs), which is what the observation_noise=True kwarg will do for you.

Looking through the rest of the nb now.

[Balandat]

Thanks a lot for expanding on the nb.

However, this could be due to their method of estimating noise, which is different than mine (I'm not sure I totally get the reasoning behind their method). Please let me know if you have any suggestions.

So if I get that right your method of estimating the noise is just doing var(x_j) = (y_obs - post_mean(x_j))**2. I'm not quite sure I understand this. You're not including any variance prediction in this estimate (well, implicitly though the mean estimate, I guess). What the paper does instead is draw s samples from the (joint) posterior (over the function values, not the predictive posterior!) and compute an MC sample of the variance E[(y - mean(y))**2]. This seems like a reasonable approach. You can use MCSampler to do just that:

from botorch.sampling import IIDNormalSampler

sampler = IIDNormalSampler(num_samples=s, resample=True)
posterior = homo_model.posterior(X_train)
samples = sampler(posterior)
observed_var = 0.5 * (samples **2).mean(dim=0)

Another thing to note is that the paper users an RBF kernel for the noise GP, which will smooth things more than the Matern kernel that's used by default in the SingleTaskGP. This could also help with convergence.

A side comment: If you use fit_gpytorch_model instead of fit_gpytorch_scipy then you get a couple of benefits: (i) no need to call train or eval explicitly - inside the function will call train at the beginning and eval after the model fitting, and (ii) there is some logic for retrying the optimization upon failure with initial conditions sampled from the priors - this will improve robustness of the fitting.

Let's see how things go if you switch out the noise estimation procedure - if the results look reasonable and things converge then we can bake this into a PR.