Repository: cornellius-gp/gpytorch
Issue: [Bug] size mismatch when using natural gradient in multi-output GP

# 🐛 Bug

size mismatch when using natural gradient in multi-output GP

## To reproduce

** Code snippet to reproduce **
```python
import math
import torch
import gpytorch
import tqdm
from matplotlib import pyplot as plt

train_x = torch.linspace(0, 1, 100)

train_y = torch.stack([
    torch.sin(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * 0.2,
    torch.cos(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * 0.2,
    torch.sin(train_x * (2 * math.pi)) + 2 * torch.cos(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * 0.2,
    -torch.cos(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * 0.2,
], -1)

print(train_x.shape, train_y.shape)

num_latents = 3
num_tasks = 4


class IndependentMultitaskGPModel(gpytorch.models.ApproximateGP):
    def __init__(self):
        # Let's use a different set of inducing points for each task
        inducing_points = torch.rand(num_tasks, 16, 1)

        # We have to mark the CholeskyVariationalDistribution as batch
        # so that we learn a variational distribution for each task
        # variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(
        #     inducing_points.size(-2), batch_shape=torch.Size([num_tasks])
        # )
        variational_distribution = gpytorch.variational.NaturalVariationalDistribution(
            inducing_points.size(-2), batch_shape=torch.Size([num_tasks])
        )

        variational_strategy = gpytorch.variational.IndependentMultitaskVariationalStrategy(
            gpytorch.variational.VariationalStrategy(
                self, inducing_points, variational_distribution, learn_inducing_locations=True
            ),
            num_tasks=4,
        )

        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.covar_module = gpytorch.kernels.ScaleKernel(
            gpytorch.kernels.RBFKernel(batch_shape=torch.Size([num_tasks])),
            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)


# model = MultitaskGPModel()

model = IndependentMultitaskGPModel()
likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=num_tasks)

# this is for running the notebook in our testing framework
import os
smoke_test = ('CI' in os.environ)
num_epochs = 1 if smoke_test else 500


model.train()
likelihood.train()

ngd_optimizer = gpytorch.optim.NGD(model.variational_parameters(),
                                   num_data=train_y.size(0),
                                   lr=0.1)
# We use SGD here, rather than Adam. Emperically, we find that SGD is better for variational regression
optimizer = torch.optim.Adam([
    {'params': model.hyperparameters()},
    {'params': likelihood.parameters()},
], lr=0.1)

# Our loss object. We're using the VariationalELBO, which essentially just computes the ELBO
mll = gpytorch.mlls.VariationalELBO(likelihood, model, num_data=train_y.size(0))

# We use more CG iterations here because the preconditioner introduced in the NeurIPS paper seems to be less
# effective for VI.
epochs_iter = tqdm.tqdm(range(num_epochs), desc="Epoch")
for i in epochs_iter:
    # Within each iteration, we will go over each minibatch of data
    ngd_optimizer.zero_grad()
    optimizer.zero_grad()
    output = model(train_x)
    loss = -mll(output, train_y)
    epochs_iter.set_postfix(loss=loss.item())
    loss.backward()
    ngd_optimizer.step()
    optimizer.step()

# Set into eval mode
model.eval()
likelihood.eval()

# Initialize plots
fig, axs = plt.subplots(1, num_tasks, figsize=(4 * num_tasks, 3))

# Make predictions
with torch.no_grad(), gpytorch.settings.fast_pred_var():
    test_x = torch.linspace(0, 1, 51)
    predictions = likelihood(model(test_x))
    mean = predictions.mean
    lower, upper = predictions.confidence_region()

for task, ax in enumerate(axs):
    # Plot training data as black stars
    ax.plot(train_x.detach().numpy(), train_y[:, task].detach().numpy(), 'k*')
    # Predictive mean as blue line
    ax.plot(test_x.numpy(), mean[:, task].numpy(), 'b')
    # Shade in confidence
    ax.fill_between(test_x.numpy(), lower[:, task].numpy(), upper[:, task].numpy(), alpha=0.5)
    ax.set_ylim([-3, 3])
    ax.legend(['Observed Data', 'Mean', 'Confidence'])
    ax.set_title(f'Task {task + 1}')

fig.tight_layout()
plt.show()
```

** Stack trace/error message **
```bash
Traceback (most recent call last):
torch.Size([100]) torch.Size([100, 4])
Epoch:   0%|                                            | 0/500 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "multi-svgp.py", line 99, in <module>
    output = model(train_x)
  File "/home/user/miniconda2/envs/torch_play/lib/python3.7/site-packages/gpytorch/models/approximate_gp.py", line 81, in __call__
    return self.variational_strategy(inputs, prior=prior)
  File "/home/user/miniconda2/envs/torch_play/lib/python3.7/site-packages/gpytorch/variational/independent_multitask_variational_strategy.py", line 47, in __call__
    function_dist = self.base_variational_strategy(x, prior=prior)
  File "/home/user/miniconda2/envs/torch_play/lib/python3.7/site-packages/gpytorch/variational/variational_strategy.py", line 166, in __call__
    return super().__call__(x, prior=prior)
  File "/home/user/miniconda2/envs/torch_play/lib/python3.7/site-packages/gpytorch/variational/_variational_strategy.py", line 114, in __call__
    self._variational_distribution.initialize_variational_distribution(prior_dist)
  File "/home/user/miniconda2/envs/torch_play/lib/python3.7/site-packages/gpytorch/variational/natural_variational_distribution.py", line 69, in initialize_variational_distribution
    self.natural_vec.data.copy_((prior_prec @ prior_mean).add_(noise))
RuntimeError: size mismatch, m1: [64 x 16], m2: [4 x 16] at /pytorch/aten/src/TH/generic/THTensorMath.cpp:41
```

