Mulit-Task Normalization problem

[ToennisStef]

Hi I am trying to do a Multi-Fidelity optimization using the MultiTaskGP, but I have issues when normalizing the training data.

My training data is

x-training data

train_x = tensor([[0.5582, 1.0000],
        [0.1914, 1.0000],
        [0.3829, 1.0000],
        [0.8596, 1.0000],
        [0.9101, 1.0000],
        [0.3395, 1.0000],
        [0.0159, 1.0000],
        [0.7267, 1.0000],
        [0.8438, 2.0000],
        [0.0147, 2.0000],
        [0.2824, 2.0000],
        [0.5763, 2.0000],
        [0.6735, 2.0000],
        [0.4675, 2.0000],
        [0.2349, 2.0000],
        [0.9058, 2.0000],
        [0.1193, 3.0000],
        [0.6885, 3.0000],
        [0.8751, 3.0000],
        [0.3078, 3.0000],
        [0.4897, 3.0000],
        [0.8265, 3.0000],
        [0.5150, 3.0000],
        [0.1762, 3.0000],
        [0.3495, 0.0000],
        [0.4088, 0.0000],
        [0.7570, 0.0000]])

y-training data:

train_y = tensor([[-1.3938e-01],
        [ 2.0765e+00],
        [-8.6975e-02],
        [-1.6114e-01],
        [-3.7411e+00],
        [-2.9439e-02],
        [-3.6055e+00],
        [ 2.2405e+00],
        [ 1.3858e+00],
        [ 2.6987e+00],
        [ 3.1281e+00],
        [ 1.2618e+00],
        [ 3.6563e+00],
        [ 1.6774e+00],
        [ 3.5674e+00],
        [-4.9250e+00],
        [ 9.2543e+00],
        [ 5.1593e+00],
        [ 1.0825e-01],
        [ 6.9252e+00],
        [ 4.6827e+00],
        [ 3.2985e+00],
        [ 4.3634e+00],
        [ 8.6610e+00],
        [-1.8239e-03],
        [-1.6121e-01],
        [ 6.0207e+00]])

These are exact observations form the multifidelity forrester function defined in Analytical Benchmark Problems for Multifidelity Optimization Methods.

when i do the input transform with the Normalize and outcome_tranform with Standardize the model fit is not correct:

model = MultiTaskGP(train_X=train_x, train_Y=train_y, task_feature=-1, input_transform=Normalize(d=1+1), outcome_transform=Standardize(m=1))
mll = ExactMarginalLogLikelihood(model.likelihood, model)
fit_gpytorch_mll(mll)

When i plot the results with:

eval_x = torch.linspace(0, 1, 100).unsqueeze(-1)
eval_int0 = torch.zeros((100,1))
eval_int1 = torch.ones((100,1))
eval_int2 = torch.ones((100,1))*2
eval_int3 = torch.ones((100,1))*3

eval_x0 = torch.cat((eval_x, eval_int0), dim=-1)
eval_x1 = torch.cat((eval_x, eval_int1), dim=-1)
eval_x2 = torch.cat((eval_x, eval_int2), dim=-1)
eval_x3 = torch.cat((eval_x, eval_int3), dim=-1)
#print(eval_x)
pred0 = model.posterior(eval_x0)
pred1 = model.posterior(eval_x1)
pred2 = model.posterior(eval_x2)
pred3 = model.posterior(eval_x3)

lower0, upper0 = pred0.confidence_region()
lower1, upper1 = pred1.confidence_region()
lower2, upper2 = pred2.confidence_region()
lower3, upper3 = pred3.confidence_region()

from matplotlib import pyplot as plt

with torch.no_grad():
    plt.plot(eval_x[:,0], pred0.mean, label="1")
    plt.fill_between(eval_x[:,0], lower0, upper0, alpha=0.5)
    plt.plot(eval_x[:,0], pred1.mean, label="2")
    plt.fill_between(eval_x[:,0], lower1, upper1, alpha=0.5)
    plt.plot(eval_x[:,0], pred2.mean, label="3")
    plt.fill_between(eval_x[:,0], lower2, upper2, alpha=0.5)
    plt.plot(eval_x[:,0], pred3.mean, label="4")
    plt.fill_between(eval_x[:,0], lower3, upper3, alpha=0.5)
    plt.legend()

I get this result:

but when I omit the input and outcome transform I get a much better fit:

Why is it so different? Has it something to do with the strong priors?
The normal "x" values are already normalized, just the discrete task_feature values get normalized.
What is also curious to me is that posterior predictions of the fidelities or tasks 0 to 2 are identical.
Are the task_feature values rounded somehow?
My assumption of the underlying problem is that the task_feature values internally getting transformed to integer values, and since I have normalized all input values the, the unnormalized task_feature values of 1 and 2 become 0,333... and 0,666... after normalization then internally somehow get transformed to integer values and become 0.
Is that correct?
So should you only apply the normalization to the "x" values and not the task_feature values?

Thanks to anyone reading my post! I'm grateful for any help.

[esantorella]

I don't think you should be normalizing the task values. This will raise an exception in newer versions of BoTorch. When I tried to run this, it errored out at pred2 = model.posterior(eval_x2) with

ValueError: Received invalid raw task values. Expected raw value to be in {0, 1}, but got unexpected task values: {2}.

I'm not an expert on multi-task models, so I'll see if I can get back to you on that.

[Balandat]

Correct, you should not normalize the task indices here. The newer version of botorch errors out to avoid that this is done inadvertently.

[ToennisStef]

I guess i'll better update my libraries. Thank you very much!