Improving prediction accuracy of GP with derivative observations

[m-julian]

Hello,

I am fitting two different functions (one with 9 dimensional input and one with 12 dimensional input) for which I have function and derivative observations. The 9D case predicts well, however the 12D case not so much. I have tried several things so far to get the 12D case working better:

1. use different optimizers (Adam and LBFGS). I get similar results with both optimizers, where the 9D case fits and predicts fine, but the 12D case does not.
2. use task noises instead of global noise. The function and derivatives should have low noise, so it does not really make a difference if I am using individual task noises or one global noise.
3. use priors for the hyperparameters (Gamma/Normal prior)
4. change the kernel - I have tried the Matern52 kernel (pull request here Matern52 grad #2512 ), but that performs worse for both 9D and 12D
5. Scale the function observations and derivatives separately, but that also gives worse results for both 9D and 12D. I am also not sure if this makes sense to do since scaling the function has an effect on the derivatives as well.

Any suggestions on what else I can try to improve the 12D results? I am curious how to scale data when dealing with both function observations and derivatives. In the example below, I have min-max normalized the inputs, standardized the outputs, and then scaled the derivatives accordingly. I have tried adding more data in the 12D case, but that does not improve results much.

Below is the example 9D/12D train and test data and example code that I have used for the datasets and prediction results:

csv_files.zip

import torch
import gpytorch
import pandas as pd
from gpytorch.constraints import GreaterThan
from gpytorch.kernels import RBFKernelGrad
from LBFGS import FullBatchLBFGS

OPTIMIZER_ITERATIONS = 500
LR = 0.1

if torch.cuda.is_available():
    print(f"{torch.cuda.device_count()} GPUs available")
    print(f"0 device name: {torch.cuda.get_device_name(0)}")

CUDA_DEVICE_COUNT = torch.cuda.device_count()
OUTPUT_DEVICE = torch.device('cuda:0') if torch.cuda.is_available() else torch.device("cpu")

class GPModelWithDerivatives(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(GPModelWithDerivatives, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMeanGrad()

        ndim = train_x.shape[-1]

        k = RBFKernelGrad(ard_num_dims=ndim, active_dims=list(range(ndim)))
        self.covar_module = gpytorch.kernels.ScaleKernel(k)

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultitaskMultivariateNormal(mean_x, covar_x)


df_train = pd.read_csv("csv_files/9d_train.csv", header=0, index_col=False)
df_test = pd.read_csv("csv_files/9d_test.csv", header=0, index_col=False)

# df_train = pd.read_csv("csv_files/12d_train.csv", header=0, index_col=False)
# df_test = pd.read_csv("csv_files/12d_test.csv", header=0, index_col=False)

print(df_train.shape)
print(df_test.shape)

input_cols = [i for i in df_train.columns if i.startswith("x")]
deriv_cols = [i for i in df_train.columns if i.startswith("dfdx")]

NDIM = len(input_cols)

# train data
train_x = torch.tensor(df_train[input_cols].values).double().to(OUTPUT_DEVICE)
train_dfdx = torch.tensor(df_train[deriv_cols].values).double().to(OUTPUT_DEVICE)
train_f = torch.tensor(df_train["f"].values).double().to(OUTPUT_DEVICE)

# test data
test_x = torch.tensor(df_test[input_cols].values).double().to(OUTPUT_DEVICE)
test_dfdx = torch.tensor(df_test[deriv_cols].values).double().to(OUTPUT_DEVICE)
test_f = torch.tensor(df_test["f"].values).double().to(OUTPUT_DEVICE)

# min-max normalize inputs, standardize outputs
a = 0
b = 1

train_x_min = train_x.min(0)[0]
train_x_max = train_x.max(0)[0]

# scale inputs
normalized_train_x = a + (((train_x - train_x_min) * (b-a)) / (train_x_max - train_x_min))
# scale outputs
train_f_mean = train_f.mean()
train_f_std = train_f.std()
standardized_train_f = ((train_f - train_f_mean) / train_f_std).reshape(-1, 1)
# scale derivatives
scaled_train_dfdx = train_dfdx * ((train_x_max - train_x_min) / ((b-a) * train_f_std))
# statck f and dfdx
stacked_scaled_train_fdfdx = torch.cat([standardized_train_f, scaled_train_dfdx], -1)

# scale same as training set
normalized_test_inputs = a + (((test_x - train_x_min) * (b-a)) / (train_x_max - train_x_min))

likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=NDIM+1, has_global_noise=True, has_task_noise=False, noise_constraint=GreaterThan(1e-7))
model = GPModelWithDerivatives(normalized_train_x, stacked_scaled_train_fdfdx, likelihood)

# use double precision, move to output device
likelihood = likelihood.double()
model = model.double()
likelihood = likelihood.to(OUTPUT_DEVICE)
model = model.to(OUTPUT_DEVICE)
mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model).to(OUTPUT_DEVICE)

# optimizer = torch.optim.Adam(model.parameters(), lr=LR)
optimizer = FullBatchLBFGS(model.parameters(), lr=LR)

model.train()
likelihood.train()

with gpytorch.settings.max_cholesky_size(100_000), gpytorch.settings.fast_computations(False, False, False):

    # commented out Adam opt
    # for i in range(1, OPTIMIZER_ITERATIONS+1):
    #     optimizer.zero_grad()
    #     output = model(normalized_train_x)
    #     loss = -mll(output, stacked_scaled_train_fdfdx)

    #     print(i, loss)

    #     loss.backward()
    #     optimizer.step()
    
    
    def closure():
        optimizer.zero_grad()
        output = model(normalized_train_x)
        loss = -mll(output, stacked_scaled_train_fdfdx)
        return loss

    loss = closure()
    loss.backward()

    for i in range(1, OPTIMIZER_ITERATIONS+1):
        
        options = {"closure": closure, "current_loss": loss, "max_ls":10}
        loss, _, _, _, _, _, _, fail = optimizer.step(options)
        
        print(i, loss)

        if fail:
            print("Convergence reached")
            break

    print("Hyperparameters")
    print("outputscale", model.covar_module.outputscale)
    print("lengthscales", model.covar_module.base_kernel.lengthscale)
    print("global noise", model.likelihood.noise)
    print()

model.eval()
likelihood.eval()

with torch.no_grad(), gpytorch.settings.max_cholesky_size(100_000), gpytorch.settings.fast_computations(False, False, False):

    predictions = likelihood(model(normalized_test_inputs))
    mean = predictions.mean.detach()
    
    # get all f predictions
    f_predictions = ((mean[:, 0] * train_f_std) + train_f_mean).squeeze()
    # get all dfdx predictions
    dfdx_predictions = mean[:, 1:] * (((b-a) * train_f_std) / (train_x_max - train_x_min))
    
    f_errors = torch.abs(test_f - f_predictions)
    dfdx_errors = torch.abs(test_dfdx - dfdx_predictions)

    print("f top 10 errors")
    print(torch.sort(f_errors, descending=True)[0][:10])
    print("f MAE")
    print(f_errors.mean())
    print()
    print("dfdx max error")
    print(dfdx_errors.max(0)[0])
    print("dfdx MAE")
    print(dfdx_errors.mean(0))

9D results:

253 tensor(-1.3004, device='cuda:0', dtype=torch.float64, grad_fn=<NegBackward0>)
254 tensor(-1.3004, device='cuda:0', dtype=torch.float64, grad_fn=<NegBackward0>)
Convergence reached
Hyperparameters
outputscale tensor(2974.1778, device='cuda:0', dtype=torch.float64,
       grad_fn=<SoftplusBackward0>)
lengthscales tensor([[2.3000, 2.2820, 3.3470, 2.6066, 2.1392, 1.2014, 2.7161, 2.1914, 0.3912]],
       device='cuda:0', dtype=torch.float64, grad_fn=<SoftplusBackward0>)
global noise tensor([1.0000e-07], device='cuda:0', dtype=torch.float64,
       grad_fn=<AddBackward0>)

f top 10 errors
tensor([6.6291, 3.3133, 2.9717, 2.1583, 2.0997, 1.6892, 1.4444, 1.3003, 1.2711,
        1.0543], device='cuda:0', dtype=torch.float64)
f MAE
tensor(0.0924, device='cuda:0', dtype=torch.float64)

dfdx max error
tensor([ 9.9158, 70.9245, 15.3375, 35.0028, 63.0137, 92.1295, 66.9040, 56.2857,
        28.0372], device='cuda:0', dtype=torch.float64)
dfdx MAE
tensor([0.5125, 0.6553, 0.3373, 0.5299, 0.4999, 0.4495, 0.6263, 0.4713, 0.3844],
       device='cuda:0', dtype=torch.float64)

12D

231 tensor(1.3412, device='cuda:0', dtype=torch.float64, grad_fn=<NegBackward0>)
232 tensor(1.3412, device='cuda:0', dtype=torch.float64, grad_fn=<NegBackward0>)
Convergence reached
Hyperparameters
outputscale tensor(291.4916, device='cuda:0', dtype=torch.float64,
       grad_fn=<SoftplusBackward0>)
lengthscales tensor([[2.9073, 1.6494, 0.4564, 2.0578, 2.5941, 1.4789, 2.9995, 2.2718, 2.0199,
         1.0809, 0.3390, 0.2243]], device='cuda:0', dtype=torch.float64,
       grad_fn=<SoftplusBackward0>)
global noise tensor([1.0000e-07], device='cuda:0', dtype=torch.float64,
       grad_fn=<AddBackward0>)

f top 10 errors
tensor([91.2320, 58.8100, 39.2446, 36.3542, 36.0043, 31.9246, 29.2634, 26.6150,
        24.6648, 24.2088], device='cuda:0', dtype=torch.float64)
f MAE
tensor(3.2475, device='cuda:0', dtype=torch.float64)

dfdx max error
tensor([ 76.6679, 142.5832, 934.0112, 168.7883,  77.4476,  74.0563, 133.3079,
         60.0463,  61.7621, 158.0040, 448.3243, 444.6349], device='cuda:0',
       dtype=torch.float64)
dfdx MAE
tensor([10.9056, 14.4664, 47.2319, 10.7409,  8.5159,  6.3225, 10.9720,  8.4276,
         5.9810, 12.3910, 32.7540, 35.9129], device='cuda:0',
       dtype=torch.float64)

[m-julian]

@jacobrgardner @gpleiss @Balandat Sorry for the tag, but I have ran out of ideas on this! Has anyone come across fitting problems when dealing with function + derivative observations?

Since the first post, I have checked the input/output data again and everything seems to be correct. I have also tried using fit_gpytorch_mll from botorch, but that also gave similar results. Not exactly sure why the fit for the 12D is that different because fitting works perfectly fine for all smaller dimensional problems that I've worked on. If you have any suggestions on things to try to improve prediction accuracy, that would be great!