Classification predictions

[KrisNguyen135]

(Full Colab code here.) I'm trying to set up a model for a simple 3-class, 2-dimensional classification problem using a softmax likelihood like so:

class GPClassifier(ApproximateGP):
    def __init__(self, train_x, num_classes):
        variational_distribution = CholeskyVariationalDistribution(train_x.size(0))
        
        variational_strategy = UnwhitenedVariationalStrategy(
            self, train_x, variational_distribution, learn_inducing_locations=False
        )
        variational_strategy = MultitaskVariationalStrategy(
            variational_strategy, num_tasks=num_classes
        )
        
        super(GPClassifier, self).__init__(variational_strategy)
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(
            gpytorch.kernels.RBFKernel()
        )
    
    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        latent_pred = gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
        return latent_pred
    
    
model = GPClassifier(X, num_groups)
likelihood = gpytorch.likelihoods.SoftmaxLikelihood(
    num_classes=num_groups, mixing_weights=False
)
mll = VariationalELBO(likelihood, model, y.numel())

and I'm having some trouble visualizing the model's predictions on a grid.

When I plot the latent means using f_pred.loc where

model.eval()
likelihood.eval()

with torch.no_grad():
    f_pred = model(X_test)

I got the same plot for all three latents, which makes me suspect I'm doing something wrong.

When I plot the probs and logits of the output from the likelihood, I got different things corresponding to the three latents, but they are not as smooth as I would like. I'm looking for something similar to what we get when we do regression on the labels like this tutorial.

Any pointer would be appreciated. Thanks in advance!

[gpleiss]

I got the same plot for all three latents, which makes me suspect I'm doing something wrong.

You have to set batch_shape=torch.Size([num_classes]) on the CholeskyVariationalDistribution. Otherwise, the GPs for each class are going to have the same output. See the variational multi-output tutorial for more details.

When I plot the probs and logits of the output from the likelihood, I got different things corresponding to the three latents, but they are not as smooth as I would like.

I think this is just an optimization issue. I raised the learning rate to 0.1 and upped the number of iterations to 1000, and predictions came out much smoother looking.

A few other minor issues:

1. MultitaskVariationalStrategy is deprecated. Use IndependentMultitaskVariationalStrategy instead.
2. I'd use VariationalStrategy rather than UnwhitenedVariationalStrategy.
3. Adding a batch_shape = torch.Size([num_classes]) to the kernel and mean objects will give the model some additional flexibility.

[KrisNguyen135]

Thank you! The suggestions work great!