Repository: cornellius-gp/gpytorch
Pull request #163 (2018-06-15 18:02:13)

AdditiveKernel and ProductKernel now accept several kernels to compose
This was blocking some implementation a student was working on here. 

`AdditiveKernel` and `ProductKernel` now operate on a list of kernels, and accept several kernels in their constructors.


Patch:
diff --git a/examples/fast_variances_ski_(LOVE).ipynb b/examples/fast_variances_ski_(LOVE).ipynb
deleted file mode 100644
index eb6c1b4c..00000000
--- a/examples/fast_variances_ski_(LOVE).ipynb
+++ /dev/null
@@ -1,442 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Fast Predictive Distributions with SKI+LOVE\n",
-    "\n",
-    "## Overview\n",
-    "\n",
-    "In this notebook, we'll give a brief tutorial on how to use Lanczos Variance Estimates (LOVE) to achieve fast predictive distributions as described in this paper https://arxiv.org/abs/1803.06058. To see LOVE in use with exact GPs, see `fast_variances_exact_(LOVE).ipynb`. For a tutorial on using the fast sampling mechanism described in the paper, see `fast_sampling_ski_(LOVE).ipynb`.\n",
-    "\n",
-    "LOVE is an algorithm for approximating the predictive covariances of a Gaussian process in constant time after linear time precomputation. In this notebook, we will train a deep kernel learning model with SKI on one of the UCI datasets used in the paper, and then compare the time required to make predictions with each model.\n",
-    "\n",
-    "**NOTE**: The timing results reported in the paper compare the time required to compute (co)variances __only__. Because excluding the mean computations from the timing results requires hacking the internals of GPyTorch, the timing results presented in this notebook include the time required to compute predictive means, which are not accelerated by LOVE. Nevertheless, as we will see, LOVE achieves impressive speed-ups."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 121,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import math\n",
-    "import torch\n",
-    "import gpytorch\n",
-    "from matplotlib import pyplot as plt\n",
-    "from torch import nn, optim\n",
-    "from torch.autograd import Variable\n",
-    "from gpytorch.kernels import RBFKernel, GridInterpolationKernel\n",
-    "from gpytorch.means import ConstantMean\n",
-    "from gpytorch.likelihoods import GaussianLikelihood\n",
-    "from gpytorch.random_variables import GaussianRandomVariable\n",
-    "\n",
-    "# Make plots inline\n",
-    "%matplotlib inline"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Loading Data\n",
-    "\n",
-    "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.\n",
-    "\n",
-    "**Note**: Running the next cell will attempt to download a ~400 KB dataset file to the current directory."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {
-    "collapsed": true
-   },
-   "outputs": [],
-   "source": [
-    "import urllib.request\n",
-    "import os.path\n",
-    "from scipy.io import loadmat\n",
-    "from math import floor\n",
-    "\n",
-    "if 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",
-    "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",
-    "# Use the first 80% of the data for training, and the last 20% for testing.\n",
-    "train_n = int(floor(0.8*len(X)))\n",
-    "\n",
-    "train_x = X[:train_n, :].contiguous().cuda()\n",
-    "train_y = y[:train_n].contiguous().cuda()\n",
-    "\n",
-    "test_x = X[train_n:, :].contiguous().cuda()\n",
-    "test_y = y[train_n:].contiguous().cuda()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Defining the DKL Feature Extractor\n",
-    "\n",
-    "Next, we define the deep feature extractor we'll be using for DKL. In this case, we use a fully connected network with the architecture `d -> 1000 -> 500 -> 50 -> 2`, as described in the original DKL paper. All of the code below uses standard PyTorch implementations of neural network layers."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {
-    "collapsed": true
-   },
-   "outputs": [],
-   "source": [
-    "data_dim = train_x.size(-1)\n",
-    "\n",
-    "class LargeFeatureExtractor(nn.Sequential):           \n",
-    "    def __init__(self):                                      \n",
-    "        super(LargeFeatureExtractor, self).__init__()        \n",
-    "        self.add_module('linear1', nn.Linear(data_dim, 1000))\n",
-    "        self.add_module('relu1', nn.ReLU())                  \n",
-    "        self.add_module('linear2', nn.Linear(1000, 500))     \n",
-    "        self.add_module('relu2', nn.ReLU())                  \n",
-    "        self.add_module('linear3', nn.Linear(500, 50))       \n",
-    "        self.add_module('relu3', nn.ReLU())                  \n",
-    "        self.add_module('linear4', nn.Linear(50, 2))         \n",
-    "                                                             \n",
-    "feature_extractor = LargeFeatureExtractor().cuda()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Defining the GP Model\n",
-    "\n",
-    "We now define the GP model. For more details on the use of GP models, see our simpler examples. This model uses a `GridInterpolationKernel` (SKI) with an RBF base kernel. The forward method passes the input data `x` through the neural network feature extractor defined above, scales the resulting features to be between 0 and 1, and then calls the kernel."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {
-    "collapsed": true
-   },
-   "outputs": [],
-   "source": [
-    "class GPRegressionModel(gpytorch.models.ExactGP):\n",
-    "        def __init__(self, train_x, train_y, likelihood):\n",
-    "            super(GPRegressionModel, self).__init__(train_x, train_y, likelihood)\n",
-    "            self.mean_module = ConstantMean(constant_bounds=[-1e-5, 1e-5])\n",
-    "            self.base_covar_module = RBFKernel(log_lengthscale_bounds=(-5, 6))\n",
-    "\n",
-    "            self.covar_module = GridInterpolationKernel(self.base_covar_module, grid_size=100,\n",
-    "                                                        grid_bounds=[(-1.1, 1.1), (-1.1, 1.1)])\n",
-    "            self.register_parameter('log_outputscale', nn.Parameter(torch.Tensor([0])), bounds=(-5, 6))\n",
-    "            self.feature_extractor = feature_extractor\n",
-    "\n",
-    "        def forward(self, x):\n",
-    "            projected_x = self.feature_extractor(x)\n",
-    "            projected_x = projected_x - projected_x.min(0)[0]\n",
-    "            projected_x = 2 * (projected_x / projected_x.max(0)[0]) - 1\n",
-    "            mean_x = self.mean_module(projected_x)\n",
-    "            covar_x = self.covar_module(projected_x)\n",
-    "            covar_x = covar_x.mul(self.log_outputscale.exp())\n",
-    "            return GaussianRandomVariable(mean_x, covar_x)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {
-    "collapsed": true
-   },
-   "outputs": [],
-   "source": [
-    "likelihood = GaussianLikelihood(log_noise_bounds=(-5, 5)).cuda()\n",
-    "model = GPRegressionModel(train_x, train_y, likelihood).cuda()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Training the model\n",
-    "\n",
-    "The cell below trains the DKL model above, finding optimal hyperparameters using Type-II MLE. We run 20 iterations of training using the `Adam` optimizer built in to PyTorch. With a decent GPU, this should only take a few seconds."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {
-    "scrolled": false
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Iter 1/20 - Loss: 0.949\n",
-      "Iter 2/20 - Loss: 0.892\n",
-      "Iter 3/20 - Loss: 0.843\n",
-      "Iter 4/20 - Loss: 0.795\n",
-      "Iter 5/20 - Loss: 0.747\n",
-      "Iter 6/20 - Loss: 0.699\n",
-      "Iter 7/20 - Loss: 0.651\n",
-      "Iter 8/20 - Loss: 0.604\n",
-      "Iter 9/20 - Loss: 0.556\n",
-      "Iter 10/20 - Loss: 0.508\n",
-      "Iter 11/20 - Loss: 0.460\n",
-      "Iter 12/20 - Loss: 0.412\n",
-      "Iter 13/20 - Loss: 0.364\n",
-      "Iter 14/20 - Loss: 0.315\n",
-      "Iter 15/20 - Loss: 0.271\n",
-      "Iter 16/20 - Loss: 0.219\n",
-      "Iter 17/20 - Loss: 0.174\n",
-      "Iter 18/20 - Loss: 0.126\n",
-      "Iter 19/20 - Loss: 0.075\n",
-      "Iter 20/20 - Loss: 0.023\n",
-      "CPU times: user 4.08 s, sys: 1.32 s, total: 5.39 s\n",
-      "Wall time: 5.36 s\n"
-     ]
-    }
-   ],
-   "source": [
-    "# Find optimal model hyperparameters\n",
-    "model.train()\n",
-    "likelihood.train()\n",
-    "\n",
-    "# Use the adam optimizer\n",
-    "optimizer = torch.optim.Adam([\n",
-    "    {'params': model.parameters()},  # Includes GaussianLikelihood parameters\n",
-    "], lr=0.1)\n",
-    "\n",
-    "# \"Loss\" for GPs - the marginal log likelihood\n",
-    "mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)\n",
-    "\n",
-    "training_iterations = 20\n",
-    "def train():\n",
-    "    for i in range(training_iterations):\n",
-    "        # Zero backprop gradients\n",
-    "        optimizer.zero_grad()\n",
-    "        # Get output from model\n",
-    "        output = model(train_x)\n",
-    "        # Calc loss and backprop derivatives\n",
-    "        loss = -mll(output, train_y)\n",
-    "        loss.backward()\n",
-    "        print('Iter %d/%d - Loss: %.3f' % (i + 1,\n",
-    "                                           training_iterations,\n",
-    "                                           loss.data[0]))\n",
-    "        optimizer.step()\n",
-    "        \n",
-    "# See dkl_mnist.ipynb for explanation of this flag\n",
-    "with gpytorch.settings.use_toeplitz(False):\n",
-    "    %time train()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Make Predictions using Standard SKI Code\n",
-    "\n",
-    "The next cell gets the predictive covariance for the test set (and also technically gets the predictive mean, stored in `preds.mean()`) using the standard SKI testing code, with no acceleration or precomputation. \n",
-    "\n",
-    "**Note:** Full predictive covariance matrices (and the computations needed to get them) can be quite memory intensive. Depending on the memory available on your GPU, you may need to reduce the size of the test set for the code below to run. If you run out of memory, try replacing `test_x` below with something like `test_x[:1000]` to use the first 1000 test points only, and then restart the notebook."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 113,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import time\n",
-    "\n",
-    "# Set into eval mode\n",
-    "model.eval()\n",
-    "likelihood.eval()\n",
-    "with torch.no_grad(), gpytorch.settings.use_toeplitz(False):\n",
-    "    start_time = time.time()\n",
-    "    preds = model(test_x)\n",
-    "    exact_covar = preds.covar().evaluate()\n",
-    "    exact_covar_time = time.time() - start_time"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 114,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Time to compute exact mean + covariances: 1.96s\n"
-     ]
-    }
-   ],
-   "source": [
-    "print('Time to compute exact mean + covariances: {:.2f}s'.format(exact_covar_time))"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Clear Memory and any Precomputed Values\n",
-    "\n",
-    "The next cell clears as much as possible to avoid influencing the timing results of the fast predictive variances code. Strictly speaking, the timing results above and the timing results to follow should be run in entirely separate notebooks. However,  this will suffice for this simple example."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 112,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "GaussianLikelihood()"
-      ]
-     },
-     "execution_count": 112,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "# Clear as much 'stuff' as possible\n",
-    "import gc\n",
-    "gc.collect()\n",
-    "torch.cuda.empty_cache()\n",
-    "model.train()\n",
-    "likelihood.train()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Compute Predictions with LOVE, but Before Precomputation\n",
-    "\n",
-    "Next we compute predictive covariances (and the predictive means) for LOVE, but starting from scratch. That is, we don't yet have access to the precomputed cache discussed in the paper. This should still be faster than the full covariance computation code above.\n",
-    "\n",
-    "In this simple example, we allow a rank 10 root decomposition, although increasing this to rank 20-40 should not affect the timing results substantially."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 115,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Set into eval mode\n",
-    "model.eval()\n",
-    "likelihood.eval()\n",
-    "with torch.no_grad(), gpytorch.settings.use_toeplitz(False), gpytorch.beta_features.fast_pred_var(), gpytorch.settings.max_root_decomposition_size(10):\n",
-    "    start_time = time.time()\n",
-    "    preds = model(test_x)\n",
-    "    fast_time_no_cache = time.time() - start_time"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Compute Predictions with LOVE After Precomputation\n",
-    "\n",
-    "The above cell additionally computed the caches required to get fast predictions. From this point onwards, unless we put the model back in training mode, predictions should be extremely fast. The cell below re-runs the above code, but takes full advantage of both the mean cache and the LOVE cache for variances."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 116,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "with torch.no_grad(), gpytorch.settings.use_toeplitz(False), gpytorch.beta_features.fast_pred_var(), gpytorch.settings.max_root_decomposition_size(10):\n",
-    "    start_time = time.time()\n",
-    "    preds = model(test_x)\n",
-    "    fast_covar = preds.covar().evaluate()\n",
-    "    fast_time_with_cache = time.time() - start_time"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 119,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Time to compute mean + covariances (no cache) 0.16s\n",
-      "Time to compute mean + variances (cache): 0.06s\n"
-     ]
-    }
-   ],
-   "source": [
-    "print('Time to compute mean + covariances (no cache) {:.2f}s'.format(fast_time_no_cache))\n",
-    "print('Time to compute mean + variances (cache): {:.2f}s'.format(fast_time_with_cache))"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Compute Error between Exact and Fast Variances\n",
-    "\n",
-    "Finally, we compute the mean absolute error between the fast variances computed by LOVE (stored in fast_covar), and the exact variances computed previously. \n",
-    "\n",
-    "Note that these tests were run with a root decomposition of rank 10, which is about the minimum you would realistically ever run with. Despite this, the fast variance estimates are quite good. If more accuracy was needed, increasing `max_root_decomposition_size` to 30 or 40 would provide even better estimates."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 120,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "MAE between exact covar matrix and fast covar matrix: 2.6490910386200994e-05\n"
-     ]
-    }
-   ],
-   "source": [
-    "print('MAE between exact covar matrix and fast covar matrix: {}'.format((exact_covar - fast_covar).abs().mean()))"
-   ]
-  }
- ],
- "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.6.3"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}
diff --git a/gpytorch/kernels/__init__.py b/gpytorch/kernels/__init__.py
index f982faeb..aa4fb0d6 100644
--- a/gpytorch/kernels/__init__.py
+++ b/gpytorch/kernels/__init__.py
@@ -1,4 +1,4 @@
-from .kernel import Kernel
+from .kernel import Kernel, AdditiveKernel, ProductKernel
 from .rbf_kernel import RBFKernel
 from .matern_kernel import MaternKernel
 from .periodic_kernel import PeriodicKernel
