Repository: cornellius-gp/gpytorch
Pull request #330 (2018-10-18 23:07:50)

feature: add rational quadratic kernel
Implement rational quadratic kernel

add `log_alpha` hyperparameter, otherwise implement similarly to `RBFKernel`

gpleiss: Can you add a unit test? See the other kernels for examples.

Patch:
diff --git a/gpytorch/kernels/__init__.py b/gpytorch/kernels/__init__.py
index 26ece727..76c54e9f 100644
--- a/gpytorch/kernels/__init__.py
+++ b/gpytorch/kernels/__init__.py
@@ -20,6 +20,7 @@ from .polynomial_kernel_grad import PolynomialKernelGrad
 from .product_structure_kernel import ProductStructureKernel
 from .rbf_kernel import RBFKernel
 from .rbf_kernel_grad import RBFKernelGrad
+from .rq_kernel import RQKernel
 from .scale_kernel import ScaleKernel
 from .spectral_mixture_kernel import SpectralMixtureKernel
 from .white_noise_kernel import WhiteNoiseKernel
@@ -49,6 +50,7 @@ __all__ = [
     "ProductStructureKernel",
     "RBFKernel",
     "RBFKernelGrad",
+    "RQKernel",
     "ScaleKernel",
     "SpectralMixtureKernel",
     "WhiteNoiseKernel",
diff --git a/gpytorch/kernels/rq_kernel.py b/gpytorch/kernels/rq_kernel.py
new file mode 100644
index 00000000..5187b7ca
--- /dev/null
+++ b/gpytorch/kernels/rq_kernel.py
@@ -0,0 +1,89 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import torch
+from .kernel import Kernel
+from ..constraints import Positive
+
+
+class RQKernel(Kernel):
+    r"""
+    Computes a covariance matrix based on the rational quadratic kernel
+    between inputs :math:`\mathbf{x_1}` and :math:`\mathbf{x_2}`:
+
+    .. math::
+
+       \begin{equation*}
+          k_{\text{RQ}}(\mathbf{x_1}, \mathbf{x_2}) =  \left(1 + \frac{1}{2\alpha}
+          (\mathbf{x_1} - \mathbf{x_2})^\top \Theta^{-2} (\mathbf{x_1} - \mathbf{x_2}) \right)^{-\alpha}
+       \end{equation*}
+
+    where :math:`\Theta` is a :attr:`lengthscale` parameter, and :math:`\alpha` is the
+    rational quadratic relative weighting parameter.
+    See :class:`gpytorch.kernels.Kernel` for descriptions of the lengthscale options.
+
+    .. note::
+
+        This kernel does not have an `outputscale` parameter. To add a scaling parameter,
+        decorate this kernel with a :class:`gpytorch.kernels.ScaleKernel`.
+
+    Args:
+        :attr:`ard_num_dims` (int, optional):
+            Set this if you want a separate lengthscale for each
+            input dimension. It should be `d` if :attr:`x1` is a `n x d` matrix. Default: `None`
+        :attr:`batch_shape` (torch.Size, optional):
+            Set this if you want a separate lengthscale for each
+            batch of input data. It should be `b` if :attr:`x1` is a `b x n x d` tensor. Default: `torch.Size([])`.
+        :attr:`active_dims` (tuple of ints, optional):
+            Set this if you want to compute the covariance of only a few input dimensions. The ints
+            corresponds to the indices of the dimensions. Default: `None`.
+        :attr:`lengthscale_prior` (Prior, optional):
+            Set this if you want to apply a prior to the lengthscale parameter.  Default: `None`.
+        :attr:`lengthscale_constraint` (Constraint, optional):
+            Set this if you want to apply a constraint to the lengthscale parameter. Default: `Positive`.
+        :attr:`alpha_constraint` (Constraint, optional):
+            Set this if you want to apply a constraint to the alpha parameter. Default: `Positive`.
+        :attr:`eps` (float):
+            The minimum value that the lengthscale can take (prevents divide by zero errors). Default: `1e-6`.
+
+    Attributes:
+        :attr:`lengthscale` (Tensor):
+            The lengthscale parameter. Size/shape of parameter depends on the
+            :attr:`ard_num_dims` and :attr:`batch_shape` arguments.
+        :attr:`alpha` (Tensor):
+            The rational quadratic relative weighting parameter. Size/shape of parameter depends
+            on the :attr:`batch_shape` argument
+    """
+
+    def __init__(self, alpha_constraint=None, **kwargs):
+        super(RQKernel, self).__init__(has_lengthscale=True, **kwargs)
+        self.register_parameter(name="raw_alpha", parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1)))
+        if alpha_constraint is None:
+            alpha_constraint = Positive()
+
+        self.register_constraint("raw_alpha", alpha_constraint)
+
+    def forward(self, x1, x2, diag=False, **params):
+        def postprocess_rq(dist):
+            alpha = self.alpha
+            for _ in range(1, len(dist.shape) - len(self.batch_shape)):
+                alpha = alpha.unsqueeze(-1)
+            return (1 + dist.div(2 * alpha)).pow(-alpha)
+
+        x1_ = x1.div(self.lengthscale)
+        x2_ = x2.div(self.lengthscale)
+        return self.covar_dist(
+            x1_, x2_, square_dist=True, diag=diag, dist_postprocess_func=postprocess_rq, postprocess=True, **params
+        )
+
+    @property
+    def alpha(self):
+        return self.raw_alpha_constraint.transform(self.raw_alpha)
+
+    @alpha.setter
+    def alpha(self, value):
+        if not torch.is_tensor(value):
+            value = torch.as_tensor(value).to(self.raw_lengthscale)
+        self.initialize(raw_alpha=self.raw_alpha_constraint.inverse_transform(value))


