How can I retrieve the posterior of samples that may share samples with self.datapoints in PairwiseGP?

[florisvdf]

Hi,

It seems like the default behavior of the PairwiseGP forward and posterior methods is to construct the kernel matrix by computing the covariance using all samples in self.datapoints and all samples that are passed to either the forward or the posterior method.

botorch/botorch/models/pairwise_gp.py

Lines 1022 to 1030 in 67f34cb

	X, X_new = self._transform_batch_shape(transformed_dp, transformed_new_dp)
	covar_chol, _ = self._transform_batch_shape(self.covar_chol, X_new)
	hl, _ = self._transform_batch_shape(self.likelihood_hess, X_new)
	hlcov_eye, _ = self._transform_batch_shape(self.hlcov_eye, X_new)
	# otherwise compute predictive mean and covariance
	covar_xnew_x = self._calc_covar(X_new, X)
	covar_x_xnew = covar_xnew_x.transpose(-1, -2)
	covar_xnew = self._calc_covar(X_new, X_new)

This means that when one does inference on datapoints that share samples with self.datapoints, the kernel matrix will contain entries that correspond to duplicate samples. In many cases, this results in a kernel matrix that is no longer PSD.

What is the intended way to predict the posterior distribution over samples that also exist in self.datapoints?

Thanks!

[ItsMrLin]

Hi @florisvdf, thanks for reaching out. Inside PairwiseGP, we try to handle numeric stability carefully so that by simply calling model.posterior(train_X) should give you the training posterior (given by laplace approximation) just fine, especially with the default kernel.

If you wish you may explicitly calculate the training data posterior as N(model.utility, torch.inverse(model.covar_inv + model.likelihood_hess)) but it should just give your the same thing as model.posterior.

Is there a particular example you are running into issue with using PairwiseGP?

[florisvdf]

Hi @ItsMrLin, thanks for the fast reply. I'm using a composite kernel as the model.covar_module. In my particular use case, I have a tensor named train_inputs that stores the samples used for training and a tensor named inputs that stores both the training samples and the test samples. I train the PairwiseGP using train_inputs, and when the training is finished, I call model.forward to generate the posterior of the inputs tensor, like so:

model.eval()
model.likelihood.eval()

with torch.no_grad():
    posterior = model.forward(inputs)

This then raises

linear_operator.utils.errors.NotPSDError: Matrix not positive definite after repeatedly adding jitter up to 1.0e-01.

However, if I create a different tensor storing only the test samples, test_inputs, and then generate the posterior with model.forward(test_inputs), this error does not occur.

[ItsMrLin]

Depending on what base kernels are used for the composite kernel (e.g., we have seen linear kernels become numerically unstable with PairwiseGP). Here are a few remedies on top of my head that you may try:

• Normalized the train_inputs and test_inputs to be within the unit hypercube [0, 1]^d if you haven't done this already
• Try a different, numerically stabler kernel
• Consider adding priors on the kernel's hyperparameters to prevent extreme values in resulting covariance matrix
• When appropriate, passing in a bigger jitter when constructing a PairwiseGP

[florisvdf]

Hi @ItsMrLin, I think I should have provided a bit more detail about the problem that I’m running into and how this relates to the implementation of PairwiseGP. To illustrate the problem, I list below three inference scenarios, i.e. predicting a posterior after training the GP.

Scenario 3 is where I think problems arise. For that scenario, I wrote what I think reflects the behavior of PairwiseGP, and after that what I think should be the intended behavior.

I should also note that I’m a newbie when it comes to GPs and that I’m definitely not well versed in linear algebra, so please let me know if I’m making any conceptual mistakes.

$X_{t}$, training samples

$X_{s}$, test samples (no samples shared with $X_{t}$)

$X_{n}$, the union of the training samples, $X_{t}$, and the test samples, $X_{s}$

Scenario 1: Passing only train inputs

$$ \begin{gather} f(X_{t}) ∼ N \Bigl(m(X_{t}), K(X_{t}, X_{t})\Bigr) \end{gather} $$

Retrieve $f(X_{t})$, the posterior over the training samples.

Scenario 2: Passing only test inputs

