Spectral Mixture Kernel Performance Issues

[mbelalsh]

Hello @gpleiss. I trained the Multi-task GP model with RQKernel, and got decent results; the negative mll went down to -6.25. I read some of the papers of AG Wilson wherein he said that RBF (and some popular) kernels do not discover patterns; they are just filters. Also, you recommended that I should read one of the papers of AG Wilson related to Pattern discovery and extrapolation.

• After that, I thought of giving it a try, but the results are not good as to the negative mll gets stuck around 15.69 which is worse than what I was able to achieve with RQKernel (-6.25).
• I know that the parameters of SM kernel are sensitive to initialization so I initialized those using initialize_from_data method. Still the negative mll getting stuck around 15.69 after 20 epochs and never improves after that.
• I also switched everything to double precision to avoid any numerical stability issues.
• I did not change anything in my code for both Kernels except some kernel-related steps (like initialization for SpectralMixtureKernel, etc.). Other than that, I kept everything the same (like inducing points, etc.) to have a fair comparison.

Any practical advice to get better performance?

• My main Multi-task GP model is given below. I am using RQKernel or SpectralMixtureKernel depending on learnKernel flag.

class MultitaskGPModel(gpytorch.models.ApproximateGP):
    def __init__(self,
                 num_tasks=1,
                 num_inducing_points=200,
                 num_inputs=1,
                 learnKernel:bool = 0,    
                 base_kernel_type=gpytorch.kernels.RQKernel,
                 inducing_points=None,
                 ):
        # We have to mark the CholeskyVariationalDistribution as batch
        # so that we learn a variational distribution for each task
        """
        learnKernel: 1 if you want to use Spectral Mixture Kernel & 0 otherwise  
        """
        if inducing_points is None:
            print(num_inducing_points)
            inducing_points = torch.rand(num_tasks, num_inducing_points, num_inputs, dtype=torch.double)

        variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(
            inducing_points.size(-2), batch_shape=torch.Size([num_tasks])
        )

        # We have to wrap the VariationalStrategy in a MultitaskVariationalStrategy
        # so that the output will be a MultitaskMultivariateNormal rather than a batch output
        variational_strategy = gpytorch.variational.IndependentMultitaskVariationalStrategy(
            gpytorch.variational.VariationalStrategy(
                self, inducing_points, variational_distribution, learn_inducing_locations=True
            ), num_tasks=num_tasks
        )

        super().__init__(variational_strategy)

        # The mean and covariance modules should be marked as batch so we learn a different set of hyperparameters
        # self.mean_module = gpytorch.means.ConstantMean(
        #     batch_shape=torch.Size([num_tasks])
        # )
        self.mean_module = gpytorch.means.ZeroMean(
            batch_shape=torch.Size([num_tasks])
        )
        
        if learnKernel == 0:
            self.covar_module = gpytorch.kernels.ScaleKernel(
                base_kernel_type(batch_shape=torch.Size([num_tasks])),
                batch_shape=torch.Size([num_tasks]),
                ard_num_dims=num_inputs,
            )
        else:
            self.covar_module = gpytorch.kernels.SpectralMixtureKernel(
                num_mixtures=6, ard_num_dims=num_inputs, 
                batch_shape=torch.Size([num_tasks]),
            )   

    def forward(self, x):
        # The forward function should be written as if we were dealing with each output
        # dimension in batch
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

• My training loop is given below. As you can see that I am initializing SpectralMixtureKernel parameters. The Functional Kernel Learning paper also discuss some initialization related issues in its Appendix D. That is why, I was careful with the initialization.

def train(self, train_epochs, dirName:str):
        loss_fn = VariationalELBO(self.likelihood, self.model, num_data=len(self.ds_train))

        optimizer = torch.optim.Adam([
            {"params": self.model.parameters()},
            {"params": self.likelihood.parameters()},
        ], lr=self.lr)

        #lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.3)

        self.model.train()
        self.likelihood.train()
        # Initializing parameters of SM 
        self.model.covar_module.initialize_from_data(torch.from_numpy(self.X_train).double().to(device),
                                                     torch.from_numpy(self.y_train).double().to(device)) 

        for i in range(train_epochs):
            avg_loss = []
            for x, y in self.dl_train:
                x, y = x.to(device), y.to(device)
                y_pred = self.model(x)

                optimizer.zero_grad()
                loss = -loss_fn(y_pred, y)
                loss.backward()
                optimizer.step()
                avg_loss.append(loss.item())

            # if i > 0 and i % 50 == 0:
            #     lr_scheduler.step()

            if i % 10 == 0:
                print(f"Test Results - Epoch: {i}, Train Likelihood: {np.mean(avg_loss):.2f}")
                self.save(f"{i}_{np.mean(avg_loss):.2f}", dirName)

[Turakar]

AFAIK, the marginal log likelihood is not a good measure of performance accross different kernel choices. Instead, one should compare a proper metric like MSE or something that also takes the uncertainty into account (docs on metrics in GPyTorch).

[mbelalsh]

Thanks for this.

• I did hear that before from Geoff. Also, I read a paper on that by AG Wilson.
• With that being said, I have measured their performances. The one with minimum negative mll (i.e., the model with RQKernel) performs better than the one with SpectralMixtureKernel in terms of RMSE on testing data.

[gpleiss]

The spectral mixture kernel is very sensitive to initialization, and optimization can be tricky. Sometimes the initialize_from_data method works; sometimes it doesn't.

I would try doing 100 random initializations, and seeing which one works best, especially using the metrics that @Turakar suggests.

[mbelalsh]

Thanks Geoff. I will try this and get back to you.