@@ -12,11 +12,13 @@ from .multiplicative_grid_interpolation_kernel import MultiplicativeGridInterpol
 from .white_noise_kernel import WhiteNoiseKernel
 
 __all__ = [
+    AdditiveKernel,
     Kernel,
     LinearKernel,
     RBFKernel,
     MaternKernel,
     PeriodicKernel,
+    ProductKernel,
     SpectralMixtureKernel,
     IndexKernel,
     GridInterpolationKernel,
diff --git a/gpytorch/kernels/kernel.py b/gpytorch/kernels/kernel.py
index 5fc9c411..3beeaf07 100644
--- a/gpytorch/kernels/kernel.py
+++ b/gpytorch/kernels/kernel.py
@@ -5,7 +5,8 @@ from __future__ import unicode_literals
 
 import torch
 from ..module import Module
-from ..lazy import LazyEvaluatedKernelVariable
+from ..lazy import LazyEvaluatedKernelVariable, ZeroLazyVariable
+import gpytorch
 
 
 class Kernel(Module):
@@ -63,20 +64,22 @@ class Kernel(Module):
 
 
 class AdditiveKernel(Kernel):
-    def __init__(self, kernel_1, kernel_2):
+    def __init__(self, *kernels):
         super(AdditiveKernel, self).__init__()
-        self.kernel_1 = kernel_1
-        self.kernel_2 = kernel_2
+        self.kernels = kernels
 
     def forward(self, x1, x2):
-        return self.kernel_1(x1, x2).evaluate_kernel() + self.kernel_2(x1, x2).evaluate_kernel()
+        res = ZeroLazyVariable()
+        for kern in self.kernels:
+            res = res + kern(x1, x2).evaluate_kernel()
+
+        return res
 
 
 class ProductKernel(Kernel):
-    def __init__(self, kernel_1, kernel_2):
+    def __init__(self, *kernels):
         super(ProductKernel, self).__init__()
-        self.kernel_1 = kernel_1
-        self.kernel_2 = kernel_2
+        self.kernels = kernels
 
     def forward(self, x1, x2):
-        return self.kernel_1(x1, x2).evaluate_kernel() * self.kernel_2(x1, x2).evaluate_kernel()
+        return gpytorch.utils.prod([k(x1, x2).evaluate_kernel() for k in self.kernels])
diff --git a/gpytorch/lazy/lazy_evaluated_kernel_variable.py b/gpytorch/lazy/lazy_evaluated_kernel_variable.py
index 0a67f90b..96662570 100644
--- a/gpytorch/lazy/lazy_evaluated_kernel_variable.py
+++ b/gpytorch/lazy/lazy_evaluated_kernel_variable.py
@@ -38,9 +38,7 @@ class LazyEvaluatedKernelVariable(LazyVariable):
         Implementing it this way allows us to compute predictions more efficiently
         in cases where only the variances are required.
         """
-        from ..kernels import Kernel, GridInterpolationKernel
-        if isinstance(self.kernel, GridInterpolationKernel):
-            return self.evaluate_kernel().diag()
+        from ..kernels import Kernel
 
         if hasattr(self, "_cached_kernel_diag"):
             return self._cached_kernel_diag


Test patch:
diff --git a/test/kernels/test_additive_kernel.py b/test/kernels/test_additive_kernel.py
index 03c6254d..48471e22 100644
--- a/test/kernels/test_additive_kernel.py
+++ b/test/kernels/test_additive_kernel.py
@@ -6,7 +6,7 @@ from __future__ import unicode_literals
 import math
 import torch
 import unittest
-from gpytorch.kernels import RBFKernel
+from gpytorch.kernels import RBFKernel, AdditiveKernel, ProductKernel
 
 
 class TestAdditiveKernel(unittest.TestCase):
@@ -40,6 +40,38 @@ class TestAdditiveKernel(unittest.TestCase):
         res = kernel(a, b).evaluate()
         self.assertLess(torch.norm(res - actual), 2e-5)
 
+    def test_computes_product_of_three_radial_basis_function(self):
+        a = torch.Tensor([4, 2, 8]).view(3, 1)
+        b = torch.Tensor([0, 2]).view(2, 1)
+        lengthscale = 2
+
+        kernel_1 = RBFKernel().initialize(log_lengthscale=math.log(lengthscale))
+        kernel_2 = RBFKernel().initialize(log_lengthscale=math.log(lengthscale))
+        kernel_3 = RBFKernel().initialize(log_lengthscale=math.log(lengthscale))
+        kernel = ProductKernel(kernel_1, kernel_2, kernel_3)
+
+        actual = torch.Tensor([[16, 4], [4, 0], [64, 36]]).mul_(-0.5).div_(lengthscale ** 2).exp() ** 3
+
+        kernel.eval()
+        res = kernel(a, b).evaluate()
+        self.assertLess(torch.norm(res - actual), 2e-5)
+
+    def test_computes_sum_of_three_radial_basis_function(self):
+        a = torch.Tensor([4, 2, 8]).view(3, 1)
+        b = torch.Tensor([0, 2]).view(2, 1)
+        lengthscale = 2
+
+        kernel_1 = RBFKernel().initialize(log_lengthscale=math.log(lengthscale))
+        kernel_2 = RBFKernel().initialize(log_lengthscale=math.log(lengthscale))
+        kernel_3 = RBFKernel().initialize(log_lengthscale=math.log(lengthscale))
+        kernel = AdditiveKernel(kernel_1, kernel_2, kernel_3)
+
+        actual = torch.Tensor([[16, 4], [4, 0], [64, 36]]).mul_(-0.5).div_(lengthscale ** 2).exp() * 3
+
+        kernel.eval()
+        res = kernel(a, b).evaluate()
+        self.assertLess(torch.norm(res - actual), 2e-5)
+
     def test_computes_sum_radial_basis_function_gradient(self):
         a = torch.Tensor([4, 2, 8]).view(3, 1)
         b = torch.Tensor([0, 2, 2]).view(3, 1)
@@ -59,7 +91,34 @@ class TestAdditiveKernel(unittest.TestCase):
 
         output = kernel(a, b).evaluate()
         output.backward(gradient=torch.eye(3))
-        res = kernel.kernel_1.log_lengthscale.grad + kernel.kernel_2.log_lengthscale.grad
+        res = kernel.kernels[0].log_lengthscale.grad + kernel.kernels[1].log_lengthscale.grad
+        self.assertLess(torch.norm(res - actual_param_grad), 2e-5)
+
+    def test_computes_sum_three_radial_basis_function_gradient(self):
+        a = torch.Tensor([4, 2, 8]).view(3, 1)
+        b = torch.Tensor([0, 2, 2]).view(3, 1)
+        lengthscale = 2
+
+        param = math.log(lengthscale) * torch.ones(3, 3)
+        param.requires_grad_()
+        diffs = a.expand(3, 3) - b.expand(3, 3).transpose(0, 1)
+        actual_output = (-0.5 * (diffs / param.exp()) ** 2).exp()
+        actual_output.backward(torch.eye(3))
+        actual_param_grad = param.grad.sum() * 3
+
+        kernel_1 = RBFKernel().initialize(log_lengthscale=math.log(lengthscale))
+        kernel_2 = RBFKernel().initialize(log_lengthscale=math.log(lengthscale))
+        kernel_3 = RBFKernel().initialize(log_lengthscale=math.log(lengthscale))
+        kernel = AdditiveKernel(kernel_1, kernel_2, kernel_3)
+        kernel.eval()
+
+        output = kernel(a, b).evaluate()
+        output.backward(gradient=torch.eye(3))
+        res = (
+            kernel.kernels[0].log_lengthscale.grad
+            + kernel.kernels[1].log_lengthscale.grad
+            + kernel.kernels[2].log_lengthscale.grad
+        )
         self.assertLess(torch.norm(res - actual_param_grad), 2e-5)