prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN 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!
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `gpytorch/mlls/exact_marginal_log_likelihood.py`
  Content:
  ```
  1 #!/usr/bin/env python3
  2 
  3 from ..distributions import MultivariateNormal
  4 from ..likelihoods import _GaussianLikelihoodBase
  5 from .marginal_log_likelihood import MarginalLogLikelihood
  6 
  7 
  8 class ExactMarginalLogLikelihood(MarginalLogLikelihood):
  9     """
  10     The exact marginal log likelihood (MLL) for an exact Gaussian process with a
  11     Gaussian likelihood.
  12 
  13     .. note::
  14         This module will not work with anything other than a :obj:`~gpytorch.likelihoods.GaussianLikelihood`
  15         and a :obj:`~gpytorch.models.ExactGP`. It also cannot be used in conjunction with
  16         stochastic optimization.
  17 
  18     :param ~gpytorch.likelihoods.GaussianLikelihood likelihood: The Gaussian likelihood for the model
  19     :param ~gpytorch.models.ExactGP model: The exact GP model
  20 
  21     Example:
  22         >>> # model is a gpytorch.models.ExactGP
  23         >>> # likelihood is a gpytorch.likelihoods.Likelihood
  24         >>> mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)
  25         >>>
  26         >>> output = model(train_x)
  27         >>> loss = -mll(output, train_y)
  28         >>> loss.backward()
  29     """
  30 
  31     def __init__(self, likelihood, model):
  32         if not isinstance(likelihood, _GaussianLikelihoodBase):
  33             raise RuntimeError("Likelihood must be Gaussian for exact inference")
  34         super(ExactMarginalLogLikelihood, self).__init__(likelihood, model)
  35 
  36     def _add_other_terms(self, res, params):
  37         # Add additional terms (SGPR / learned inducing points, heteroskedastic likelihood models)
  38         for added_loss_term in self.model.added_loss_terms():
  39             res = res.add(added_loss_term.loss(*params))
  40 
  41         # Add log probs of priors on the (functions of) parameters
  42         for name, module, prior, closure, _ in self.named_priors():
  43             res.add_(prior.log_prob(closure(module)).sum())
  44 
  45         return res
  46 
  47     def forward(self, function_dist, target, *params):
  48         r"""
  49         Computes the MLL given :math:`p(\mathbf f)` and :math:`\mathbf y`.
  50 
  51         :param ~gpytorch.distributions.MultivariateNormal function_dist: :math:`p(\mathbf f)`
  52             the outputs of the latent function (the :obj:`gpytorch.models.ExactGP`)
  53         :param torch.Tensor target: :math:`\mathbf y` The target values
  54         :rtype: torch.Tensor
  55         :return: Exact MLL. Output shape corresponds to batch shape of the model/input data.
  56         """
  57         if not isinstance(function_dist, MultivariateNormal):
  58             raise RuntimeError("ExactMarginalLogLikelihood can only operate on Gaussian random variables")
  59 
  60         # Get the log prob of the marginal distribution
  61         output = self.likelihood(function_dist, *params)
  62         res = output.log_prob(target)
  63         res = self._add_other_terms(res, params)
  64 
  65         # Scale by the amount of data we have
  66         num_data = target.size(-1)
  67         return res.div_(num_data)
  68 
  69     def pyro_factor(self, output, target, *params):
  70         import pyro
  71 
  72         mll = target.size(-1) * self(output, target, *params)
  73         pyro.factor("gp_mll", mll)
  74         return mll
  75
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: 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 @@
           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):
  
verification_info: {"golden_diff": "diff --git a/gpytorch/mlls/exact_marginal_log_likelihood.py b/gpytorch/mlls/exact_marginal_log_likelihood.py\n--- a/gpytorch/mlls/exact_marginal_log_likelihood.py\n+++ b/gpytorch/mlls/exact_marginal_log_likelihood.py\n@@ -63,7 +63,7 @@\n         res = self._add_other_terms(res, params)\n \n         # Scale by the amount of data we have\n-        num_data = target.size(-1)\n+        num_data = function_dist.event_shape.numel()\n         return res.div_(num_data)\n \n     def pyro_factor(self, output, target, *params):\n", "issue": "Slow convergence of MultiTask regressor and unreliable results.\nHello everyone,\r\n\r\nCurrently 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.\r\nI 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 ). \r\n\r\nSimplified code of GP layer and the main model.\r\n\r\n```\r\nclass GPBasedModel(torch.nn.Module):\r\n    def __init__(self, model_name, likelihood, out_features=10):\r\n        super(GPBasedModel, self).__init__()\r\n        \r\n        self.body, feature_dim = create_body(...)\r\n\r\n        train_feats = torch.FloatTensor(torch.zeros(32, feature_dim)).cuda()\r\n        train_y = torch.FloatTensor(torch.zeros(32,out_features)).cuda()\r\n        \r\n        self.gp_layer = BatchIndependentMultitaskGPModel(train_feats, train_y, likelihood, out_dim)\r\n        \r\n        self.feature_norm = torch.nn.Sequential(\r\n                                                torch.nn.BatchNorm1d(input_dim)\r\n                                                )\r\n            \r\n    def forward(self, batch):\r\n        features = self.body(batch)\r\n        features = features.view(features.size(0), -1)\r\n\r\n        features = self.feature_norm(features) #z-score features\r\n                \r\n        self.gp_layer.set_train_data(inputs=features)\r\n        res = self.gp_layer(features)\r\n        \r\n        return res \r\n\r\nclass BatchIndependentMultitaskGPModel(gpytorch.models.ExactGP):\r\n    def __init__(self, train_x, train_y, likelihood, out_dim):\r\n        super().__init__(train_x, train_y, likelihood)\r\n        self.mean_module = gpytorch.means.ConstantMean(batch_shape=torch.Size([out_dim]))\r\n        self.covar_module = gpytorch.kernels.ScaleKernel(\r\n            gpytorch.kernels.RBFKernel(batch_shape=torch.Size([out_dim])),\r\n            batch_shape=torch.Size([out_dim])\r\n        )\r\n\r\n    def forward(self, x):\r\n        mean_x = self.mean_module(x)\r\n        covar_x = self.covar_module(x)\r\n        return gpytorch.distributions.MultitaskMultivariateNormal.from_batch_mvn(\r\n            gpytorch.distributions.MultivariateNormal(mean_x, covar_x)\r\n        )\r\n\r\nlikelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=10)\r\nmll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model.gp_layer)\r\n\r\noptimizer = torch.optim.Adam([\r\n    {'params': model.body.parameters(), 'lr': 1e-3},\r\n    {'params': model.meta_head.parameters(), 'lr': 1e-3},\r\n    {'params': model.gp_layer.hyperparameters(), 'lr': 1e-1},\r\n    ], lr=1e-3)\r\n\r\n```\r\n\r\nlikelihood defined as MultitaskGaussian and loss function as ExactMarginalLogLikelihood. \r\n\r\nHowever, 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).\r\n\r\nI 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.\r\n\r\nSo my question is kinda simple: is smth wrong with the data, my code or GP model is not suitable for that case? \r\n\r\nThanks in advance!\n", "before_files": [{"content": "#!/usr/bin/env python3\n\nfrom ..distributions import MultivariateNormal\nfrom ..likelihoods import _GaussianLikelihoodBase\nfrom .marginal_log_likelihood import MarginalLogLikelihood\n\n\nclass ExactMarginalLogLikelihood(MarginalLogLikelihood):\n    \"\"\"\n    The exact marginal log likelihood (MLL) for an exact Gaussian process with a\n    Gaussian likelihood.\n\n    .. note::\n        This module will not work with anything other than a :obj:`~gpytorch.likelihoods.GaussianLikelihood`\n        and a :obj:`~gpytorch.models.ExactGP`. It also cannot be used in conjunction with\n        stochastic optimization.\n\n    :param ~gpytorch.likelihoods.GaussianLikelihood likelihood: The Gaussian likelihood for the model\n    :param ~gpytorch.models.ExactGP model: The exact GP model\n\n    Example:\n        >>> # model is a gpytorch.models.ExactGP\n        >>> # likelihood is a gpytorch.likelihoods.Likelihood\n        >>> mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)\n        >>>\n        >>> output = model(train_x)\n        >>> loss = -mll(output, train_y)\n        >>> loss.backward()\n    \"\"\"\n\n    def __init__(self, likelihood, model):\n        if not isinstance(likelihood, _GaussianLikelihoodBase):\n            raise RuntimeError(\"Likelihood must be Gaussian for exact inference\")\n        super(ExactMarginalLogLikelihood, self).__init__(likelihood, model)\n\n    def _add_other_terms(self, res, params):\n        # Add additional terms (SGPR / learned inducing points, heteroskedastic likelihood models)\n        for added_loss_term in self.model.added_loss_terms():\n            res = res.add(added_loss_term.loss(*params))\n\n        # Add log probs of priors on the (functions of) parameters\n        for name, module, prior, closure, _ in self.named_priors():\n            res.add_(prior.log_prob(closure(module)).sum())\n\n        return res\n\n    def forward(self, function_dist, target, *params):\n        r\"\"\"\n        Computes the MLL given :math:`p(\\mathbf f)` and :math:`\\mathbf y`.\n\n        :param ~gpytorch.distributions.MultivariateNormal function_dist: :math:`p(\\mathbf f)`\n            the outputs of the latent function (the :obj:`gpytorch.models.ExactGP`)\n        :param torch.Tensor target: :math:`\\mathbf y` The target values\n        :rtype: torch.Tensor\n        :return: Exact MLL. Output shape corresponds to batch shape of the model/input data.\n        \"\"\"\n        if not isinstance(function_dist, MultivariateNormal):\n            raise RuntimeError(\"ExactMarginalLogLikelihood can only operate on Gaussian random variables\")\n\n        # Get the log prob of the marginal distribution\n        output = self.likelihood(function_dist, *params)\n        res = output.log_prob(target)\n        res = self._add_other_terms(res, params)\n\n        # Scale by the amount of data we have\n        num_data = target.size(-1)\n        return res.div_(num_data)\n\n    def pyro_factor(self, output, target, *params):\n        import pyro\n\n        mll = target.size(-1) * self(output, target, *params)\n        pyro.factor(\"gp_mll\", mll)\n        return mll\n", "path": "gpytorch/mlls/exact_marginal_log_likelihood.py"}], "after_files": [{"content": "#!/usr/bin/env python3\n\nfrom ..distributions import MultivariateNormal\nfrom ..likelihoods import _GaussianLikelihoodBase\nfrom .marginal_log_likelihood import MarginalLogLikelihood\n\n\nclass ExactMarginalLogLikelihood(MarginalLogLikelihood):\n    \"\"\"\n    The exact marginal log likelihood (MLL) for an exact Gaussian process with a\n    Gaussian likelihood.\n\n    .. note::\n        This module will not work with anything other than a :obj:`~gpytorch.likelihoods.GaussianLikelihood`\n        and a :obj:`~gpytorch.models.ExactGP`. It also cannot be used in conjunction with\n        stochastic optimization.\n\n    :param ~gpytorch.likelihoods.GaussianLikelihood likelihood: The Gaussian likelihood for the model\n    :param ~gpytorch.models.ExactGP model: The exact GP model\n\n    Example:\n        >>> # model is a gpytorch.models.ExactGP\n        >>> # likelihood is a gpytorch.likelihoods.Likelihood\n        >>> mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)\n        >>>\n        >>> output = model(train_x)\n        >>> loss = -mll(output, train_y)\n        >>> loss.backward()\n    \"\"\"\n\n    def __init__(self, likelihood, model):\n        if not isinstance(likelihood, _GaussianLikelihoodBase):\n            raise RuntimeError(\"Likelihood must be Gaussian for exact inference\")\n        super(ExactMarginalLogLikelihood, self).__init__(likelihood, model)\n\n    def _add_other_terms(self, res, params):\n        # Add additional terms (SGPR / learned inducing points, heteroskedastic likelihood models)\n        for added_loss_term in self.model.added_loss_terms():\n            res = res.add(added_loss_term.loss(*params))\n\n        # Add log probs of priors on the (functions of) parameters\n        for name, module, prior, closure, _ in self.named_priors():\n            res.add_(prior.log_prob(closure(module)).sum())\n\n        return res\n\n    def forward(self, function_dist, target, *params):\n        r\"\"\"\n        Computes the MLL given :math:`p(\\mathbf f)` and :math:`\\mathbf y`.\n\n        :param ~gpytorch.distributions.MultivariateNormal function_dist: :math:`p(\\mathbf f)`\n            the outputs of the latent function (the :obj:`gpytorch.models.ExactGP`)\n        :param torch.Tensor target: :math:`\\mathbf y` The target values\n        :rtype: torch.Tensor\n        :return: Exact MLL. Output shape corresponds to batch shape of the model/input data.\n        \"\"\"\n        if not isinstance(function_dist, MultivariateNormal):\n            raise RuntimeError(\"ExactMarginalLogLikelihood can only operate on Gaussian random variables\")\n\n        # Get the log prob of the marginal distribution\n        output = self.likelihood(function_dist, *params)\n        res = output.log_prob(target)\n        res = self._add_other_terms(res, params)\n\n        # Scale by the amount of data we have\n        num_data = function_dist.event_shape.numel()\n        return res.div_(num_data)\n\n    def pyro_factor(self, output, target, *params):\n        import pyro\n\n        mll = target.size(-1) * self(output, target, *params)\n        pyro.factor(\"gp_mll\", mll)\n        return mll\n", "path": "gpytorch/mlls/exact_marginal_log_likelihood.py"}]}