Repository: cornellius-gp/gpytorch

Issue:
Natural Gradient Optimizer
Hello, 

I have seen in the [literature](https://arxiv.org/abs/1803.09151) for GP algorithms that use inducing points and variational inference to perform inference can sometimes converge slowly when only using optimizers such as stochastic gradient descent. This is apparently due to the difficulty of optimizing the variational parameters as well as the likelihood and kernel parameters. The authors in that paper suggest that using the **natural gradient optimizer** has shown to have big improvements especially in large-scale GP methods like the **SVGP** and **DeepGP**.

I've seen the NatGrad optimizer has already been implemented in this [branch](https://github.com/cornellius-gp/gpytorch/tree/natural) within the GPyTorch library. I wanted to know what was everyone's experience using it and why is it not in the main branch? Perhaps the nature of using matrix-vector-multiplication methods don't suffer from the same issue of convergence or joint optimization? Or it's just a matter of code coverage?



Discussion:

Patch:
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 56397ac1..7fd396db 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -35,6 +35,7 @@ GPyTorch's documentation
    distributions
    priors
    variational
+   optim
 
 .. toctree::
    :maxdepth: 1
diff --git a/docs/source/optim.rst b/docs/source/optim.rst
new file mode 100644
index 00000000..e4791ed3
--- /dev/null
+++ b/docs/source/optim.rst
@@ -0,0 +1,15 @@
+.. role:: hidden
+    :class: hidden-section
+
+gpytorch.optim
+===================================
+
+.. automodule:: gpytorch.optim
+.. currentmodule:: gpytorch.optim
+
+
+LazyTensor
+----------------
+
+.. autoclass:: NGD
+   :members:
diff --git a/docs/source/variational.rst b/docs/source/variational.rst
index 37cea8f5..42b4e326 100644
--- a/docs/source/variational.rst
+++ b/docs/source/variational.rst
@@ -139,3 +139,53 @@ distributions.
 
 .. autoclass:: MeanFieldVariationalDistribution
    :members:
+
+
+Variational Distributions for Natural Gradient Descent
+---------------------------------------------------------
+
+Special :obj:`~_VariationalDistribution` objects designed specifically for
+use with natural gradient descent techniques.
+
+If the variational distribution is defined by :math:`\mathcal{N}(\mathbf m, \mathbf S)`, then
+a :obj:`~gpytorch.variational.NaturalVariationalDistribution` uses the parameterization:
+
+.. math::
+
+    \begin{align*}
+        \boldsymbol \theta_\text{vec} &= \mathbf S^{-1} \mathbf m
+        \\
+        \boldsymbol \Theta_\text{mat} &= -\frac{1}{2} \mathbf S^{-1}.
+    \end{align*}
+
+The gradients with respect to the variational parameters calculated by this
+class are instead the **natural** gradients. Thus, optimising its
+parameters using gradient descent (:obj:`~torch.optim.SGDOptimizer`)
+becomes natural gradient descent (see e.g. `Salimbeni et al., 2018`_).
+
+GPyTorch offers several :obj:`_NaturalVariationalDistribution` classes, each of which uses a
+different representation of the natural parameters. The different parameterizations trade off
+speed and stability.
+
+.. note::
+    Natural gradient descent is very stable with variational regression,
+    and fast: if the hyperparameters are fixed, the variational parameters
+    converge in 1 iteration. However, it can be unstable with non-conjugate
+    likelihoods and alternative objective functions.
+
+.. _Salimbeni et al., 2018:
+    https://arxiv.org/abs/1803.09151
+
+
+:hidden:`NaturalVariationalDistribution`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: NaturalVariationalDistribution
+   :members:
+
+
+:hidden:`TrilNaturalVariationalDistribution`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: TrilNaturalVariationalDistribution
+   :members:
diff --git a/examples/04_Variational_and_Approximate_GPs/Natural_Gradient_Descent.ipynb b/examples/04_Variational_and_Approximate_GPs/Natural_Gradient_Descent.ipynb
new file mode 100644
index 00000000..da1f3738
--- /dev/null
+++ b/examples/04_Variational_and_Approximate_GPs/Natural_Gradient_Descent.ipynb
@@ -0,0 +1,389 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Using Natural Gradient Descent with Variational Models\n",
+    "\n",
+    "## Overview\n",
+    "\n",
+    "In this notebook, we'll demonstrate how to use **natural gradient descent** when optimizing variational GPyTorch models. This will be the same as the [SVGP regression notebook](./SVGP_Regression_CUDA.ipynb), except we will be using a different optimizer.\n",
+    "\n",
+    "## What is natural gradient descent (NGD)?\n",
+    "\n",
+    "Without going into too much detail, using SGD or Adam isn't the best way to optimize the parameters of variational Gaussian distributions. Essentially, SGD takes steps assuming that the loss geometry of the parameters is Euclidean. This is a bad assumption for the parameters of many distributions, especially Gaussians. See [Agustinus Kristiadi's blog post](https://wiseodd.github.io/techblog/2018/03/14/natural-gradient/) for more detail.\n",
+    "\n",
+    "Instead, it makes more sense to take gradient steps that have been scaled better for the geometry of the variational distribution's parameters. Specifically, if $\\mathbf m$ and $\\mathbf S$ are the mean and covariance of our variational Gaussian posterior approximation, then we will achieve faster convergence if we take the following steps:\n",
+    "\n",
+    "\\begin{align}\n",
+    "\\begin{bmatrix} \\mathbf m \\\\ \\mathbf S \\end{bmatrix}\n",
+    "\\leftarrow\n",
+    "\\begin{bmatrix} \\mathbf m \\\\ \\mathbf S \\end{bmatrix}\n",
+    "-\n",
+    "\\alpha \\mathcal{\\mathbf F}^{-1} \\nabla \\begin{bmatrix} \\mathbf m \\\\   \\mathbf S \\end{bmatrix}\n",
+    "\\end{align}\n",
+    "\n",
+    "where $\\alpha$ is a step size and $\\mathbf F$ is the **Fisher information matrix** corresponding to this distribution. This is known as **natural gradient descent**, or **NGD**.\n",
+    "\n",
+    "It turns out that for Gaussian distributions (and, more broadly, for all distributions in the exponential family), there are efficient update equations for NGD. See the following papers for more information:\n",
+    "- [Salimbeni, Hugh, Stefanos Eleftheriadis, and James Hensman. \"Natural gradients in practice: Non-conjugate variational inference in gaussian process models.\" AISTATS (2018).](https://arxiv.org/abs/1803.09151)\n",
+    "- [Hensman, James, Magnus Rattray, and Neil D. Lawrence. \"Fast variational inference in the conjugate exponential family.\" NeurIPS (2012).](http://papers.nips.cc/paper/4766-fast-variational-inference-in-the-conjugate-exponential-family)\n",
+    "\n",
+    "## Jointly optimizing variational parameters/hyperparameters"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import tqdm\n",
+    "import math\n",
+    "import torch\n",
+    "import gpytorch\n",
+    "from matplotlib import pyplot as plt\n",
+    "\n",
+    "# Make plots inline\n",
+    "%matplotlib inline"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "For this example notebook, we'll be using the `elevators` UCI dataset used in the paper. Running the next cell downloads a copy of the dataset that has already been scaled and normalized appropriately. For this notebook, we'll simply be splitting the data using the first 80% of the data as training and the last 20% as testing."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import urllib.request\n",
+    "import os\n",
+    "from scipy.io import loadmat\n",
+    "from math import floor\n",
+    "\n",
+    "\n",
+    "# this is for running the notebook in our testing framework\n",
+    "smoke_test = ('CI' in os.environ)\n",
+    "\n",
+    "\n",
+    "if not smoke_test and not os.path.isfile('../elevators.mat'):\n",
+    "    print('Downloading \\'elevators\\' UCI dataset...')\n",
+    "    urllib.request.urlretrieve('https://drive.google.com/uc?export=download&id=1jhWL3YUHvXIaftia4qeAyDwVxo6j1alk', '../elevators.mat')\n",
+    "\n",
+    "\n",
+    "if smoke_test:  # this is for running the notebook in our testing framework\n",
+    "    X, y = torch.randn(1000, 3), torch.randn(1000)\n",
+    "else:\n",
+    "    data = torch.Tensor(loadmat('../elevators.mat')['data'])\n",
+    "    X = data[:, :-1]\n",
+    "    X = X - X.min(0)[0]\n",
+    "    X = 2 * (X / X.max(0)[0]) - 1\n",
+    "    y = data[:, -1]\n",
+    "\n",
+    "\n",
+    "train_n = int(floor(0.8 * len(X)))\n",
+    "train_x = X[:train_n, :].contiguous()\n",
+    "train_y = y[:train_n].contiguous()\n",
+    "\n",
+    "test_x = X[train_n:, :].contiguous()\n",
+    "test_y = y[train_n:].contiguous()\n",
+    "\n",
+    "if torch.cuda.is_available():\n",
+    "    train_x, train_y, test_x, test_y = train_x.cuda(), train_y.cuda(), test_x.cuda(), test_y.cuda()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The following steps create the dataloader objects. See the [SVGP regression notebook](./SVGP_Regression_CUDA.ipynb) for details."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from torch.utils.data import TensorDataset, DataLoader\n",
+    "train_dataset = TensorDataset(train_x, train_y)\n",
+    "train_loader = DataLoader(train_dataset, batch_size=1024, shuffle=True)\n",
+    "\n",
+    "test_dataset = TensorDataset(test_x, test_y)\n",
+    "test_loader = DataLoader(test_dataset, batch_size=1024, shuffle=False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## SVGP models for NGD\n",
+    "\n",
+    "There are **three** key differences between NGD-SVGP models and the standard SVGP models from the [SVGP regression notebook](./SVGP_Regression_CUDA.ipynb).\n",
+    "\n",
+    "\n",
+    "### Difference #1: NaturalVariationalDistribution\n",
+    "\n",
+    "Rather than using `gpytorch.variational.CholeskyVarationalDistribution` (or other variational distribution objects), you have to use one of the two objects:\n",
+    "\n",
+    "- `gpytorch.variational.NaturalVariationalDistribution` (typically faster optimization convergence, but less stable for non-conjugate likelihoods)\n",
+    "- `gpytorch.variational.TrilNaturalVariationalDistribution` (typically slower optimization convergence, but more stable for non-conjugate likelihoods)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class GPModel(gpytorch.models.ApproximateGP):\n",
+    "    def __init__(self, inducing_points):\n",
+    "        variational_distribution = gpytorch.variational.NaturalVariationalDistribution(inducing_points.size(0))\n",
+    "        variational_strategy = gpytorch.variational.VariationalStrategy(\n",
+    "            self, inducing_points, variational_distribution, learn_inducing_locations=True\n",
+    "        )\n",
+    "        super(GPModel, self).__init__(variational_strategy)\n",
+    "        self.mean_module = gpytorch.means.ConstantMean()\n",
+    "        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())\n",
+    "        \n",
+    "    def forward(self, x):\n",
+    "        mean_x = self.mean_module(x)\n",
+    "        covar_x = self.covar_module(x)\n",
+    "        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)\n",
+    "\n",
+    "\n",
+    "inducing_points = train_x[:500, :]\n",
+    "model = GPModel(inducing_points=inducing_points)\n",
+    "likelihood = gpytorch.likelihoods.GaussianLikelihood()\n",
+    "\n",
+    "if torch.cuda.is_available():\n",
+    "    model = model.cuda()\n",
+    "    likelihood = likelihood.cuda()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Difference #2: Two optimizers - one for the variational parameters; one for the hyperparameters\n",
+    "\n",
+    "NGD steps only update the variational parameters. Therefore, we need two separate optimizers: one for the variational parameters (using NGD) and one for the other hyperparameters (using Adam or whatever you want).\n",
+    "\n",
+    "Some things to note about the NGD variational optimizer:\n",
+    "\n",
+    "- **You must use `gpytorch.optim.NGD` as the variational NGD optimizer!** Adaptive gradient algorithms will mess up the natural gradient steps. (Any stochastic optimizer works for the hyperparameters.)\n",
+    "- **Use a large learning rate for the variational optimizer.** Typically, 0.1 is a good learning rate."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "variational_ngd_optimizer = gpytorch.optim.NGD(model.variational_parameters(), num_data=train_y.size(0), lr=0.1)\n",
+    "\n",
+    "hyperparameter_optimizer = torch.optim.Adam([\n",
+    "    {'params': model.hyperparameters()},\n",
+    "    {'params': likelihood.parameters()},\n",
+    "], lr=0.01)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Difference #3: The updated training loop\n",
+    "\n",
+    "In the training loop, we have to update both optimizers (`variational_ngd_optimizer` and `hyperparameter_optimizer`). "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "126c78ce830a4352aba355851d55ad5a",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "HBox(children=(FloatProgress(value=0.0, description='Epoch', max=4.0, style=ProgressStyle(description_width='i…"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "HBox(children=(FloatProgress(value=0.0, description='Minibatch', max=13.0, style=ProgressStyle(description_wid…"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "HBox(children=(FloatProgress(value=0.0, description='Minibatch', max=13.0, style=ProgressStyle(description_wid…"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "HBox(children=(FloatProgress(value=0.0, description='Minibatch', max=13.0, style=ProgressStyle(description_wid…"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "HBox(children=(FloatProgress(value=0.0, description='Minibatch', max=13.0, style=ProgressStyle(description_wid…"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "model.train()\n",
+    "likelihood.train()\n",
+    "mll = gpytorch.mlls.VariationalELBO(likelihood, model, num_data=train_y.size(0))\n",
+    "\n",
+    "num_epochs = 1 if smoke_test else 4\n",
+    "epochs_iter = tqdm.notebook.tqdm(range(num_epochs), desc=\"Epoch\")\n",
+    "for i in epochs_iter:\n",
+    "    minibatch_iter = tqdm.notebook.tqdm(train_loader, desc=\"Minibatch\", leave=False)\n",
+    "    \n",
+    "    for x_batch, y_batch in minibatch_iter:\n",
+    "        ### Perform NGD step to optimize variational parameters\n",
+    "        variational_ngd_optimizer.zero_grad()\n",
+    "        hyperparameter_optimizer.zero_grad()\n",
+    "        output = model(x_batch)\n",
+    "        loss = -mll(output, y_batch)\n",
+    "        minibatch_iter.set_postfix(loss=loss.item())\n",
+    "        loss.backward()\n",
+    "        variational_ngd_optimizer.step()\n",
+    "        hyperparameter_optimizer.step()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "You could also modify the optimization loop to alternate between the NGD step and the hyperparameter updates:\n",
+    "\n",
+    "```python\n",
+    "    for x_batch, y_batch in minibatch_iter:\n",
+    "        ### Perform NGD step to optimize variational parameters\n",
+    "        variational_ngd_optimizer.zero_grad()\n",
+    "        output = model(x_batch)\n",
+    "        loss = -mll(output, y_batch)\n",
+    "        minibatch_iter.set_postfix(loss=loss.item())\n",
+    "        loss.backward()\n",
+    "        variational_ngd_optimizer.step()\n",
+    "        \n",
+    "        ### Perform Adam step to optimize hyperparameters\n",
+    "        hyperparameter_optimizer.zero_grad()\n",
+    "        output = model(x_batch)\n",
+    "        loss = -mll(output, y_batch)\n",
+    "        loss.backward()\n",
+    "        hyperparameter_optimizer.step()\n",
+    "```\n",
+    "\n",
+    "### Evaluation\n",
+    "\n",
+    "That's it! This model should converge must faster/better than the standard SVGP model - and it can often get better performance (especially when using more inducing points)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Test MAE: 0.07514254748821259\n"
+     ]
+    }
+   ],
+   "source": [
+    "model.eval()\n",
+    "likelihood.eval()\n",
+    "means = torch.tensor([0.])\n",
+    "with torch.no_grad():\n",
+    "    for x_batch, y_batch in test_loader:\n",
+    "        preds = model(x_batch)\n",
+    "        means = torch.cat([means, preds.mean.cpu()])\n",
+    "means = means[1:]\n",
+    "print('Test MAE: {}'.format(torch.mean(torch.abs(means - test_y.cpu()))))"
+   ]
+  }
+ ],
+ "metadata": {
+  "anaconda-cloud": {},
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/examples/04_Variational_and_Approximate_GPs/index.rst b/examples/04_Variational_and_Approximate_GPs/index.rst
index 29531757..daea7518 100644
--- a/examples/04_Variational_and_Approximate_GPs/index.rst
+++ b/examples/04_Variational_and_Approximate_GPs/index.rst
@@ -23,6 +23,7 @@ The variational documentation has more information on how to use these objects.
 Here we provide some examples which highlight some of the common use cases:
 
 - **Large-scale regression** (when exact methods are too memory intensive): see the `stochastic variational regression example`_.
+- **Variational inference with natural gradient descent** (for faster/better optimization): see the `ngd example`_.
 - **Variational distribution options** for different scalability/expressiveness: see the `strategy/distribution comparison`_.
 - **Alternative optimization objectives** for the GP's predictive distribution: see the `approximate GP objective functions notebook`_.
   This example compares and contrasts the variational ELBO with the predictive log likelihood of Jankowiak et al., 2020.
@@ -36,6 +37,7 @@ Here we provide some examples which highlight some of the common use cases:
 
    SVGP_Regression_CUDA.ipynb
    Modifying_the_variational_strategy_and_distribution.ipynb
+   Natural_Gradient_Descent.ipynb
    Approximate_GP_Objective_Functions.ipynb
    Non_Gaussian_Likelihoods.ipynb
    SVGP_Multitask_GP_Regression.ipynb
@@ -47,6 +49,9 @@ Here we provide some examples which highlight some of the common use cases:
 .. _stochastic variational regression example:
   ./SVGP_Regression_CUDA.ipynb
 
+.. _ngd example:
+  ./Natural_Gradient_Descent.ipynb
+
 .. _approximate GP objective functions notebook:
   ./Approximate_GP_Objective_Functions.ipynb
 
diff --git a/gpytorch/__init__.py b/gpytorch/__init__.py
index 61d7b15a..827e39fb 100644
--- a/gpytorch/__init__.py
+++ b/gpytorch/__init__.py
@@ -8,6 +8,7 @@ from . import (
     means,
     mlls,
     models,
+    optim,
     priors,
     settings,
     utils,
@@ -41,6 +42,7 @@ __all__ = [
     "means",
     "mlls",
     "models",
+    "optim",
     "priors",
     "utils",
     "variational",
diff --git a/gpytorch/lazy/chol_lazy_tensor.py b/gpytorch/lazy/chol_lazy_tensor.py
index 9d19c320..f31229c6 100644
--- a/gpytorch/lazy/chol_lazy_tensor.py
+++ b/gpytorch/lazy/chol_lazy_tensor.py
@@ -45,7 +45,7 @@ class CholLazyTensor(RootLazyTensor):
     @cached
     def inverse(self):
         Linv = self.root.inverse()  # this could be slow in some cases w/ structured lazies
-        return CholLazyTensor(Linv, upper=not self.upper)
+        return CholLazyTensor(TriangularLazyTensor(Linv, upper=not self.upper), upper=not self.upper)
 
     def inv_matmul(self, right_tensor, left_tensor=None):
         is_vector = right_tensor.ndim == 1
diff --git a/gpytorch/lazy/lazy_tensor.py b/gpytorch/lazy/lazy_tensor.py
index af416b38..bc4e5e2d 100644
--- a/gpytorch/lazy/lazy_tensor.py
+++ b/gpytorch/lazy/lazy_tensor.py
@@ -1017,8 +1017,9 @@ class LazyTensor(ABC):
         # Special case: use Cholesky to compute these terms
         if settings.fast_computations.log_prob.off() or (self.size(-1) <= settings.max_cholesky_size.value()):
             from .chol_lazy_tensor import CholLazyTensor
+            from .triangular_lazy_tensor import TriangularLazyTensor
 
-            cholesky = CholLazyTensor(self.cholesky())
+            cholesky = CholLazyTensor(TriangularLazyTensor(self.cholesky()))
             return cholesky.inv_quad_logdet(inv_quad_rhs=inv_quad_rhs, logdet=logdet, reduce_inv_quad=reduce_inv_quad)
 
         # Default: use modified batch conjugate gradients to compute these terms
diff --git a/gpytorch/models/pyro/_pyro_mixin.py b/gpytorch/models/pyro/_pyro_mixin.py
index d38d57a9..a33c4af5 100644
--- a/gpytorch/models/pyro/_pyro_mixin.py
+++ b/gpytorch/models/pyro/_pyro_mixin.py
@@ -1,8 +1,9 @@
 #!/usr/bin/env python3
 
-import pyro
 import torch
 
+import pyro
+
 
 class _PyroMixin(object):
     def pyro_guide(self, input, beta=1.0, name_prefix=""):
diff --git a/gpytorch/module.py b/gpytorch/module.py
index ef052096..96ba783e 100644
--- a/gpytorch/module.py
+++ b/gpytorch/module.py
@@ -138,9 +138,12 @@ class Module(nn.Module):
         return _extract_named_added_loss_terms(module=self, memo=None, prefix="")
 
     def named_hyperparameters(self):
-        for name, param in self.named_parameters():
-            if "variational_" not in name:
-                yield name, param
+        from .variational._variational_distribution import _VariationalDistribution
+
+        for module_prefix, module in self.named_modules():
+            if not isinstance(module, _VariationalDistribution):
+                for elem in module.named_parameters(prefix=module_prefix, recurse=False):
+                    yield elem
 
     def named_priors(self, memo=None, prefix=""):
         """Returns an iterator over the module's priors, yielding the name of the prior,
@@ -159,9 +162,12 @@ class Module(nn.Module):
         return _extract_named_constraints(module=self, memo=None, prefix="")
 
     def named_variational_parameters(self):
-        for name, param in self.named_parameters():
-            if "variational_" in name:
-                yield name, param
+        from .variational._variational_distribution import _VariationalDistribution
+
+        for module_prefix, module in self.named_modules():
+            if isinstance(module, _VariationalDistribution):
+                for elem in module.named_parameters(prefix=module_prefix, recurse=False):
+                    yield elem
 
     def register_added_loss_term(self, name):
         self._added_loss_terms[name] = None
diff --git a/gpytorch/optim/__init__.py b/gpytorch/optim/__init__.py
new file mode 100644
index 00000000..3b3f0ab6
--- /dev/null
+++ b/gpytorch/optim/__init__.py
@@ -0,0 +1,5 @@
+#!/usr/bin/env python3
+
+from .ngd import NGD
+
+__all__ = ["NGD"]
diff --git a/gpytorch/optim/ngd.py b/gpytorch/optim/ngd.py
new file mode 100644
index 00000000..0847243c
--- /dev/null
+++ b/gpytorch/optim/ngd.py
@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+
+from typing import Iterable, Union
+
+import torch
+
+
+class NGD(torch.optim.Optimizer):
+    r"""Implements a natural gradient descent step.
+    It **can only** be used in conjunction with a :obj:`~gpytorch.variational._NaturalVariationalDistribution`.
+
+    Nesterov momentum is based on the formula from
+    `On the importance of initialization and momentum in deep learning`__.
+
+    .. seealso::
+        - :obj:`gpytorch.variational.NaturalVariationalDistribution`
+        - :obj:`gpytorch.variational.TrilNaturalVariationalDistribution`
+        - The `natural gradient descent tutorial
+          <examples/04_Variational_and_Approximate_GPs/Natural_Gradient_Descent.ipynb>`_
+          for use instructions.
+
+    Example:
+        >>> ngd_optimizer = torch.optim.NGD(model.variational_parameters(), lr=0.1, momentum=0.9)
+        >>> ngd_optimizer.zero_grad()
+        >>> mll(gp_model(input), target).backward()
+        >>> ngd_optimizer.step()
+    """
+
+    def __init__(self, params: Iterable[Union[torch.nn.Parameter, dict]], num_data: int, lr: float = 0.1):
+        self.num_data = num_data
+        super().__init__(params, defaults=dict(lr=lr))
+
+    @torch.no_grad()
+    def step(self) -> torch.Tensor:
+        """Performs a single optimization step.
+        """
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                p.add_(p.grad, alpha=(-group["lr"] * self.num_data))
+
+        return None
diff --git a/gpytorch/variational/__init__.py b/gpytorch/variational/__init__.py
index 069b051a..8e54411b 100644
--- a/gpytorch/variational/__init__.py
+++ b/gpytorch/variational/__init__.py
@@ -13,7 +13,9 @@ from .independent_multitask_variational_strategy import (
 )
 from .lmc_variational_strategy import LMCVariationalStrategy
 from .mean_field_variational_distribution import MeanFieldVariationalDistribution
+from .natural_variational_distribution import NaturalVariationalDistribution, _NaturalVariationalDistribution
 from .orthogonally_decoupled_variational_strategy import OrthogonallyDecoupledVariationalStrategy
+from .tril_natural_variational_distribution import TrilNaturalVariationalDistribution
 from .unwhitened_variational_strategy import UnwhitenedVariationalStrategy
 from .variational_strategy import VariationalStrategy
 from .whitened_variational_strategy import WhitenedVariationalStrategy
@@ -34,4 +36,7 @@ __all__ = [
     "CholeskyVariationalDistribution",
     "MeanFieldVariationalDistribution",
     "DeltaVariationalDistribution",
+    "_NaturalVariationalDistribution",
+    "NaturalVariationalDistribution",
+    "TrilNaturalVariationalDistribution",
 ]
diff --git a/gpytorch/variational/_variational_distribution.py b/gpytorch/variational/_variational_distribution.py
index d63d23bb..f8081187 100644
--- a/gpytorch/variational/_variational_distribution.py
+++ b/gpytorch/variational/_variational_distribution.py
@@ -24,7 +24,7 @@ class _VariationalDistribution(Module, ABC):
         Constructs and returns the variational distribution
 
         :rtype: :obj:`~gpytorch.distributions.MultivariateNormal`
-        :return: The distribution :math:q(\mathbf u)"
+        :return: The distribution :math:`q(\mathbf u)`
         """
         raise NotImplementedError
 
@@ -33,7 +33,7 @@ class _VariationalDistribution(Module, ABC):
         r"""
         Method for initializing the variational distribution, based on the prior distribution.
 
-        :param ~gpytorch.distribution.Distribution prior_dist: The prior distribution :math:`p(\mathbf u)`.
+        :param ~gpytorch.distributions.Distribution prior_dist: The prior distribution :math:`p(\mathbf u)`.
         """
         raise NotImplementedError
 
diff --git a/gpytorch/variational/cholesky_variational_distribution.py b/gpytorch/variational/cholesky_variational_distribution.py
index 53bc016c..da83041a 100644
--- a/gpytorch/variational/cholesky_variational_distribution.py
+++ b/gpytorch/variational/cholesky_variational_distribution.py
@@ -18,9 +18,10 @@ class CholeskyVariationalDistribution(_VariationalDistribution):
 
     :param int num_inducing_points: Size of the variational distribution. This implies that the variational mean
         should be this size, and the variational covariance matrix should have this many rows and columns.
-    :param torch.Size batch_shape: (Optional.) Specifies an optional batch size
+    :param batch_shape: Specifies an optional batch size
         for the variational parameters. This is useful for example when doing additive variational inference.
-    :param float mean_init_std: (default=1e-3) Standard deviation of gaussian noise to add to the mean initialization.
+    :type batch_shape: :obj:`torch.Size`, optional
+    :param float mean_init_std: (Default: 1e-3) Standard deviation of gaussian noise to add to the mean initialization.
     """
 
     def __init__(self, num_inducing_points, batch_shape=torch.Size([]), mean_init_std=1e-3, **kwargs):
diff --git a/gpytorch/variational/delta_variational_distribution.py b/gpytorch/variational/delta_variational_distribution.py
index c3bbaf92..a5dd8f5e 100644
--- a/gpytorch/variational/delta_variational_distribution.py
+++ b/gpytorch/variational/delta_variational_distribution.py
@@ -12,10 +12,11 @@ class DeltaVariationalDistribution(_VariationalDistribution):
     with a single particle. It is equivalent to doing MAP inference.
 
     :param int num_inducing_points: Size of the variational distribution. This implies that the variational mean
-        should be this size.
-    :param torch.Size batch_shape: (Optional.) Specifies an optional batch size
+        should be this size, and the variational covariance matrix should have this many rows and columns.
+    :param batch_shape: Specifies an optional batch size
         for the variational parameters. This is useful for example when doing additive variational inference.
-    :param float mean_init_std: (default=1e-3) Standard deviation of gaussian noise to add to the mean initialization.
+    :type batch_shape: :obj:`torch.Size`, optional
+    :param float mean_init_std: (Default: 1e-3) Standard deviation of gaussian noise to add to the mean initialization.
     """
 
     def __init__(self, num_inducing_points, batch_shape=torch.Size([]), mean_init_std=1e-3, **kwargs):
diff --git a/gpytorch/variational/mean_field_variational_distribution.py b/gpytorch/variational/mean_field_variational_distribution.py
index 2f0447de..be62de70 100644
--- a/gpytorch/variational/mean_field_variational_distribution.py
+++ b/gpytorch/variational/mean_field_variational_distribution.py
@@ -15,9 +15,10 @@ class MeanFieldVariationalDistribution(_VariationalDistribution):
 
     :param int num_inducing_points: Size of the variational distribution. This implies that the variational mean
         should be this size, and the variational covariance matrix should have this many rows and columns.
-    :param torch.Size batch_shape: (Optional.) Specifies an optional batch size
+    :param batch_shape: Specifies an optional batch size
         for the variational parameters. This is useful for example when doing additive variational inference.
-    :param float mean_init_std: (default=1e-3) Standard deviation of gaussian noise to add to the mean initialization.
+    :type batch_shape: :obj:`torch.Size`, optional
+    :param float mean_init_std: (Default: 1e-3) Standard deviation of gaussian noise to add to the mean initialization.
     """
 
     def __init__(self, num_inducing_points, batch_shape=torch.Size([]), mean_init_std=1e-3, **kwargs):
diff --git a/gpytorch/variational/natural_variational_distribution.py b/gpytorch/variational/natural_variational_distribution.py
new file mode 100644
index 00000000..6f6ff5f4
--- /dev/null
+++ b/gpytorch/variational/natural_variational_distribution.py
@@ -0,0 +1,145 @@
+#!/usr/bin/env python3
+
+import abc
+
+import torch
+
+from ..distributions import MultivariateNormal
+from ..lazy import CholLazyTensor, TriangularLazyTensor
+from ._variational_distribution import _VariationalDistribution
+
+
+class _NaturalVariationalDistribution(_VariationalDistribution, abc.ABC):
+    r"""Any :obj:`~gpytorch.variational._VariationalDistribution` which calculates
+    natural gradients with respect to its parameters.
+    """
+    pass
+
+
+class NaturalVariationalDistribution(_NaturalVariationalDistribution):
+    r"""A multivariate normal :obj:`~gpytorch.variational._VariationalDistribution`,
+    parameterized by **natural** parameters.
+
+    .. note::
+       The :obj:`~gpytorch.variational.NaturalVariationalDistribution` can only
+       be used with :obj:`gpytorch.optim.NGD`, or other optimizers
+       that follow exactly the gradient direction. Failure to do so will cause
+       the natural matrix :math:`\mathbf \Theta_\text{mat}` to stop being
+       positive definite, and a :obj:`~RuntimeError` will be raised.
+
+    .. seealso::
+        The `natural gradient descent tutorial
+        <examples/04_Variational_and_Approximate_GPs/Natural_Gradient_Descent.ipynb>`_
+        for use instructions.
+
+        The :obj:`~gpytorch.variational.TrilNaturalVariationalDistribution` for
+        a more numerically stable parameterization, at the cost of needing more
+        iterations to make variational regression converge.
+
+    :param int num_inducing_points: Size of the variational distribution. This implies that the variational mean
+        should be this size, and the variational covariance matrix should have this many rows and columns.
+    :param batch_shape: Specifies an optional batch size
+        for the variational parameters. This is useful for example when doing additive variational inference.
+    :type batch_shape: :obj:`torch.Size`, optional
+    :param float mean_init_std: (Default: 1e-3) Standard deviation of gaussian noise to add to the mean initialization.
+
+    """
+
+    def __init__(self, num_inducing_points, batch_shape=torch.Size([]), mean_init_std=1e-3, **kwargs):
+        super().__init__(num_inducing_points=num_inducing_points, batch_shape=batch_shape, mean_init_std=mean_init_std)
+        scaled_mean_init = torch.zeros(num_inducing_points)
+        neg_prec_init = torch.eye(num_inducing_points, num_inducing_points).mul(-0.5)
+        scaled_mean_init = scaled_mean_init.repeat(*batch_shape, 1)
+        neg_prec_init = neg_prec_init.repeat(*batch_shape, 1, 1)
+
+        # eta1 and eta2 parameterization of the variational distribution
+        self.register_parameter(name="natural_vec", parameter=torch.nn.Parameter(scaled_mean_init))
+        self.register_parameter(name="natural_mat", parameter=torch.nn.Parameter(neg_prec_init))
+
+    def forward(self):
+        mean, chol_covar = _NaturalToMuVarSqrt.apply(self.natural_vec, self.natural_mat)
+        res = MultivariateNormal(mean, CholLazyTensor(TriangularLazyTensor(chol_covar)))
+        return res
+
+    def initialize_variational_distribution(self, prior_dist):
+        prior_prec = prior_dist.covariance_matrix.inverse()
+        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_mat.data.copy_(prior_prec.mul(-0.5))
+
+
+def _triangular_inverse(A, upper=False):
+    eye = torch.eye(A.size(-1), dtype=A.dtype, device=A.device)
+    return eye.triangular_solve(A, upper=upper).solution
+
+
+def _phi_for_cholesky_(A):
+    "Modifies A to be the phi function used in differentiating through Cholesky"
+    A.tril_().diagonal(offset=0, dim1=-2, dim2=-1).mul_(0.5)
+    return A
+
+
+def _cholesky_backward(dout_dL, L, L_inverse):
+    # c.f. https://github.com/pytorch/pytorch/blob/25ba802ce4cbdeaebcad4a03cec8502f0de9b7b3/
+    #      tools/autograd/templates/Functions.cpp
+    A = L.transpose(-1, -2) @ dout_dL
+    phi = _phi_for_cholesky_(A)
+    grad_input = (L_inverse.transpose(-1, -2) @ phi) @ L_inverse
+    # Symmetrize gradient
+    return grad_input.add(grad_input.transpose(-1, -2)).mul_(0.5)
+
+
+class _NaturalToMuVarSqrt(torch.autograd.Function):
+    @staticmethod
+    def _forward(nat_mean, nat_covar):
+        try:
+            L_inv = torch.cholesky(-2.0 * nat_covar, upper=False)
+        except RuntimeError as e:
+            if str(e).startswith("cholesky"):
+                raise RuntimeError(
+                    "Non-negative-definite natural covariance. You probably "
+                    "updated it using an optimizer other than gpytorch.optim.NGD (such as Adam). "
+                    "This is not supported."
+                )
+            else:
+                raise e
+        L = _triangular_inverse(L_inv, upper=False)
+        S = L.transpose(-1, -2) @ L
+        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)
+
+    @staticmethod
+    def forward(ctx, nat_mean, nat_covar):
+        mu, L = _NaturalToMuVarSqrt._forward(nat_mean, nat_covar)
+        ctx.save_for_backward(mu, L)
+        return mu, L
+
+    @staticmethod
+    def _backward(dout_dmu, dout_dL, mu, L, C):
+        """Calculate dout/d(eta1, eta2), which are:
+        eta1 = mu
+        eta2 = mu*mu^T + LL^T = mu*mu^T + Sigma
+
+        Thus:
+        dout/deta1 = dout/dmu + dout/dL dL/deta1
+        dout/deta2 = dout/dL dL/deta1
+
+        For L = chol(eta2 - eta1*eta1^T).
+        dout/dSigma = _cholesky_backward(dout/dL, L)
+        dout/deta2 = dout/dSigma
+        dSigma/deta1 = -2* (dout/dSigma) mu
+        """
+        dout_dSigma = _cholesky_backward(dout_dL, L, C)
+        dout_deta1 = dout_dmu - 2 * (dout_dSigma @ mu.unsqueeze(-1)).squeeze(-1)
+        return dout_deta1, dout_dSigma
+
+    @staticmethod
+    def backward(ctx, dout_dmu, dout_dL):
+        "Calculates the natural gradient with respect to nat_mean, nat_covar"
+        mu, L = ctx.saved_tensors
+        C = _triangular_inverse(L, upper=False)
+        return _NaturalToMuVarSqrt._backward(dout_dmu, dout_dL, mu, L, C)
diff --git a/gpytorch/variational/tril_natural_variational_distribution.py b/gpytorch/variational/tril_natural_variational_distribution.py
new file mode 100644
index 00000000..f6d85abd
--- /dev/null
+++ b/gpytorch/variational/tril_natural_variational_distribution.py
@@ -0,0 +1,118 @@
+import torch
+
+from ..distributions import MultivariateNormal
+from ..lazy import CholLazyTensor, TriangularLazyTensor
+from .natural_variational_distribution import (
+    _NaturalToMuVarSqrt,
+    _NaturalVariationalDistribution,
+    _phi_for_cholesky_,
+    _triangular_inverse,
+)
+
+
+class TrilNaturalVariationalDistribution(_NaturalVariationalDistribution):
+    r"""A multivariate normal :obj:`~gpytorch.variational._VariationalDistribution`,
+    parameterized by the natural vector, and a triangular decomposition of the
+    natural matrix (which is not the Cholesky).
+
+    .. note::
+       The :obj:`~gpytorch.variational.TrilNaturalVariationalDistribution` should only
+       be used with :obj:`gpytorch.optim.NGD`, or other optimizers
+       that follow exactly the gradient direction.
+
+    .. seealso::
+        The `natural gradient descent tutorial
+        <examples/04_Variational_and_Approximate_GPs/Natural_Gradient_Descent.ipynb>`_
+        for use instructions.
+
+        The :obj:`~gpytorch.variational.NaturalVariationalDistribution`, which
+        needs less iterations to make variational regression converge, at the
+        cost of introducing numerical instability.
+
+    .. note::
+        The relationship of the parameter :math:`\mathbf \Theta_\text{tril_mat}`
+        to the natural parameter :math:`\mathbf \Theta_\text{mat}` from
+        :obj:`~gpytorch.variational.NaturalVariationalDistribution` is
+        :math:`\mathbf \Theta_\text{mat} = -1/2 {\mathbf \Theta_\text{tril_mat}}^T {\mathbf \Theta_\text{tril_mat}}`.
+        Note that this is not the form of the Cholesky decomposition of :math:`\boldsymbol \Theta_\text{mat}`.
+
+    :param int num_inducing_points: Size of the variational distribution. This implies that the variational mean
+        should be this size, and the variational covariance matrix should have this many rows and columns.
+    :param batch_shape: Specifies an optional batch size
+        for the variational parameters. This is useful for example when doing additive variational inference.
+    :type batch_shape: :obj:`torch.Size`, optional
+    :param float mean_init_std: (Default: 1e-3) Standard deviation of gaussian noise to add to the mean initialization.
+    """
+
+    def __init__(self, num_inducing_points, batch_shape=torch.Size([]), mean_init_std=1e-3, **kwargs):
+        super().__init__(num_inducing_points=num_inducing_points, batch_shape=batch_shape, mean_init_std=mean_init_std)
+        scaled_mean_init = torch.zeros(num_inducing_points)
+        neg_prec_init = torch.eye(num_inducing_points, num_inducing_points)
+        scaled_mean_init = scaled_mean_init.repeat(*batch_shape, 1)
+        neg_prec_init = neg_prec_init.repeat(*batch_shape, 1, 1)
+
+        # eta1 and tril_dec(eta2) parameterization of the variational distribution
+        self.register_parameter(name="natural_vec", parameter=torch.nn.Parameter(scaled_mean_init))
+        self.register_parameter(name="natural_tril_mat", parameter=torch.nn.Parameter(neg_prec_init))
+
+    def forward(self):
+        mean, chol_covar = _TrilNaturalToMuVarSqrt.apply(self.natural_vec, self.natural_tril_mat)
+        return MultivariateNormal(mean, CholLazyTensor(TriangularLazyTensor(chol_covar)))
+
+    def initialize_variational_distribution(self, prior_dist):
+        prior_cov = prior_dist.lazy_covariance_matrix
+        chol = prior_cov.cholesky().evaluate()
+        tril_mat = _triangular_inverse(chol, upper=False)
+
+        natural_vec = prior_cov.inv_matmul(prior_dist.mean.unsqueeze(-1)).squeeze(-1)
+        noise = torch.randn_like(natural_vec).mul_(self.mean_init_std)
+
+        self.natural_vec.data.copy_(natural_vec.add_(noise))
+        self.natural_tril_mat.data.copy_(tril_mat)
+
+
+class _TrilNaturalToMuVarSqrt(torch.autograd.Function):
+    @staticmethod
+    def _forward(nat_mean, tril_nat_covar):
+        L = _triangular_inverse(tril_nat_covar, upper=False)
+        mu = L @ (L.transpose(-1, -2) @ nat_mean.unsqueeze(-1))
+        return mu.squeeze(-1), L
+        # return nat_mean, L
+
+    @staticmethod
+    def forward(ctx, nat_mean, tril_nat_covar):
+        mu, L = _TrilNaturalToMuVarSqrt._forward(nat_mean, tril_nat_covar)
+        ctx.save_for_backward(mu, L, tril_nat_covar)
+        return mu, L
+
+    @staticmethod
+    def backward(ctx, dout_dmu, dout_dL):
+        mu, L, C = ctx.saved_tensors
+        dout_dnat1, dout_dnat2 = _NaturalToMuVarSqrt._backward(dout_dmu, dout_dL, mu, L, C)
+        """
+        Now we need to do the Jacobian-Vector Product for the transformation:
+        L = inv(chol(inv(-2 theta_cov)))
+
+        C^T C = -2 theta_cov
+
+        so we need to do forward differentiation, starting with sensitivity (sensitivities marked with .dots.)
+        .theta_cov. = dout_dnat2
+
+        and ending with sensitivity .C.
+
+        if B = inv(-2 theta_cov) then:
+
+        .B.  =  d inv(-2 theta_cov)/dtheta_cov * .theta_cov.  =  -B (-2 .theta_cov.) B
+
+        if L = chol(B), B = LL^T then (https://homepages.inf.ed.ac.uk/imurray2/pub/16choldiff/choldiff.pdf):
+
+        .L. = L phi(L^{-1} .B. (L^{-1})^T) = L phi(2 L^T .theta_cov. L)
+
+        Then C = inv(L), so
+
+        .C. = -C .L. C = phi(-2 L^T .theta_cov. L)C
+        """
+        A = L.transpose(-2, -1) @ dout_dnat2 @ L
+        phi = _phi_for_cholesky_(A.mul_(-2))
+        dout_dtril = phi @ C
+        return dout_dnat1, dout_dtril
diff --git a/gpytorch/variational/whitened_variational_strategy.py b/gpytorch/variational/whitened_variational_strategy.py
index 6ad8f4ca..ffed6ef9 100644
--- a/gpytorch/variational/whitened_variational_strategy.py
+++ b/gpytorch/variational/whitened_variational_strategy.py
@@ -15,6 +15,7 @@ from ..lazy import (
     MatmulLazyTensor,
     PsdSumLazyTensor,
     RootLazyTensor,
+    TriangularLazyTensor,
 )
 from ..module import Module
 from ..utils.memoize import add_to_cache, cached, clear_cache_hook
@@ -137,7 +138,7 @@ class WhitenedVariationalStrategy(UnwhitenedVariationalStrategy):
             # If we're less than a certain size, we'll compute the Cholesky decomposition of induc_induc_covar
             cholesky = False
             if settings.fast_computations.log_prob.off() or (num_induc <= settings.max_cholesky_size.value()):
-                induc_induc_covar = CholLazyTensor(induc_induc_covar.cholesky())
+                induc_induc_covar = CholLazyTensor(TriangularLazyTensor(induc_induc_covar.cholesky()))
                 cholesky = True
 
             # Cache the CG results


Test patch:
diff --git a/test/optim/__init__.py b/test/optim/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/test/optim/test_ngd.py b/test/optim/test_ngd.py
new file mode 100644
index 00000000..0714eef3
--- /dev/null
+++ b/test/optim/test_ngd.py
@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+
+import unittest
+
+import torch
+
+from gpytorch.optim import NGD
+from gpytorch.test.base_test_case import BaseTestCase
+
+
+class TestNGD(unittest.TestCase, BaseTestCase):
+    def test_ngd_step_no_groups(self):
+        parameters = [
+            torch.nn.Parameter(torch.tensor(2.5)),
+            torch.nn.Parameter(torch.tensor([1.0, -0.5])),
+        ]
+        # parameters[0].grad = torch.tensor(1.)
+        # parameters[1].grad = torch.tensor([2., -1.])
+
+        optimizer = NGD(parameters, num_data=5, lr=0.1)
+        optimizer.zero_grad()
+        loss = parameters[0] + torch.dot(parameters[1], torch.tensor([2.0, -1]))
+        loss.backward()
+        optimizer.step()
+
+        self.assertAllClose(parameters[0], torch.tensor(2.0))
+        self.assertAllClose(parameters[1], torch.tensor([0.0, 0.0]))
+
+    def test_ngd_step_groups(self):
+        parameters = [
+            {"params": [torch.nn.Parameter(torch.tensor(2.5))], "lr": 0.2},
+            {"params": [torch.nn.Parameter(torch.tensor([1.0, -0.5]))]},
+        ]
+
+        optimizer = NGD(parameters, num_data=5, lr=0.1)
+        optimizer.zero_grad()
+        loss = parameters[0]["params"][0] + torch.dot(parameters[1]["params"][0], torch.tensor([2.0, -1]))
+        loss.backward()
+        optimizer.step()
+
+        self.assertAllClose(parameters[0]["params"][0], torch.tensor(1.5))
+        self.assertAllClose(parameters[1]["params"][0], torch.tensor([0.0, 0.0]))
diff --git a/test/variational/test_natural_variational_distribution.py b/test/variational/test_natural_variational_distribution.py
new file mode 100644
index 00000000..a35b78ff
--- /dev/null
+++ b/test/variational/test_natural_variational_distribution.py
@@ -0,0 +1,157 @@
+import unittest
+
+import torch
+
+import gpytorch
+from gpytorch.distributions import MultivariateNormal
+from gpytorch.lazy import CholLazyTensor, TriangularLazyTensor
+from gpytorch.variational import NaturalVariationalDistribution, TrilNaturalVariationalDistribution
+
+
+class Float64Test(unittest.TestCase):
+    def setUp(self):
+        self.prev_type = torch.get_default_dtype()
+        torch.set_default_dtype(torch.float64)
+
+    def tearDown(self):
+        torch.set_default_dtype(self.prev_type)
+
+
+class TestNatVariational(Float64Test):
+    def test_invertible_init(self, D=5):
+        mu = torch.randn(D)
+        cov = torch.randn(D, D).tril_()
+        dist = MultivariateNormal(mu, CholLazyTensor(TriangularLazyTensor(cov)))
+
+        v_dist = NaturalVariationalDistribution(D, mean_init_std=0.0)
+        v_dist.initialize_variational_distribution(dist)
+
+        out_dist = v_dist()
+
+        assert torch.allclose(out_dist.mean, dist.mean)
+        assert torch.allclose(out_dist.covariance_matrix, dist.covariance_matrix)
+
+    def test_natgrad(self, D=5):
+        mu = torch.randn(D)
+        cov = torch.randn(D, D).tril_()
+        dist = MultivariateNormal(mu, CholLazyTensor(TriangularLazyTensor(cov)))
+        sample = dist.sample()
+
+        v_dist = NaturalVariationalDistribution(D)
+        v_dist.initialize_variational_distribution(dist)
+        mu = v_dist().mean.detach()
+
+        v_dist().log_prob(sample).squeeze().backward()
+
+        eta1 = mu.clone().requires_grad_(True)
+        eta2 = (mu[:, None] * mu + cov @ cov.t()).requires_grad_(True)
+        L = torch.cholesky(eta2 - eta1[:, None] * eta1)
+        dist2 = MultivariateNormal(eta1, CholLazyTensor(TriangularLazyTensor(L)))
+        dist2.log_prob(sample).squeeze().backward()
+
+        assert torch.allclose(v_dist.natural_vec.grad, eta1.grad)
+        assert torch.allclose(v_dist.natural_mat.grad, eta2.grad)
+
+    def test_optimization_optimal_error(self, num_inducing=16, num_data=32, D=2):
+        inducing_points = torch.randn(num_inducing, D)
+
+        class SVGP(gpytorch.models.ApproximateGP):
+            def __init__(self):
+                v_dist = NaturalVariationalDistribution(num_inducing)
+                v_strat = gpytorch.variational.UnwhitenedVariationalStrategy(self, inducing_points, v_dist)
+                super().__init__(v_strat)
+                self.mean_module = gpytorch.means.ZeroMean()
+                self.covar_module = gpytorch.kernels.RBFKernel()
+
+            def forward(self, x):
+                return MultivariateNormal(self.mean_module(x), self.covar_module(x))
+
+        model = SVGP().train()
+        likelihood = gpytorch.likelihoods.GaussianLikelihood().train()
+        mll = gpytorch.mlls.VariationalELBO(likelihood, model, num_data)
+        X = torch.randn((num_data, D))
+        y = torch.randn(num_data)
+
+        def loss():
+            return -mll(model(X), y)
+
+        optimizer = torch.optim.SGD(
+            model.variational_strategy._variational_distribution.parameters(), lr=float(num_data)
+        )
+
+        optimizer.zero_grad()
+        loss().backward()
+        optimizer.step()  # Now we should be at the optimum
+
+        optimizer.zero_grad()
+        loss().backward()
+        natgrad_natural_vec2, natgrad_natural_mat2 = (
+            model.variational_strategy._variational_distribution.natural_vec.grad.clone(),
+            model.variational_strategy._variational_distribution.natural_mat.grad.clone(),
+        )
+        # At the optimum, the (natural) gradients are zero:
+        assert torch.allclose(natgrad_natural_vec2, torch.zeros(()))
+        assert torch.allclose(natgrad_natural_mat2, torch.zeros(()))
+
+
+class TestTrilNatVariational(Float64Test):
+    def test_invertible_init(self, D=5):
+        mu = torch.randn(D)
+        cov = torch.randn(D, D).tril_()
+        dist = MultivariateNormal(mu, CholLazyTensor(TriangularLazyTensor(cov)))
+
+        v_dist = TrilNaturalVariationalDistribution(D, mean_init_std=0.0)
+        v_dist.initialize_variational_distribution(dist)
+
+        out_dist = v_dist()
+
+        assert torch.allclose(out_dist.mean, dist.mean)
+        assert torch.allclose(out_dist.covariance_matrix, dist.covariance_matrix)
+
+    def test_natgrad(self, D=5):
+        mu = torch.randn(D)
+        cov = torch.randn(D, D)
+        cov = cov @ cov.t()
+        dist = MultivariateNormal(mu, CholLazyTensor(TriangularLazyTensor(cov.cholesky())))
+        sample = dist.sample()
+
+        v_dist = TrilNaturalVariationalDistribution(D, mean_init_std=0.0)
+        v_dist.initialize_variational_distribution(dist)
+        v_dist().log_prob(sample).squeeze().backward()
+        dout_dnat1 = v_dist.natural_vec.grad
+        dout_dnat2 = v_dist.natural_tril_mat.grad
+
+        # mean_init_std=0. because we need to ensure both have the same distribution
+        v_dist_ref = NaturalVariationalDistribution(D, mean_init_std=0.0)
+        v_dist_ref.initialize_variational_distribution(dist)
+        v_dist_ref().log_prob(sample).squeeze().backward()
+        dout_dnat1_noforward_ref = v_dist_ref.natural_vec.grad
+        dout_dnat2_noforward_ref = v_dist_ref.natural_mat.grad
+
+        def f(natural_vec, natural_tril_mat):
+            "Transform natural_tril_mat to L"
+            Sigma = torch.inverse(-2 * natural_tril_mat)
+            mu = natural_vec
+            return mu, Sigma.cholesky().inverse().tril()
+
+        (mu_ref, natural_tril_mat_ref), (dout_dmu_ref, dout_dnat2_ref) = jvp(
+            f,
+            (v_dist_ref.natural_vec.detach(), v_dist_ref.natural_mat.detach()),
+            (dout_dnat1_noforward_ref, dout_dnat2_noforward_ref),
+        )
+
+        assert torch.allclose(natural_tril_mat_ref, v_dist.natural_tril_mat), "Sigma transformation"
+        assert torch.allclose(dout_dnat2_ref, dout_dnat2), "Sigma gradient"
+
+        assert torch.allclose(mu_ref, v_dist.natural_vec), "mu transformation"
+        assert torch.allclose(dout_dmu_ref, dout_dnat1), "mu gradient"
+
+
+def jvp(f, x, v):
+    "Simulate forward-mode AD using two reverse-mode AD"
+    x = tuple(xx.requires_grad_(True) for xx in x)
+    v = tuple(vv.requires_grad_(True) for vv in v)
+    y = f(*x)
+    grad_x = torch.autograd.grad(y, x, v, create_graph=True)
+    jvp_val = torch.autograd.grad(grad_x, v, v)
+    return y, jvp_val