Repository: cornellius-gp/gpytorch

Issue:
Slow convergence of MultiTask regressor and unreliable results.
Hello everyone,

Currently I am trying to learn a model to predict multiple real-valued properties of a cellular image. I have a baseline that uses CNN (e.g. resnet) as feature extractor and FC head as predictor.
I would like to try to use GP as predictor and I implemented some simple network following multitask tutorial and gpshot repo ( https://github.com/BayesWatch/deep-kernel-transfer/blob/master/methods/gpshot_regression.py ). 

Simplified code of GP layer and the main model.

```
class GPBasedModel(torch.nn.Module):
    def __init__(self, model_name, likelihood, out_features=10):
        super(GPBasedModel, self).__init__()
        
        self.body, feature_dim = create_body(...)

        train_feats = torch.FloatTensor(torch.zeros(32, feature_dim)).cuda()
        train_y = torch.FloatTensor(torch.zeros(32,out_features)).cuda()
        
        self.gp_layer = BatchIndependentMultitaskGPModel(train_feats, train_y, likelihood, out_dim)
        
        self.feature_norm = torch.nn.Sequential(
                                                torch.nn.BatchNorm1d(input_dim)
                                                )
            
    def forward(self, batch):
        features = self.body(batch)
        features = features.view(features.size(0), -1)

        features = self.feature_norm(features) #z-score features
                
        self.gp_layer.set_train_data(inputs=features)
        res = self.gp_layer(features)
        
        return res 

class BatchIndependentMultitaskGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood, out_dim):
        super().__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean(batch_shape=torch.Size([out_dim]))
        self.covar_module = gpytorch.kernels.ScaleKernel(
            gpytorch.kernels.RBFKernel(batch_shape=torch.Size([out_dim])),
            batch_shape=torch.Size([out_dim])
        )

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultitaskMultivariateNormal.from_batch_mvn(
            gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
        )

likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=10)
mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model.gp_layer)

optimizer = torch.optim.Adam([
    {'params': model.body.parameters(), 'lr': 1e-3},
    {'params': model.meta_head.parameters(), 'lr': 1e-3},
    {'params': model.gp_layer.hyperparameters(), 'lr': 1e-1},
    ], lr=1e-3)

```

likelihood defined as MultitaskGaussian and loss function as ExactMarginalLogLikelihood. 

However, results look really bad: convergence is painfully slow comparing to the model with FC head & results are much worse (distribution of predicted values does not look alike real ones) + MLL is relatively high (~1000-2000).

I have tried to z-score both target & features (add BN  for features) but it didnt change results a lot: despite MLL dropped to ~45 but it stays around this value during training.

So my question is kinda simple: is smth wrong with the data, my code or GP model is not suitable for that case? 

Thanks in advance!


Discussion:
My guess is there might be some trouble with how you are handling the training data setting. In general, you should be able to train exact GPs with a DNN feature extractor without having to use set_train_data. For example, I notice that your code never sets the train labels, so they’ll always be all zeros won’t they?

The way I’d recommend using the feature extractor is to set the training data to the original input data / labels, and then having your GP forward data through the feature extractor / body in the GP forward method itself, rather than in an outside Module that then updates the GP’s training data. See our exact DKL examples for an example of this exact kind of model.

Resolving patch:
diff --git a/gpytorch/mlls/exact_marginal_log_likelihood.py b/gpytorch/mlls/exact_marginal_log_likelihood.py
--- a/gpytorch/mlls/exact_marginal_log_likelihood.py
+++ b/gpytorch/mlls/exact_marginal_log_likelihood.py
@@ -63,7 +63,7 @@ def forward(self, function_dist, target, *params):
         res = self._add_other_terms(res, params)
 
         # Scale by the amount of data we have
-        num_data = target.size(-1)
+        num_data = function_dist.event_shape.numel()
         return res.div_(num_data)
 
     def pyro_factor(self, output, target, *params):


Test patch: