How to Extract Covariance Matrix Block Diagonal?

[awang124]

I'm performing GP regression for 2-tasks (independent) on 2-dimensional inputs, i.e. my X and Y both have shape (N, 2). After training, I currently obtain my posterior mean and variance as follows:

model.eval()
likelihood.eval()
with torch.no_grad(), gpytorch.settings.fast_pred_var():
    pred = likelihood(model(X_test))
    mean = pred.mean.cpu().numpy()
    var = pred.variance.cpu().numpy()

This produces mean and var. I know var is shape (N, 2), and its columns each contain variance for a single task. However, I'm wondering if it's possible to extract each input's crosscovariance elements as well, without forming my entire test data's covariance matrix pred.covariance_matrix, of shape (2N, 2N), which is far too big for my task to store. Essentially, I'm wondering if it's possible to extract only the 2x2 block diagonal of this matrix (since I'm treating each test input independently), without the full matrix.

Apologies if I've missed an obvious way to achieve this. Any help is appreciated!

[kayween]

Hi @awang124, I am assuming your are using a multi-task GP model similar to this tutorial.

To extract the 2 x 2 block diagonal entries, you could do something as follows. The trick here is that you could get a lazy representation of the 2N x 2N covariance matrix without materializing it. Then, you can slide it like a regular tensor.

covar = pred.lazy_covariance_matrix

idx = torch.arange(0, covar.shape[0]).reshape(-1, 2)  # (2 x N, 2)
block_diag_entries = covar[idx.unsqueeze(-1), idx.unsqueeze(-2)].to_dense()  # (N, 2, 2)

print(block_diag_entries)