## Expected Behavior

Getting similar results as in the [non-natural gradient notebook](https://docs.gpytorch.ai/en/v1.2.0/examples/04_Variational_and_Approximate_GPs/SVGP_Multitask_GP_Regression.html#Train-the-model)

## System information

**Please complete the following information:**
-  GPyTorch Version 1.2.0 
- PyTorch Version 1.6.0
-  Debian 10

## Additional context
I followed the [NGD tutorial](https://docs.gpytorch.ai/en/v1.2.0/examples/04_Variational_and_Approximate_GPs/Natural_Gradient_Descent.html) and replace the necessary code in the multi-output SVGP cases. Seems the problems come to the shape of `prior_prec` and `prior_mean`. 


Changed files:
test/variational/test_variational_strategy.py, gpytorch/variational/natural_variational_distribution.py

diff --git a/gpytorch/variational/natural_variational_distribution.py b/gpytorch/variational/natural_variational_distribution.py
index 6f6ff5f4..4927a5f3 100644
--- a/gpytorch/variational/natural_variational_distribution.py
+++ b/gpytorch/variational/natural_variational_distribution.py
@@ -6,6 +6,7 @@ import torch
 
 from ..distributions import MultivariateNormal
 from ..lazy import CholLazyTensor, TriangularLazyTensor
+from ..utils.cholesky import psd_safe_cholesky
 from ._variational_distribution import _VariationalDistribution
 
 
@@ -66,7 +67,7 @@ class NaturalVariationalDistribution(_NaturalVariationalDistribution):
         prior_mean = prior_dist.mean
         noise = torch.randn_like(prior_mean).mul_(self.mean_init_std)
 
-        self.natural_vec.data.copy_((prior_prec @ prior_mean).add_(noise))
+        self.natural_vec.data.copy_((prior_prec @ prior_mean.unsqueeze(-1)).squeeze(-1).add_(noise))
         self.natural_mat.data.copy_(prior_prec.mul(-0.5))
 
 
@@ -95,7 +96,7 @@ class _NaturalToMuVarSqrt(torch.autograd.Function):
     @staticmethod
     def _forward(nat_mean, nat_covar):
         try:
-            L_inv = torch.cholesky(-2.0 * nat_covar, upper=False)
+            L_inv = psd_safe_cholesky(-2.0 * nat_covar, upper=False)
         except RuntimeError as e:
             if str(e).startswith("cholesky"):
                 raise RuntimeError(
@@ -110,7 +111,7 @@ class _NaturalToMuVarSqrt(torch.autograd.Function):
         mu = (S @ nat_mean.unsqueeze(-1)).squeeze(-1)
         # Two choleskys are annoying, but we don't have good support for a
         # LazyTensor of form L.T @ L
-        return mu, torch.cholesky(S, upper=False)
+        return mu, psd_safe_cholesky(S, upper=False)
 
     @staticmethod
     def forward(ctx, nat_mean, nat_covar):
diff --git a/test/variational/test_variational_strategy.py b/test/variational/test_variational_strategy.py
index 2d28e69e..4b879af4 100644
--- a/test/variational/test_variational_strategy.py
+++ b/test/variational/test_variational_strategy.py
@@ -86,5 +86,23 @@ class TestDeltaRobustVGP(TestRobustVGP):
         return gpytorch.variational.DeltaVariationalDistribution
 
 
+class TestNGDVariationalGP(TestVariationalGP):
+    @property
+    def distribution_cls(self):
+        return gpytorch.variational.NaturalVariationalDistribution
+
+    def test_training_iteration(self, *args, **kwargs):
+        cg_mock, cholesky_mock, ciq_mock = VariationalTestCase.test_training_iteration(self, *args, **kwargs)
+        self.assertFalse(cg_mock.called)
+        self.assertEqual(cholesky_mock.call_count, 6)  # Three for each forward pass
+        self.assertFalse(ciq_mock.called)
+
+    def test_eval_iteration(self, *args, **kwargs):
+        cg_mock, cholesky_mock, ciq_mock = VariationalTestCase.test_eval_iteration(self, *args, **kwargs)
+        self.assertFalse(cg_mock.called)
+        self.assertEqual(cholesky_mock.call_count, 3)  # One to compute cache + 2 to compute variational distribution
+        self.assertFalse(ciq_mock.called)
+
+
 if __name__ == "__main__":
     unittest.main()