Test patch:
diff --git a/test/kernels/test_rq_kernel.py b/test/kernels/test_rq_kernel.py
new file mode 100644
index 00000000..49511ea9
--- /dev/null
+++ b/test/kernels/test_rq_kernel.py
@@ -0,0 +1,225 @@
+#!/usr/bin/env python3
+
+import math
+import torch
+import unittest
+from gpytorch.kernels import RQKernel
+from gpytorch.test.base_kernel_test_case import BaseKernelTestCase
+
+
+class TestRQKernel(unittest.TestCase, BaseKernelTestCase):
+    def create_kernel_no_ard(self, **kwargs):
+        return RQKernel(**kwargs)
+
+    def create_kernel_ard(self, num_dims, **kwargs):
+        return RQKernel(ard_num_dims=num_dims, **kwargs)
+
+    def test_ard(self):
+        a = torch.tensor([[1, 2], [2, 4]], dtype=torch.float)
+        b = torch.tensor([[1, 3], [0, 4]], dtype=torch.float)
+        lengthscales = torch.tensor([1, 2], dtype=torch.float).view(1, 2)
+
+        kernel = RQKernel(ard_num_dims=2)
+        kernel.initialize(lengthscale=lengthscales)
+        kernel.initialize(alpha=3.0)
+        kernel.eval()
+
+        scaled_a = a.div(lengthscales)
+        scaled_b = b.div(lengthscales)
+        dist = (scaled_a.unsqueeze(-2) - scaled_b.unsqueeze(-3)).pow(2).sum(dim=-1)
+        actual = dist.div_(2 * kernel.alpha).add_(1.0).pow(-kernel.alpha)
+        res = kernel(a, b).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # Diag
+        res = kernel(a, b).diag()
+        actual = actual.diag()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims
+        diff = scaled_a.transpose(-1, -2).unsqueeze(-1) - scaled_b.transpose(-1, -2).unsqueeze(-2)
+        actual = diff.pow(2).div_(2 * kernel.alpha).add_(1.0).pow(-kernel.alpha)
+        res = kernel(a, b, last_dim_is_batch=True).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims and diag
+        res = kernel(a, b, last_dim_is_batch=True).diag()
+        actual = actual.diagonal(dim1=-1, dim2=-2)
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+    def test_ard_batch(self):
+        a = torch.tensor([[[1, 2, 3], [2, 4, 0]], [[-1, 1, 2], [2, 1, 4]]], dtype=torch.float)
+        b = torch.tensor([[[1, 3, 1]], [[2, -1, 0]]], dtype=torch.float).repeat(1, 2, 1)
+        lengthscales = torch.tensor([[[1, 2, 1]]], dtype=torch.float)
+
+        kernel = RQKernel(batch_shape=torch.Size([2]), ard_num_dims=3)
+        kernel.initialize(lengthscale=lengthscales)
+        kernel.initialize(alpha=3.0)
+        kernel.eval()
+
+        scaled_a = a.div(lengthscales)
+        scaled_b = b.div(lengthscales)
+        dist = (scaled_a.unsqueeze(-2) - scaled_b.unsqueeze(-3)).pow(2).sum(dim=-1)
+        actual = dist.div_(2 * kernel.alpha).add_(1.0).pow(-kernel.alpha)
+        res = kernel(a, b).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # diag
+        res = kernel(a, b).diag()
+        actual = actual.diagonal(dim1=-1, dim2=-2)
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # # batch_dims
+        double_batch_a = scaled_a.transpose(-1, -2).unsqueeze(-1)
+        double_batch_b = scaled_b.transpose(-1, -2).unsqueeze(-2)
+        actual = double_batch_a - double_batch_b
+        alpha = kernel.alpha.view(2, 1, 1, 1)
+        actual = actual.pow_(2).div_(2 * alpha).add_(1.0).pow(-alpha)
+        res = kernel(a, b, last_dim_is_batch=True).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims and diag
+        res = kernel(a, b, last_dim_is_batch=True).diag()
+        actual = actual.diagonal(dim1=-2, dim2=-1)
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+    def test_ard_separate_batch(self):
+        a = torch.tensor([[[1, 2, 3], [2, 4, 0]], [[-1, 1, 2], [2, 1, 4]]], dtype=torch.float)
+        b = torch.tensor([[[1, 3, 1]], [[2, -1, 0]]], dtype=torch.float).repeat(1, 2, 1)
+        lengthscales = torch.tensor([[[1, 2, 1]], [[2, 1, 0.5]]], dtype=torch.float)
+
+        kernel = RQKernel(batch_shape=torch.Size([2]), ard_num_dims=3)
+        kernel.initialize(lengthscale=lengthscales)
+        kernel.initialize(alpha=3.0)
+        kernel.eval()
+
+        scaled_a = a.div(lengthscales)
+        scaled_b = b.div(lengthscales)
+        dist = (scaled_a.unsqueeze(-2) - scaled_b.unsqueeze(-3)).pow(2).sum(dim=-1)
+        actual = dist.div_(2 * kernel.alpha).add_(1.0).pow(-kernel.alpha)
+        res = kernel(a, b).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # diag
+        res = kernel(a, b).diag()
+        actual = actual.diagonal(dim1=-1, dim2=-2)
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+    def test_computes_rational_quadratic(self):
+        a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
+        b = torch.tensor([0, 2, 4], dtype=torch.float).view(3, 1)
+        lengthscale = 2
+
+        kernel = RQKernel().initialize(lengthscale=lengthscale)
+        kernel.eval()
+
+        dist = torch.tensor([[16, 4, 0], [4, 0, 4], [64, 36, 16]], dtype=torch.float).div(lengthscale ** 2)
+        actual = dist.div_(2 * kernel.alpha).add_(1.0).pow(-kernel.alpha)
+        res = kernel(a, b).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # diag
+        res = kernel(a, b).diag()
+        actual = actual.diag()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+    def test_computes_rational_quadratic_gradient(self):
+        softplus = torch.nn.functional.softplus
+        a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
+        b = torch.tensor([0, 2, 2], dtype=torch.float).view(3, 1)
+
+        kernel = RQKernel()
+        kernel.initialize(lengthscale=2.0)
+        kernel.initialize(alpha=3.0)
+        kernel.eval()
+
+        raw_lengthscale = torch.tensor(math.log(math.exp(2.0) - 1))
+        raw_lengthscale.requires_grad_()
+        raw_alpha = torch.tensor(math.log(math.exp(3.0) - 1))
+        raw_alpha.requires_grad_()
+        lengthscale, alpha = softplus(raw_lengthscale), softplus(raw_alpha)
+        dist = (a.expand(3, 3) - b.expand(3, 3).transpose(0, 1)).div(lengthscale).pow(2)
+        actual_output = dist.div(2 * alpha).add(1).pow(-alpha)
+        actual_output.backward(gradient=torch.eye(3))
+
+        output = kernel(a, b).evaluate()
+        output.backward(gradient=torch.eye(3))
+
+        res = kernel.raw_lengthscale.grad
+        self.assertLess(torch.norm(res - raw_lengthscale.grad), 1e-5)
+        res = kernel.raw_alpha.grad
+        self.assertLess(torch.norm(res - raw_alpha.grad), 1e-5)
+
+    def test_subset_active_compute_rational_quadratic(self):
+        a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
+        a_p = torch.tensor([1, 2, 3], dtype=torch.float).view(3, 1)
+        a = torch.cat((a, a_p), 1)
+        b = torch.tensor([0, 2, 4], dtype=torch.float).view(3, 1)
+        lengthscale = 2
+
+        kernel = RQKernel(active_dims=[0])
+        kernel.initialize(lengthscale=lengthscale)
+        kernel.initialize(alpha=3.0)
+        kernel.eval()
+
+        actual = torch.tensor([[16, 4, 0], [4, 0, 4], [64, 36, 16]], dtype=torch.float)
+        actual.div_(lengthscale ** 2).div_(2 * kernel.alpha).add_(1).pow_(-kernel.alpha)
+        res = kernel(a, b).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # diag
+        res = kernel(a, b).diag()
+        actual = actual.diag()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+    def test_subset_active_computes_rational_quadratic_gradient(self):
+        softplus = torch.nn.functional.softplus
+        a_1 = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
+        a_p = torch.tensor([1, 2, 3], dtype=torch.float).view(3, 1)
+        a = torch.cat((a_1, a_p), 1)
+        b = torch.tensor([0, 2, 2], dtype=torch.float).view(3, 1)
+
+        kernel = RQKernel(active_dims=[0])
+        kernel.initialize(lengthscale=2.0)
+        kernel.initialize(alpha=3.0)
+        kernel.eval()
+
+        raw_lengthscale = torch.tensor(math.log(math.exp(2.0) - 1))
+        raw_lengthscale.requires_grad_()
+        raw_alpha = torch.tensor(math.log(math.exp(3.0) - 1))
+        raw_alpha.requires_grad_()
+        lengthscale, alpha = softplus(raw_lengthscale), softplus(raw_alpha)
+        dist = (a_1.expand(3, 3) - b.expand(3, 3).transpose(0, 1)).div(lengthscale).pow(2)
+        actual_output = dist.div(2 * alpha).add(1).pow(-alpha)
+        actual_output.backward(gradient=torch.eye(3))
+
+        output = kernel(a, b).evaluate()
+        output.backward(gradient=torch.eye(3))
+
+        res = kernel.raw_lengthscale.grad
+        self.assertLess(torch.norm(res - raw_lengthscale.grad), 1e-5)
+        res = kernel.raw_alpha.grad
+        self.assertLess(torch.norm(res - raw_alpha.grad), 1e-5)
+
+    def test_initialize_lengthscale(self):
+        kernel = RQKernel()
+        kernel.initialize(lengthscale=3.14)
+        actual_value = torch.tensor(3.14).view_as(kernel.lengthscale)
+        self.assertLess(torch.norm(kernel.lengthscale - actual_value), 1e-5)
+
+    def test_initialize_lengthscale_batch(self):
+        kernel = RQKernel(batch_shape=torch.Size([2]))
+        ls_init = torch.tensor([3.14, 4.13])
+        kernel.initialize(lengthscale=ls_init)
+        actual_value = ls_init.view_as(kernel.lengthscale)
+        self.assertLess(torch.norm(kernel.lengthscale - actual_value), 1e-5)
+
+    def test_initialize_alpha(self):
+        kernel = RQKernel()
+        kernel.initialize(alpha=3.0)
+        actual_value = torch.tensor(3.0).view_as(kernel.alpha)
+        self.assertLess(torch.norm(kernel.alpha - actual_value), 1e-5)
+
+
+if __name__ == "__main__":
+    unittest.main()