Repository: cornellius-gp/gpytorch
Pull request #821 (2019-08-01 18:13:30)

Add a "Newton-Girard additive kernel" of Duvenaud et al 2011
Paper: https://arxiv.org/abs/1112.4394

This branch implements the "Additive Gaussian Process" kernel from the link above which uses the Newton-Girard formulae to efficiently compute higher-order additive kernels. The implementation is simple, just wrapping a base kernel, similar to AdditiveStructureKernel and in the forward computes the higher-order terms using the Newton-Girard formulae. This branch also implements some fairly basic tests for it.

Since this kernel can achieve state-of-the-art regression performance, it might be nice to include this kernel in GPyTorch.


Patch:
diff --git a/gpytorch/kernels/__init__.py b/gpytorch/kernels/__init__.py
index 46563a1d..91650bcb 100644
--- a/gpytorch/kernels/__init__.py
+++ b/gpytorch/kernels/__init__.py
@@ -13,6 +13,7 @@ from .lcm_kernel import LCMKernel
 from .linear_kernel import LinearKernel
 from .matern_kernel import MaternKernel
 from .multitask_kernel import MultitaskKernel
+from .newton_girard_additive_kernel import NewtonGirardAdditiveKernel
 from .periodic_kernel import PeriodicKernel
 from .polynomial_kernel import PolynomialKernel
 from .polynomial_kernel_grad import PolynomialKernelGrad
