Updated Linear Truncated Kernel (#200)

Summary:
Pull Request resolved: #200

The previous linear truncated kernel will not optimize over the lengthscale parameters. Here we initialize the Matern Kernels in the initialization function of Linear Truncated Kernel.

In this diff we also add the GP with Linear Truncated Kernel.

Reviewed By: Balandat

Differential Revision: D16069438

fbshipit-source-id: 446ef6717cf2163abf7c821d9b005cba5e2519a6

Author: VilockLi

botorch/models/fidelity/__init__.py

@@ -2,7 +2,7 @@
 
 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
 
-from .gp_regression_fidelity import SingleTaskMultiFidelityGP
+from .gp_regression_fidelity import SingleTaskGPLTKernel, SingleTaskMultiFidelityGP
 
 
-__all__ = ["SingleTaskMultiFidelityGP"]
+__all__ = ["SingleTaskMultiFidelityGP", "SingleTaskGPLTKernel"]

botorch/models/fidelity/gp_regression_fidelity.py

@@ -12,6 +12,9 @@
 from botorch.exceptions import UnsupportedError
 from botorch.models.fidelity_kernels.downsampling_kernel import DownsamplingKernel
 from botorch.models.fidelity_kernels.exponential_decay_kernel import ExpDecayKernel
+from botorch.models.fidelity_kernels.linear_truncated_fidelity import (
+    LinearTruncatedFidelityKernel,
+)
 from gpytorch.kernels.rbf_kernel import RBFKernel
 from gpytorch.kernels.scale_kernel import ScaleKernel
 from gpytorch.likelihoods.likelihood import Likelihood
@@ -103,3 +106,64 @@ def __init__(
         )
         super().__init__(train_X=train_X, train_Y=train_Y, covar_module=covar_module)
         self.to(train_X)
+
+
+class SingleTaskGPLTKernel(SingleTaskGP):
+    r"""A single task multi-fidelity GP model wiht Linear Truncated kernel.
+
+    A sub-class of SingleTaskGP model. By default the last two dimensions of train_X
+    are the fidelity parameters: training iterations, training data points.
+
+    Args:
+        train_X: A `n x (d + s)` or `batch_shape x n x (d + s) ` (batch mode) tensor
+            of training features, s is the dimension of the fidelity parameters.
+        train_Y: A `n x (o)` or `batch_shape x n x (o)` (batch mode) tensor of
+            training observations.
+        dimension: The dimension of `x`.
+        nu: The smoothness parameter fo Matern kernel: either 1/2, 3/2, or 5/2.
+            Default: '2.5'
+        train_iteration_fidelity: An indicator of whether we have the training
+            iteration fidelity variable.
+        train_data_fidelity: An indicator of whether we have the downsampling
+            fidelity variable. If train_iteration_fidelity and train_data_fidelity
+            are both True, the last and second last columns are treated as the
+            training data points fidelity parameter and training iteration
+            number fidelity parameter respectively. Otherwise the last column of
+            train_X is treated as the fidelity parameter with True indicator.
+            We assume train_X has at least one fidelity parameter.
+        likelihood: A likelihood. If omitted, use a standard
+            GaussianLikelihood with inferred noise level.
+
+    Example:
+        >>> train_X = torch.rand(20, 4)
+        >>> train_Y = train_X.pow(2).sum(dim=-1)
+        >>> model = SingleTaskGPLTKernel(train_X, train_Y)
+    """
+
+    def __init__(
+        self,
+        train_X: Tensor,
+        train_Y: Tensor,
+        nu: float = 2.5,
+        train_iteration_fidelity: bool = True,
+        train_data_fidelity: bool = True,
+        likelihood: Optional[Likelihood] = None,
+    ) -> None:
+        if not train_iteration_fidelity and not train_data_fidelity:
+            raise UnsupportedError("You should have at least one fidelity parameter.")
+        train_X, train_Y, _ = self._set_dimensions(train_X=train_X, train_Y=train_Y)
+        kernel = LinearTruncatedFidelityKernel(
+            nu=nu,
+            dimension=train_X.shape[-1],
+            train_iteration_fidelity=train_iteration_fidelity,
+            train_data_fidelity=train_data_fidelity,
+            batch_shape=self._aug_batch_shape,
+            power_prior=GammaPrior(3.0, 3.0),
+        )
+        covar_module = ScaleKernel(
+            kernel,
+            batch_shape=self._aug_batch_shape,
+            outputscale_prior=GammaPrior(2.0, 0.15),
+        )
+        super().__init__(train_X=train_X, train_Y=train_Y, covar_module=covar_module)
+        self.to(train_X)

botorch/models/fidelity_kernels/__init__.py

@@ -4,6 +4,7 @@
 
 from .downsampling_kernel import DownsamplingKernel
 from .exponential_decay_kernel import ExpDecayKernel
+from .linear_truncated_fidelity import LinearTruncatedFidelityKernel
 
 
-__all__ = ["ExpDecayKernel", "DownsamplingKernel"]
+__all__ = ["ExpDecayKernel", "DownsamplingKernel", "LinearTruncatedFidelityKernel"]

botorch/models/fidelity_kernels/linear_truncated_fidelity.py

@@ -44,6 +44,9 @@ class LinearTruncatedFidelityKernel(Kernel):
         We assume the last two dimensions of input `x` are the fidelity parameters.
 
     Args:
+        :attr:`dimension` (int):
+            The dimension of `x`. This is not needed if active_dims is specified.
+            Default: `3`
         :attr:`nu` (float):
             The smoothness parameter fo Matern kernel: either 1/2, 3/2, or 5/2.
             Default: '2.5'
@@ -73,6 +76,19 @@ class LinearTruncatedFidelityKernel(Kernel):
         :attr:`power_constraint` (Constraint, optional):
             Set this if you want to apply a constraint to the power parameter
             polynomial kernel. Default: `Positive`
+        :attr:`train_iteration_fidelity` (bool):
+            Set this to True if your data contains iteration fidelity parameter.
+            Default: 'True'
+        :attr:`train_data_fidelity` (bool):
+            Set this to True if your data contains training data fidelity parameter.
+            Default: 'True'
+        :attr: `covar_module_1` (Kernel):
+            Set this if you want a different kernel for the unbiased part.
+            Default: 'MaternKernel'
+        :attr: `covar_module_2` (Kernel):
+            Set this if you want a different kernel for the biased part.
+            Default: 'MaternKernel'
+
 
     Attributes:
         :attr:`lengthscale` (Tensor):
@@ -92,6 +108,7 @@ class LinearTruncatedFidelityKernel(Kernel):
 
     def __init__(
         self,
+        dimension: int = 3,
         nu: float = 2.5,
         train_iteration_fidelity: bool = True,
         train_data_fidelity: bool = True,
@@ -100,32 +117,29 @@ def __init__(
         power_constraint: Optional[Interval] = None,
         lengthscale_2_prior: Optional[Prior] = None,
         lengthscale_2_constraint: Optional[Interval] = None,
+        lengthscale_constraint: Optional[Interval] = None,
+        covar_module_1: Optional[Kernel] = None,
+        covar_module_2: Optional[Kernel] = None,
         **kwargs: Any,
     ):
         if not train_iteration_fidelity and not train_data_fidelity:
             raise UnsupportedError("You should have at least one fidelity parameter.")
         if nu not in {0.5, 1.5, 2.5}:
             raise ValueError("nu expected to be 0.5, 1.5, or 2.5")
-        super().__init__(has_lengthscale=True, **kwargs)
+        super().__init__(**kwargs)
         self.train_iteration_fidelity = train_iteration_fidelity
         self.train_data_fidelity = train_data_fidelity
         if power_constraint is None:
             power_constraint = Positive()
 
         if lengthscale_prior is None:
-            self.lengthscale_prior = GammaPrior(1.1, 1 / 20)
-        else:
-            self.lengthscale_prior = lengthscale_prior
+            lengthscale_prior = GammaPrior(3, 6)
 
         if lengthscale_2_prior is None:
-            self.lengthscale_2_prior = GammaPrior(5, 1 / 20)
-        else:
-            self.register_prior(
-                "lengthscale_2_prior",
-                lengthscale_2_prior,
-                lambda: self.lengthscale_2,
-                lambda v: self._set_lengthscale_2(v),
-            )
+            lengthscale_2_prior = GammaPrior(6, 2)
+
+        if lengthscale_constraint is None:
+            lengthscale_constraint = Positive()
 
         if lengthscale_2_constraint is None:
             lengthscale_2_constraint = Positive()
@@ -134,10 +148,6 @@ def __init__(
             name="raw_power",
             parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1)),
         )
