Shape/indexing issues with latent GP when using Predictive

I have a model that implements a latent GP as follows (relevant code shown), which later gets indexed appropriately and added into the rest of the model:

class MyModel(PyroModule);

    def __init__(self):

        ...

        self.K = gp.kernels.RBF(input_dim=1, variance=self.amp, lengthscale=self.ls)
        cov_beta = self.K(torch.DoubleTensor(days).to(device))
        cov_beta.view(-1)[:: self.game_days.shape[0] + 1] += jitter 
        self.Lff_beta = cov_beta.cholesky().detach()

        ...

    def forward(self, X, y):

        ...

        f = pyro.sample("f",
            MultivariateNormal(
                torch.zeros_like(self.all_days).to(device),
                scale_tril=self.Lff_beta
            ))

        ...

        theta = mu[player_index]
        theta += f[day_index]
        theta += g[venue_index, day_index]

        ...

This model fits just fine using SVI, but when I later go to predict (Note that data is the same that is used to fit the model (the last element is the response data y, so it is left out of the predict):

predictive = pyro.infer.Predictive(model, guide=guide, num_samples=1000, return_sites=("f",))
samples = predictive(*data[:-1])

the call to predictive fails:

    524         # Expected value
    525         theta = mu[player_index]
--> 526         theta += f[day_index]
    527         theta += g[venue_index, day_index]

IndexError: index 90 is out of bounds for dimension 0 with size 1

It seems like something different is happening to the shape of the MutlivariateNormal when using Predictive. Is this true, or am I doing something wrong?

So, if I switch to using an embedded GPRegression instead of manually with a MultivariateNormal, it seems to run:

class MyModel(PyroModule);

    def __init__(self):

        ...

        self.K = gp.kernels.RBF(input_dim=1, variance=self.amp, lengthscale=self.ls)
        self.beta = gp.models.GPRegression(
            self.game_days, None, self.K
        )

        ...

    def forward(self, X, y):

        ...

        f, _ = self.beta.model()
        f = pyro.deterministic('beta', f)

        ...

        theta = mu[player_index]
        theta += f[day_index]
        theta += g[venue_index, day_index]

        ...

So, the resulting Predictive from above runs without error with this model, however, the resulting sample trace from beta is entirely zeros, so I’ve messed up somewhere.

[48] samples['beta']
tensor([[[0., 0., 0.,  ..., 0., 0., 0.]],

        [[0., 0., 0.,  ..., 0., 0., 0.]],

        [[0., 0., 0.,  ..., 0., 0., 0.]],

        ...,

        [[0., 0., 0.,  ..., 0., 0., 0.]],

        [[0., 0., 0.,  ..., 0., 0., 0.]],

        [[0., 0., 0.,  ..., 0., 0., 0.]]])

I think I found a somewhat hacky solution to get the values of the latent GP:

...
f, f_var = self.beta.model()
beta = pyro.sample("beta", Normal(f, f_var).to_event())

however something is not quite right, as the value of beta.kernel.lengthscale is tensor([220.1832], yet the output shows a considerably smaller lengthscale:

I’m not sure what’s going, but this seems belong to a method (like forward), unless you want to use a fixed Lff_beta in your model. (In addition, in GP to make a prediction, we need to condition on the training/inducing data - but probably it is your intention to not do it). Regarding the shape error, it is usually indicating that you are missing plate or to_event somewhere.

1 Like

Ah, thanks. I was not sure whether that belonged in the __init__ with the kernel, or in forward.

I did try adding to_event to f:

f = pyro.sample(
            "f",
            MultivariateNormal(
                torch.zeros_like(self.game_days).to(device),
                scale_tril=Lff_beta,
            ).to_event(),
        )

but it still triggers the error:

    523         # Expected value
    524         theta = mu[player_index]
--> 525         theta += f[day_index]
    526         theta += g[venue_index, day_index]
    527         theta += self.gamma * elevation_1000

IndexError: index 301 is out of bounds for dimension 0 with size 1

There is still something fundamental wrong here, though, given the mismatch with the lengthscale and the output as shown above. The GP seems not to be using the lengthscale that is being estimated.

Regarding the conditioning for the GP prediction, I do pass X data when I instantiate the GPRegression. I assumed that it made predictions based on this? Do I need to use set_data as well?

I assumed that it made predictions based on this?

That’s right (no need to use set_data if you don’t change it). If you want some flexibilities, you can add a new method and use conditional utility. This utility is used across all GP modules in contrib.gp.models.

The GP seems not to be using the lengthscale that is being estimated.

I feel that you got the same issue as in the other thread.