@@ -39,6 +40,7 @@ __all__ = [
     "LinearKernel",
     "MaternKernel",
     "MultitaskKernel",
+    "NewtonGirardAdditiveKernel",
     "PeriodicKernel",
     "PolynomialKernel",
     "PolynomialKernelGrad",
diff --git a/gpytorch/kernels/newton_girard_additive_kernel.py b/gpytorch/kernels/newton_girard_additive_kernel.py
new file mode 100644
index 00000000..6d57c08b
--- /dev/null
+++ b/gpytorch/kernels/newton_girard_additive_kernel.py
@@ -0,0 +1,110 @@
+import torch
+from .kernel import Kernel
+from ..constraints import Positive
+
+
+class NewtonGirardAdditiveKernel(Kernel):
+    def __init__(self, base_kernel, num_dims, max_degree=None, active_dims=None, **kwargs):
+        """Create an Additive Kernel a la https://arxiv.org/abs/1112.4394 using Newton-Girard Formulae
+
+        :param base_kernel: a base 1-dimensional kernel. NOTE: put ard_num_dims=d in the base kernel...
+        :param max_degree: the maximum numbers of kernel degrees to compute
+        :param active_dims:
+        :param kwargs:
+        """
+        super(NewtonGirardAdditiveKernel, self).__init__(has_lengthscale=False,
+                                                         active_dims=active_dims,
+                                                         **kwargs
+                                                         )
+
+        self.base_kernel = base_kernel
+        self.num_dims = num_dims
+        if max_degree is None:
+            self.max_degree = self.num_dims
+        elif max_degree > self.num_dims:  # force cap on max_degree (silently)
+            self.max_degree = self.num_dims
+        else:
+            self.max_degree = max_degree
+
+        self.register_parameter(
+            name='raw_outputscale',
+            parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, self.max_degree))
+        )
+        outputscale_constraint = Positive()
+        self.register_constraint('raw_outputscale', outputscale_constraint)
+        self.outputscale_constraint = outputscale_constraint
+        self.outputscale = [1 / self.max_degree for _ in range(self.max_degree)]
+
+    @property
+    def outputscale(self):
+        return self.raw_outputscale_constraint.transform(self.raw_outputscale)
+
+    @outputscale.setter
+    def outputscale(self, value):
+        self._set_outputscale(value)
+
+    def _set_outputscale(self, value):
+        if not torch.is_tensor(value):
+            value = torch.as_tensor(value).to(self.raw_outputscale)
+
+        self.initialize(raw_outputscale=self.outputscale_constraint.inverse_transform(value))
+
+    def forward(self, x1, x2, diag=False, last_dim_is_batch=False, **params):
+        """Forward proceeds by Newton-Girard formulae"""
+        if last_dim_is_batch:
+            raise RuntimeError("NewtonGirardAdditiveKernel does not accept the last_dim_is_batch argument.")
+
+        # NOTE: comments about shape are only correct for the single-batch cases.
+        # kern_values is just the order-1 terms
+        # kern_values = D x n x n unless diag=True
+        kern_values = self.base_kernel(x1, x2, diag=diag, last_dim_is_batch=True, **params)
+        # last dim is batch, which gets moved up to pos. 1
+
+        kernel_dim = -3 if not diag else -2
+
+        shape = [1 for _ in range(len(kern_values.shape) + 1)]
+        shape[kernel_dim - 1] = -1
+        kvals = torch.range(1, self.max_degree, device=kern_values.device).reshape(*shape)
+        # kvals = R x 1 x 1 x 1 (these are indexes only)
+
+        # e_n = torch.ones(self.max_degree+1, *kern_values.shape[1:], device=kern_values.device)  # includes 0
+        # e_n: elementary symmetric polynomial of degree n (e.g. z1 z2 + z1 z3 + z2 z3)
+        # e_n is R x n x n, and the array is properly 0 indexed.
+        shape = [d_ for d_ in kern_values.shape]
+        shape[kernel_dim] = self.max_degree + 1
+        e_n = torch.empty(*shape, device=kern_values.device)
+        if kernel_dim == -3:
+            e_n[..., 0, :, :] = 1.0
+        else:
+            e_n[..., 0, :] = 1.0
+
+        # power sums s_k (e.g. sum_i^num_dims z_i^k
+        # s_k is R x n x n
+        s_k = kern_values.unsqueeze(kernel_dim - 1).pow(kvals).sum(dim=kernel_dim)
+
+        # just the constant -1
+        m1 = torch.tensor([-1], dtype=torch.float, device=kern_values.device)
+
+        shape = [1 for _ in range(len(kern_values.shape))]
+        shape[kernel_dim] = -1
+        for deg in range(1, self.max_degree + 1):  # deg goes from 1 to R (it's 1-indexed!)
+            # we avg over k [1, ..., deg] (-1)^(k-1)e_{deg-k} s_{k}
+
+            ks = torch.arange(1, deg + 1, device=kern_values.device, dtype=torch.float).reshape(*shape)  # use for pow
+            kslong = torch.arange(1, deg + 1, device=kern_values.device, dtype=torch.long)  # use for indexing
+
+            # note that s_k is 0-indexed, so we must subtract 1 from kslong
+            sum_ = (m1.pow(ks - 1)
+                    * e_n.index_select(kernel_dim, deg - kslong)
+                    * s_k.index_select(kernel_dim, kslong - 1)
+                    ).sum(dim=kernel_dim) / deg
+            if kernel_dim == -3:
+                e_n[..., deg, :, :] = sum_
+            else:
+                e_n[..., deg, :] = sum_
+
+        if kernel_dim == -3:
+            return (self.outputscale.unsqueeze(-1).unsqueeze(-1)
+                    * e_n.narrow(kernel_dim, 1, self.max_degree)).sum(dim=kernel_dim)
+        else:
+            return (self.outputscale.unsqueeze(-1) * e_n.narrow(kernel_dim, 1, self.max_degree)).sum(dim=kernel_dim)


