SingleTaskGP On Multiple Outputs - Train Inputs Has Extra Unexpected Dimension?

[bchen0]

When training a multitask GP on multiple outputs, one thing that is confusing me is the shape of gp.train_inputs.

For a GP with a single output, gp.train_inputs is of shape NxD, where N is the number of data points and D is the shape of the input vectors. However, for a GP with K outputs, gp.train_inputs is of shape KxNxD instead of NxD, and I am not sure why. Furthermore, gp.train_inputs[0] = gp.train_inputs[1] = ... gp.train_inputs[K-1].

This has led me to wonder if I am using SingleTaskGP incorrectly.

train_X = torch.Tensor([0, 0.25, 0.5, 0.75, 1,])
Y = torch.stack([
    torch.sin(train_X * (2 * math.pi)),
    torch.cos(train_X * (2 * math.pi))
], axis=1)
train_X = train_X.unsqueeze(1)

gp = SingleTaskGP(
    train_X=train_X,
    train_Y=Y,
    train_Yvar = torch.full_like(Y, 1e-6),
    input_transform=None,
    outcome_transform=None
)
mll = ExactMarginalLogLikelihood(gp.likelihood, gp)
fit_gpytorch_mll(mll)

print(gp.train_inputs[0].shape)

Just to make sure I am understanding:

1. Is the above code indeed generating a GP to predict Y = [sin x, cos x] (it seems to, but want to make sure I'm not missing something)?
2. Is there a case where gp.train_inputs[0] would be different from gp.train_inputs[1] and so forth?

[saitcakmak]

Hi @bchen0. train_inputs is a property of the underlying GPyTorch model, which follows different shape conventions than BoTorch. Under the hood, a SingleTaskGP is a batched ExactGP model, where the output dimension is represented as the right-most batch dimension (dim=-3). So, what you're observing is just as expected.

Is there a case where gp.train_inputs[0] would be different from gp.train_inputs[1] and so forth?

If you manually constructed a batched single-output model, then each train_inputs[i] could be different. Something like

gp = SingleTaskGP(
    train_X=torch.rand(batch, q, d),
    train_Y=torch.randn(batch, q, m),
)

will have train_inputs of shape batch x (m) x q x d (m dimension is only present if m>1) where each batch can have different inputs.

[bchen0]

Great, thanks!