$$ \begin{gather} \begin{bmatrix} f(X_{t}) \ f(X_{s}) \end{bmatrix} ∼N \left(\begin{bmatrix} m(X_{t}) \ m(X_{s} \end{bmatrix}, \begin{bmatrix} K(X_{t}, X_{t}) & K(X_{t}, X_{s}) \\ K(X_{s}, X_{t}) & K(X_{s}, X_{s}) \end{bmatrix}\right) \end{gather} $$

Retrieve $f(X_{s})$, the posterior over the test samples.

Scenario 3: Passing new inputs that may contain training samples

What I think happens in PairwiseGP

$$ \begin{gather} \begin{bmatrix} f(X_{t}) \ f(X_{n}) \end{bmatrix} ∼N \left(\begin{bmatrix} m(X_{t}) \ m(X_{n} \end{bmatrix}, \begin{bmatrix} K(X_{t}, X_{t}) & K(X_{t}, X_{n}) \\ K(X_{n}, X_{t}) & K(X_{n}, X_{n}) \end{bmatrix}\right) \end{gather} $$

Retrieve $f(X_{n})$, the posterior over $X_{n}$.

I think this is problematic, because our kernel matrix now contains duplicate rows and columns for all samples shared between $X_{t}$ and $X_{s}$, which I think causes the kernel matrix to lose rank, become ill-conditioned and in some cases no longer positive semi definite.

What I think should be the intended behavior

Create $X_{n*}$, such that:

$X_{n^*} = { x_{i} \in X_{n} \mid x_{i} \notin X_{t} }$

And then do

$$ \begin{gather} \begin{bmatrix} f(X_{t}) \ f(X_{n*}) \end{bmatrix} ∼N \left(\begin{bmatrix} m(X_{t}) \ m(X_{n*} \end{bmatrix}, \begin{bmatrix} K(X_{t}, X_{t}) & K(X_{t}, X_{n*}) \\ K(X_{n*}, X_{t}) & K(X_{n*}, X_{n*}) \end{bmatrix}\right) \end{gather} $$

Then retrieve $f(X_{n})$ by combining the posteriors from the shared samples from $f(X_{t})$ with $f(X_{n*})$.

I hope I explained the issue clearly. Am I correct in my understanding that this is the current behavior of PairwiseGP? And if so, what could I do to achieve the desired behavior?

[florisvdf]

Hi @ItsMrLin, do you have the time to take a look at this?

[ItsMrLin]

Hi @florisvdf, sorry for the late reply! You are correct that by having a covariance matrix over a (partially) overlapping same set of training points as would lead to a lower rank (though not necessarily non-PSD). In practice, we add a small "jitter" to the diagonal of the matrix like $M_{\varepsilon} = M + \varepsilon I_n$ and use methods like Cholesky decomposition instead of direct matrix inversion to avoid potential numeric error accumulation and propagation. These measures allow us to calculate the posterior with this supposed-to-be ill-conditioned matrix, and works for me on a test problems using 50 points and 100 comparisons where I try to use the same 50 points as the test points and try to calculate model.posterior (though seems not so for you).

This is why I was suggesting perhaps some of the easy fixes might just be trying to use a bigger jitter if it's acceptable for you or try to use a different kernel or consider only run this calculation on a new set of points.

If these changes do not work or are not applicable to your use case, would you be able to share a minimum reproducible example of the error and I'm happy to take a closer look at it :)

[florisvdf]

Hi @ItsMrLin, thanks again for your time. I tried adding jitter up to around 1e-1 and that unfortunately didn't solve my problem. I ended up removing training samples from $X_{n}$, and then retrieved the posterior over training samples by

scale = (
    model.covar_module.outputscale.unsqueeze(-1).unsqueeze(-1).detach()
)
posterior = gpytorch.distributions.MultivariateNormal(
    mean=model.utility,
    covariance_matrix=RootLinearOperator(
        _scaled_psd_safe_cholesky(
            matrix=model.covar,
            scale=scale,
            jitter=self._jitter,
        )
    ),
)

This works for now, though I'm still somewhat curious as to why the default behavior of the PairwiseGP is not to eliminate overlapping training samples from the kernel matrix. Thanks again for helping me find a solution :).

[ItsMrLin]

Hi @florisvdf, sorry to hear that bigger jitter doesn't help. We have seen similar situation happening before when using an impropriate kernel. Using a different covar_module could be another thing worth trying.

re: why the default behavior of the PairwiseGP is not to eliminate overlapping training samples from the kernel matrix -- in general, we are trying to stay faithful to the user's input. When the user is asking for the posterior over a set of points (including the training data or not), their expectation is to receive the posterior over the set of points, instead of a subset of them.

If you end up being able to calculate the posterior of only training samples separately from the test sample, you may switch the model to the training model to obtain it via:

model.train()
posterior = model(train_X)

This should give the correct posterior without any errors as it doesn't operate on the ill-conditioned covariance that we see before (code pointer).

Note that this is different from MultivariateNormal(model.utility, model.covar) as model.covar is directly coming out of the kernel function (i.e., the prior covariance), not the posterior.

[florisvdf]

Hi ItsMrLin, thank you for addressing the source of my confusion. This pretty much answers my original question.

I'm using a rather intricately composed composite kernel from an existing GP with a Gaussian likelihood. This kernel looks something like ScaleKernel(weight_i * RBFKernel(x_a) + weight_j * CompositeKernel(x_b) + weight_k * RBFKernel (x_c)), where weights weight_i, weight_j, weight_k sum to one, x_a, x_b, x_c are multimodal inputs, and CompositeKernel is yet another composite kernel which is a multiplication of three RBFKernels.

I don't want to make too many changes to this kernel, but thanks for pointing out that tweaking this kernel is one path to improving numerical instability.

One small detail I want to mention that I noticed when implementing your suggestion is that botorch throws in error passing train_X to the model in train mode, even when train_X and model.datapoints seems identical.

matches_train = torch.equal(self.datapoints, train_X) # dtpyes are also equal
with torch.no_grad():
    if matches_train:
        model.train()
        posterior = model(train_X)

Results in
RuntimeError: Must train on training data

replacing train_X with model.datapoints fixes this issue. I'm not sure if this is something specific to my code. In any case, I'll just continue with model.datapoints. Thanks a lot for your help!