Questions regarding HeteroskedasticNoiseSingleTaskGP

[ErikOrm]

Hi,

I have some short questions regarding the Heteroskedastic Noise Model. If someone had any insight that would have been of great help.

The first is regarding standardization. As I suppose that the y and Yvar input must maintain their original relationship, and y should be standardized, that means that Yvar cannot be standardized. This is something that botorch complains about through an InputWarning. And I suppose that it should also make it more unstable? Would it not make sense to also include an standardization outcome transform (in addition to the log transform) in the STGP that is employed within the HeteroskedasticSingleTaskGP?

The second question is regarding fitting the HeteroskedasticSingleTaskGP using fit_gpytorch_mll(ExactMarginalLogLikelihood(model.likelihood, model)). It works as long as it does not run into numerical issues, but when the multitrial fit routine kicks in it crashes with the following error:

Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.

I can provide the stack trace if that would help. Lastly, I wanted to implement a simple "manual" optimization routine that I could call if the Scipy opimization crashed as follows:

self.train()
optimizer = torch.optim.Adam([{"params": self.parameters()}], lr=0.1)
for _ in range(100):
    optimizer.zero_grad()
    out = self(*self.train_inputs)
    loss = -mll(out, self.train_targets)
    loss.backward()
    optimizer.step()
self.eval()

which crashes with the following error: TypeError: forward() missing 1 required positional argument: 'x'. Maybe because the noise model is not set to training somehow? Would it be possible to fit the HeteroscedasticSingleTaskGP manually this way and what would I need to add in that case?

Thanks,
Erik

[Balandat]

The first is regarding standardization. As I suppose that the y and Yvar input must maintain their original relationship, and y should be standardized, that means that Yvar cannot be standardized. This is something that botorch complains about through an InputWarning. And I suppose that it should also make it more unstable? Would it not make sense to also include an standardization outcome transform (in addition to the log transform) in the STGP that is employed within the HeteroskedasticSingleTaskGP?

Yes, that would be the right approach here. Unfortunately, this behavior is broken right now due to some in-the-weeds details on how we handle the transforms. #861 gives some details on the issue, including a workaround: #861 (comment). But the truth is that so far we haven't been able to prioritize fixing this for good.

The second question is regarding fitting the HeteroskedasticSingleTaskGP using fit_gpytorch_mll(ExactMarginalLogLikelihood(model.likelihood, model)). It works as long as it does not run into numerical issues, but when the multitrial fit routine kicks in it crashes with the following error:

Hmm interesting. I don't know exactly what's going on but I can see how some of the recovery logic in the fit_gpytorch_mll method may not work fully with the rather non-standard GP-wrapped-in-a-GP setting. If you could provide a stack trace (or even better, a reproducible toy example) then we can take a look!

Lastly, I wanted to implement a simple "manual" optimization routine that I could call if the Scipy opimization crashed as follows:

Could you provide the full stack trace / repro for this as well? It should definitely be possible to train the model manually, it looks like some inputs are just not getting passed through (may require passing in some of the inputs in your training setup but I'd have to see a repro to see what exactly is going on).

[ErikOrm]

Thank you greatly Max for your quick reply. Here is a minimal example to showcase where it fails for me. The first problem appears if you would set USE_MANUAL_FIT=True, and the second if you set it =False The first only appears if it runs into numerical issues, which it does for me with this seed.

import torch
from botorch.models import HeteroskedasticSingleTaskGP
from botorch.optim import optimize_acqf
from botorch.test_functions import Hartmann
from botorch.acquisition import ExpectedImprovement
from gpytorch.mlls import ExactMarginalLogLikelihood
from botorch.fit import fit_gpytorch_mll

USE_MANUAL_FIT = False

torch.manual_seed(0)

hartmann6 = Hartmann(negate=False)
input_noise_std = 0.01
bounds = torch.stack([torch.zeros(6), torch.ones(6)])
def noisy_hartmann6(X): # returns the function value and noise variance estimate
    y = hartmann6(X).unsqueeze(1)
    return y + input_noise_std * torch.randn_like(y), input_noise_std ** 2 * torch.ones_like(y)

train_x = torch.rand(10, 6, dtype = torch.float64)
train_y, train_yvar = noisy_hartmann6(train_x)

def standardize(y, y_var):        
    mean = y.mean()
    std = y.std()
    return (y - mean) / std, (y_var / std ** 2)
    
def manual_fit(model, mll):
    model.train()
    mll.train()
    optimizer = torch.optim.Adam([{"params": model.parameters()}], lr=0.1)
    for _ in range(100):
        optimizer.zero_grad()
        out = model(*model.train_inputs)
        loss = -mll(out, model.train_targets)
        loss.backward()
        optimizer.step()
    model.eval()

    
for i in range(100):
    standardized_y, standardized_yvar = standardize(train_y, train_yvar)
    
    model = HeteroskedasticSingleTaskGP(train_X=train_x, train_Y=standardized_y, train_Yvar=standardized_yvar)
    
    mll = ExactMarginalLogLikelihood(model.likelihood, model)

    if USE_MANUAL_FIT:
        manual_fit(model, mll)
    else:
        fit_gpytorch_mll(mll)
    
    acqf = ExpectedImprovement(model=model, best_f=train_y.min(), maximize=False)
    
    candidate, acq_value = optimize_acqf(acq_function=acqf, bounds=bounds, q=1, num_restarts=10, raw_samples=512)
    val, noise_var = noisy_hartmann6(candidate)
    
    print("x:", candidate.detach().numpy().squeeze(), "y:", val.detach().numpy().squeeze())
    
    train_x = torch.cat((train_x, candidate), dim = 0)
    train_y = torch.cat((train_y, val), dim = 0)
    train_yvar = torch.cat((train_yvar, noise_var), dim = 0)