-        self.register_parameter(
-            name="raw_lengthscale_2",
-            parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1)),
-        )
 
         if power_prior is not None:
             self.register_prior(
@@ -146,10 +156,35 @@ def __init__(
                 lambda: self.power,
                 lambda v: self._set_power(v),
             )
-        self.nu = nu
-        self.register_constraint("raw_lengthscale_2", lengthscale_2_constraint)
         self.register_constraint("raw_power", power_constraint)
 
+        m = self.train_iteration_fidelity + self.train_data_fidelity
+
+        if self.active_dims is not None:
+            dimension = len(self.active_dims)
+
+        if covar_module_1 is None:
+            self.covar_module_1 = MaternKernel(
+                nu=nu,
+                batch_shape=self.batch_shape,
+                lengthscale_prior=lengthscale_prior,
+                ard_num_dims=dimension - m,
+                lengthscale_constraint=lengthscale_constraint,
+            )
+        else:
+            self.covar_module_1 = covar_module_1
+
+        if covar_module_2 is None:
+            self.covar_module_2 = MaternKernel(
+                nu=nu,
+                batch_shape=self.batch_shape,
+                lengthscale_prior=lengthscale_2_prior,
+                ard_num_dims=dimension - m,
+                lengthscale_constraint=lengthscale_2_constraint,
+            )
+        else:
+            self.covar_module_2 = covar_module_2
+
     @property
     def power(self) -> torch.Tensor:
         return self.raw_power_constraint.transform(self.raw_power)
@@ -163,46 +198,20 @@ def _set_power(self, value: torch.Tensor) -> None:
             value = torch.as_tensor(value).to(self.raw_power)
         self.initialize(raw_power=self.raw_power_constraint.inverse_transform(value))
 
-    @property
-    def lengthscale_2(self) -> torch.Tensor:
-        return self.raw_lengthscale_2_constraint.transform(self.raw_lengthscale_2)
-
-    @lengthscale_2.setter
-    def lengthscale_2(self, value: torch.Tensor) -> None:
-        self._set_lengthscale_2(value)
-
-    def _set_lengthscale_2(self, value: torch.Tensor) -> None:
-        if not torch.is_tensor(value):
-            value = torch.as_tensor(value).to(self.raw_lengthscale_2)
-        self.initialize(
-            raw_lengthscale_2=self.raw_lengthscale_2_constraint.inverse_transform(value)
-        )
-
     def forward(self, x1: torch.Tensor, x2: torch.Tensor, **params) -> torch.Tensor:
-        m = self.train_iteration_fidelity + self.train_data_fidelity
         power = self.power.view(*self.batch_shape, 1, 1)
-        active_dimsM = list(range(x1.size()[-1] - m))
-        covar_module_1 = MaternKernel(
-            nu=self.nu,
-            batch_shape=self.batch_shape,
-            lengthscale_prior=self.lengthscale_prior,
-            active_dims=active_dimsM,
-            ard_num_dims=x1.shape[-1] - m,
-        )
-        covar_module_2 = MaternKernel(
-            nu=self.nu,
-            batch_shape=self.batch_shape,
-            lengthscale_prior=self.lengthscale_2_prior,
-            active_dims=active_dimsM,
-            ard_num_dims=x1.shape[-1] - m,
-        )
-        covar_0 = covar_module_1(x1, x2)
+
+        m = self.train_iteration_fidelity + self.train_data_fidelity
+        active_dimsM = list(range(x1.shape[-1] - m))
+        x1_ = x1[..., active_dimsM]
+        x2_ = x2[..., active_dimsM]
+        covar_0 = self.covar_module_1(x1_, x2_)
+        covar_1 = self.covar_module_2(x1_, x2_)
         x11_ = x1[..., -1].unsqueeze(-1)
         x21t_ = x2[..., -1].unsqueeze(-1).transpose(-1, -2)
-        covar_1 = covar_module_2(x1, x2)
         if self.train_iteration_fidelity and self.train_data_fidelity:
-            covar_2 = covar_module_2(x1, x2)
-            covar_3 = covar_module_2(x1, x2)
+            covar_2 = self.covar_module_2(x1_, x2_)
+            covar_3 = self.covar_module_2(x1_, x2_)
             x12_ = x1[..., -2].unsqueeze(-1)
             x22t_ = x2[..., -2].unsqueeze(-1).transpose(-1, -2)
             res = (

test/models/fidelity/test_gp_regression_fidelity.py

@@ -8,7 +8,10 @@
 import torch
 from botorch import fit_gpytorch_model
 from botorch.exceptions import UnsupportedError
-from botorch.models.fidelity.gp_regression_fidelity import SingleTaskMultiFidelityGP
+from botorch.models.fidelity.gp_regression_fidelity import (
+    SingleTaskGPLTKernel,
+    SingleTaskMultiFidelityGP,
+)
 from botorch.models.gp_regression import FixedNoiseGP
 from botorch.posteriors import GPyTorchPosterior
 from botorch.sampling import SobolQMCNormalSampler
@@ -45,6 +48,9 @@ def _get_random_data_with_fidelity(
 
 
 class TestSingleTaskGPFidelity(unittest.TestCase):
+    def _get_model(self):
+        return SingleTaskMultiFidelityGP
+
     def _get_model_and_data(
         self,
         train_iteration_fidelity,
@@ -66,14 +72,16 @@ def _get_model_and_data(
             "train_iteration_fidelity": train_iteration_fidelity,
             "train_data_fidelity": train_data_fidelity,
         }
-        model = SingleTaskMultiFidelityGP(**model_kwargs)
+        gp_model = self._get_model()
+        model = gp_model(**model_kwargs)
         return model, model_kwargs
 
     def test_exception_message(self, cuda=False):
         train_X = torch.rand(20, 4, device=torch.device("cuda" if cuda else "cpu"))
         train_Y = train_X.pow(2).sum(dim=-1)
+        gp_model = self._get_model()
         with self.assertRaises(UnsupportedError):
-            SingleTaskMultiFidelityGP(
+            gp_model(
                 train_X,
                 train_Y,
                 train_iteration_fidelity=False,
@@ -331,3 +339,8 @@ def test_fantasize(self, cuda=False):
     def test_fantasize_cuda(self):
         if torch.cuda.is_available():
             self.test_fantasize(cuda=True)
+
+
+class TestSingleTaskGPLTKernel(TestSingleTaskGPFidelity):
+    def _get_model(self):
+        return SingleTaskGPLTKernel