Mixture model using ExactGP by summing Cov matrices

[luisroque]

Hi,

I want to use the ExactGP implementation to build a mixture model that combines several GPs. Example: GP1 applies to time series [0, 0, 1, 1] and GP2 applies to time series [1, 1, 1, 0]. This means that for time series 2 (index starting at 0) I would like to add both GPs: GP1+GP2. For that, I would assume that the best approach would be summing the covariances instead of summing the GPs, as:

GP(\mu_1 + \mu_2 , K_1 + K_2) = GP(\mu_1, K_1) + GP(\mu_2 , K_2)

After summing the Cov matrices I would assume that I could implement the model as:

class MultitaskGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(MultitaskGPModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.MultitaskMean(
            gpytorch.means.ConstantMean(), num_tasks=n_tasks
        )
        self.covar_module = gpytorch.kernels.MultitaskKernel(
            gpytorch.kernels.RBFKernel() + gpytorch.kernels.PeriodicKernel(), num_tasks=n_tasks, rank=1
        )

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultitaskMultivariateNormal(mean_x, covar_x)


likelihood1 = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=n_tasks)
model1 = MultitaskGPModel(train_x, train_y1, likelihood1)

likelihood2 = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=n_tasks)
model2 = MultitaskGPModel(train_x, train_y2, likelihood2)

model = gpytorch.models.IndependentModelList(model1, model2)
likelihood = gpytorch.likelihoods.LikelihoodList(model1.likelihood, model2.likelihood)

How could I add the Cov matrices following the above structure?

Thanks,
Luis

[wjmaddox]

Not sure I completely follow here, but couldn't you write the covariance and mean functions as a sum of both GP models?

[luisroque]

Hi @wjmaddox, I end up implementing it as below. You were thinking of something like this or do you have any other suggestions?

Mixing covariances:

  covs = []
  for i in range(1, n_groups+1):
      cov = gpytorch.kernels.RBFKernel() + gpytorch.kernels.PeriodicKernel()
      covs.append(cov)
  
  # apply mixtures to covariances
  selected_covs = []
  mixed_covs = []
  for i in range(n_series):
      mixture_weights = known_mixtures[i]
      for w_ix in range(n_groups):
          w = mixture_weights[w_ix]
          if w == 1.0:
              selected_covs.append(covs[w_ix])
      mixed_cov = selected_covs[0]
      for cov in range(1,len(selected_covs)):
          mixed_cov += selected_covs[cov] # because GP(cov1 + cov2) = GP(cov1) + GP(cov2)
      mixed_covs.append(mixed_cov) 
      selected_covs = [] # clear out cov list

and then creating the GPs:

class ExactGPModel(gpytorch.models.ExactGP):

  def __init__(self, train_x, train_y, likelihood, cov):
      super().__init__(train_x, train_y, likelihood)
      self.mean_module = LinearMean(train_x.size(-1))
      self.covar_module = cov
  
  def forward(self, x):
      mean_x = self.mean_module(x)
      covar_x = self.covar_module(x)
      return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
        

model_list = []
likelihood_list = []
for i in range(groups['train']['s']):
    likelihood_list.append(gpytorch.likelihoods.GaussianLikelihood())
    model_list.append(ExactGPModel(train_x, train_y[:,i], likelihood_list[i], mixed_covs[i]))

    
model = gpytorch.models.IndependentModelList(*model_list)
likelihood = gpytorch.likelihoods.LikelihoodList(*likelihood_list)