Storing covariance matrix in optimization

[cmorasar]

Hi all,

In my particular application, I am interested in optimizing a GP model which has a neural network as the mean function by using a MSE-based loss function. The posterior mean of the GP requires the computation of the inverse of the training covariance matrix. Since I do not optimize the kernel parameters (I manually fix them at the beginning of optimization), this matrix remains constant throughout optimization and so I would like to store it at the very beginning to save the cost of inverting it in each forward pass. How can I do it?

This is my attempt to do so in the mean_cache function, which is used in prediction:

@cached(name="mean_cache")
def mean_cache(self, model):
       mvn = self.likelihood(self.train_prior_dist, self.train_inputs)
       train_mean = mvn.loc

       if model.train_covar_inv == None:
           ident = torch.eye(self.train_inputs[0].shape[0])
           model.train_covar_inv = mvn.lazy_covariance_matrix.inv_matmul(ident)

       train_labels_offset = (self.train_labels - train_mean).unsqueeze(-1)
       mean_cache = model.train_covar_inv.matmul(train_labels_offset).squeeze(-1)

As you can see, I created the attribute model.train_covar_inv where I aim to store the inverse of the covariance matrix in the first forward pass by using the inv_matmul method on a identity matrix. The reason for which I do it this way is because I did not see any method for computing the inverse solely.

However, I do not get the same results as when I use the original code:

@cached(name="mean_cache")
def mean_cache(self, model):
       mvn = self.likelihood(self.train_prior_dist, self.train_inputs)
       train_mean, train_train_covar = mvn.loc, mvn.lazy_covariance_matrix

       train_labels_offset = (self.train_labels - train_mean).unsqueeze(-1)
       mean_cache = train_train_covar.inv_matmul(train_labels_offset).squeeze(-1)

I would like to know why this might be happening and if there's something that I'm missing.

Thank you.