Test patch:
diff --git a/test/kernels/test_newton_girard_additive_kernel.py b/test/kernels/test_newton_girard_additive_kernel.py
new file mode 100644
index 00000000..cfeefc01
--- /dev/null
+++ b/test/kernels/test_newton_girard_additive_kernel.py
@@ -0,0 +1,156 @@
+#!/usr/bin/env python3
+
+from unittest import TestCase
+from gpytorch.test.base_kernel_test_case import BaseKernelTestCase
+from gpytorch.kernels import RBFKernel, ScaleKernel, AdditiveKernel, NewtonGirardAdditiveKernel
+from gpytorch.likelihoods import GaussianLikelihood
+from gpytorch.mlls import ExactMarginalLogLikelihood
+from gpytorch.models import ExactGP
+from gpytorch.means import ConstantMean
+from gpytorch.distributions import MultivariateNormal
+import torch
+
+
+class TestNewtonGirardAdditiveKernel(TestCase, BaseKernelTestCase):
+    def create_kernel_no_ard(self, **kwargs):
+        return NewtonGirardAdditiveKernel(RBFKernel(), 4, 2, **kwargs)
+
+    def create_kernel_ard(self, num_dims, **kwargs):
+        return NewtonGirardAdditiveKernel(RBFKernel(ard_num_dims=num_dims), num_dims, 2, **kwargs)
+
+    def test_degree1(self):
+        AddK = NewtonGirardAdditiveKernel(RBFKernel(ard_num_dims=3), 3, 1)
+        self.assertEqual(AddK.base_kernel.lengthscale.numel(), 3)
+        self.assertEqual(AddK.outputscale.numel(), 1)
+
+        testvals = torch.tensor([[1, 2, 3], [7, 5, 2]], dtype=torch.float)
+        add_k_val = AddK(testvals, testvals).evaluate()
+
+        manual_k = ScaleKernel(AdditiveKernel(RBFKernel(active_dims=0),
+                                              RBFKernel(active_dims=1),
+                                              RBFKernel(active_dims=2)))
+        manual_k.initialize(outputscale=1.)
+        manual_add_k_val = manual_k(testvals, testvals).evaluate()
+
+        # np.testing.assert_allclose(add_k_val.detach().numpy(), manual_add_k_val.detach().numpy(), atol=1e-5)
+        self.assertTrue(torch.allclose(add_k_val, manual_add_k_val, atol=1e-5))
+
+    def test_degree2(self):
+        AddK = NewtonGirardAdditiveKernel(RBFKernel(ard_num_dims=3), 3, 2)
+        self.assertEqual(AddK.base_kernel.lengthscale.numel(), 3)
+        self.assertEqual(AddK.outputscale.numel(), 2)
+
+        testvals = torch.tensor([[1, 2, 3], [7, 5, 2]], dtype=torch.float)
+        add_k_val = AddK(testvals, testvals).evaluate()
+
+        manual_k1 = ScaleKernel(AdditiveKernel(RBFKernel(active_dims=0),
+                                               RBFKernel(active_dims=1),
+                                               RBFKernel(active_dims=2)))
+        manual_k1.initialize(outputscale=1 / 2)
+        manual_k2 = ScaleKernel(AdditiveKernel(RBFKernel(active_dims=[0, 1]),
+                                               RBFKernel(active_dims=[1, 2]),
+                                               RBFKernel(active_dims=[0, 2])))
+        manual_k2.initialize(outputscale=1 / 2)
+        manual_k = AdditiveKernel(manual_k1, manual_k2)
+        manual_add_k_val = manual_k(testvals, testvals).evaluate()
+
+        # np.testing.assert_allclose(add_k_val.detach().numpy(), manual_add_k_val.detach().numpy(), atol=1e-5)
+        self.assertTrue(torch.allclose(add_k_val, manual_add_k_val, atol=1e-5))
+
+    def test_degree3(self):
+        # just make sure it doesn't break here.
+        AddK = NewtonGirardAdditiveKernel(RBFKernel(ard_num_dims=3), 3, 3)
+        self.assertEqual(AddK.base_kernel.lengthscale.numel(), 3)
+        self.assertEqual(AddK.outputscale.numel(), 3)
+
+        testvals = torch.tensor([[1, 2, 3], [7, 5, 2]], dtype=torch.float)
+        add_k_val = AddK(testvals, testvals).evaluate()
+
+        manual_k1 = ScaleKernel(AdditiveKernel(RBFKernel(active_dims=0),
+                                               RBFKernel(active_dims=1),
+                                               RBFKernel(active_dims=2)))
+        manual_k1.initialize(outputscale=1 / 3)
+        manual_k2 = ScaleKernel(AdditiveKernel(RBFKernel(active_dims=[0, 1]),
+                                               RBFKernel(active_dims=[1, 2]),
+                                               RBFKernel(active_dims=[0, 2])))
+        manual_k2.initialize(outputscale=1 / 3)
+
+        manual_k3 = ScaleKernel(AdditiveKernel(RBFKernel()))
+        manual_k3.initialize(outputscale=1 / 3)
+        manual_k = AdditiveKernel(manual_k1, manual_k2, manual_k3)
+        manual_add_k_val = manual_k(testvals, testvals).evaluate()
+        # np.testing.assert_allclose(add_k_val.detach().numpy(), manual_add_k_val.detach().numpy(), atol=1e-5)
+        self.assertTrue(torch.allclose(add_k_val, manual_add_k_val, atol=1e-5))
+
+    def test_optimizing(self):
+        # This tests should pass so long as nothing breaks.
+        torch.random.manual_seed(1)
+        data = torch.randn(40, 4)
+        target = torch.sin(data).sum(dim=-1)
+        d = 4
+
+        AddK = NewtonGirardAdditiveKernel(RBFKernel(ard_num_dims=d), d, max_degree=3)
+
+        class TestGPModel(ExactGP):
+            def __init__(self, train_x, train_y, likelihood, kernel):
+                super().__init__(train_x, train_y, likelihood)
+                self.mean_module = ConstantMean()
+                self.covar_module = kernel
+
+            def forward(self, x):
+                mean_x = self.mean_module(x)
+                covar_x = self.covar_module(x)
+                return MultivariateNormal(mean_x, covar_x)
+
+        model = TestGPModel(data, target, GaussianLikelihood(), ScaleKernel(AddK))
+        optim = torch.optim.Adam(model.parameters(), lr=0.1)
+        mll = ExactMarginalLogLikelihood(model.likelihood, model)
+        model.train()
+        for i in range(2):
+            optim.zero_grad()
+            out = model(data)
+            loss = -mll(out, target)
+            loss.backward()
+            optim.step()
+
+    def test_ard(self):
+        base_k = RBFKernel(ard_num_dims=3)
+        base_k.initialize(lengthscale=[1., 2., 3.])
+        AddK = NewtonGirardAdditiveKernel(base_k, 3, max_degree=1)
+
+        testvals = torch.tensor([[1, 2, 3], [7, 5, 2]], dtype=torch.float)
+        add_k_val = AddK(testvals, testvals).evaluate()
+
+        ks = []
+        for i in range(3):
+            k = RBFKernel(active_dims=i)
+            k.initialize(lengthscale=i + 1)
+            ks.append(k)
+        manual_k = ScaleKernel(AdditiveKernel(*ks))
+        manual_k.initialize(outputscale=1.)
+        manual_add_k_val = manual_k(testvals, testvals).evaluate()
+
+        # np.testing.assert_allclose(add_k_val.detach().numpy(), manual_add_k_val.detach().numpy(), atol=1e-5)
+        self.assertTrue(torch.allclose(add_k_val, manual_add_k_val, atol=1e-5))
+
+    def test_diag(self):
+        AddK = NewtonGirardAdditiveKernel(RBFKernel(ard_num_dims=3), 3, 2)
+        self.assertEqual(AddK.base_kernel.lengthscale.numel(), 3)
+        self.assertEqual(AddK.outputscale.numel(), 2)
+
+        testvals = torch.tensor([[1, 2, 3], [7, 5, 2]], dtype=torch.float)
+        add_k_val = AddK(testvals, testvals).diag()
+
+        manual_k1 = ScaleKernel(AdditiveKernel(RBFKernel(active_dims=0),
+                                               RBFKernel(active_dims=1),
+                                               RBFKernel(active_dims=2)))
+        manual_k1.initialize(outputscale=1 / 2)
+        manual_k2 = ScaleKernel(AdditiveKernel(RBFKernel(active_dims=[0, 1]),
+                                               RBFKernel(active_dims=[1, 2]),
+                                               RBFKernel(active_dims=[0, 2])))
+        manual_k2.initialize(outputscale=1 / 2)
+        manual_k = AdditiveKernel(manual_k1, manual_k2)
+        manual_add_k_val = manual_k(testvals, testvals).diag()
+
+        # np.testing.assert_allclose(add_k_val.detach().numpy(), manual_add_k_val.detach().numpy(), atol=1e-5)
+        self.assertTrue(torch.allclose(add_k_val, manual_add_k_val, atol=1e-5))