Add input scaling check to standard models (#267)

Summary:
Introduces a `settings.validate_input_scaling` flag that when active results in the input data being checked for normalization/standardization.

In this PR this check is on by default, we may want to make it optional if this is too much of a hassle in practice. Right now as used in the models this by default emits warnings and raises errors only on `NaN` values or negative variances.

Addresses #208
Pull Request resolved: #267

Reviewed By: sdaulton

Differential Revision: D17398454

Pulled By: Balandat

fbshipit-source-id: ddb819a6612971a5464218aa1eca54ee281e1366

Author: Balandat

botorch/models/gp_regression.py

@@ -29,6 +29,7 @@
 
 from ..sampling.samplers import MCSampler
 from .gpytorch import BatchedMultiOutputGPyTorchModel
+from .utils import validate_input_scaling
 
 
 MIN_INFERRED_NOISE_LEVEL = 1e-4
@@ -75,6 +76,7 @@ def __init__(
             >>> train_Y = torch.sin(train_X).sum(dim=1, keepdim=True)
             >>> model = SingleTaskGP(train_X, train_Y)
         """
+        validate_input_scaling(train_X=train_X, train_Y=train_Y)
         self._validate_tensor_args(X=train_X, Y=train_Y)
         self._set_dimensions(train_X=train_X, train_Y=train_Y)
         train_X, train_Y, _ = self._transform_tensor_args(X=train_X, Y=train_Y)
@@ -143,6 +145,7 @@ def __init__(self, train_X: Tensor, train_Y: Tensor, train_Yvar: Tensor) -> None
             >>> train_Yvar = torch.full_like(train_Y, 0.2)
             >>> model = FixedNoiseGP(train_X, train_Y, train_Yvar)
         """
+        validate_input_scaling(train_X=train_X, train_Y=train_Y, train_Yvar=train_Yvar)
         self._validate_tensor_args(X=train_X, Y=train_Y, Yvar=train_Yvar)
         self._set_dimensions(train_X=train_X, train_Y=train_Y)
         train_X, train_Y, train_Yvar = self._transform_tensor_args(
@@ -238,6 +241,7 @@ def __init__(self, train_X: Tensor, train_Y: Tensor, train_Yvar: Tensor) -> None
             >>> train_Yvar = 0.1 + se * torch.rand_like(train_Y)
             >>> model = HeteroskedasticSingleTaskGP(train_X, train_Y, train_Yvar)
         """
+        validate_input_scaling(train_X=train_X, train_Y=train_Y, train_Yvar=train_Yvar)
         self._validate_tensor_args(X=train_X, Y=train_Y, Yvar=train_Yvar)
         self._set_dimensions(train_X=train_X, train_Y=train_Y)
         noise_likelihood = GaussianLikelihood(

botorch/models/multitask.py

@@ -70,6 +70,7 @@ def __init__(
             >>> train_Y = torch.cat(f1(X1), f2(X2))
             >>> model = MultiTaskGP(train_X, train_Y, task_feature=-1)
         """
+        # TODO: Validate input normalization/scaling
         if train_X.ndimension() != 2:
             # Currently, batch mode MTGPs are blocked upstream in GPyTorch
             raise ValueError(f"Unsupported shape {train_X.shape} for train_X.")

botorch/models/utils.py

@@ -13,6 +13,7 @@
 from gpytorch.utils.broadcasting import _mul_broadcast_shape
 from torch import Tensor
 
+from .. import settings
 from ..exceptions import InputDataError, InputDataWarning
 
 
@@ -179,3 +180,42 @@ def check_standardization(
             if raise_on_fail:
                 raise InputDataError(msg)
             warnings.warn(msg, InputDataWarning)
+
+
+def validate_input_scaling(
+    train_X: Tensor,
+    train_Y: Tensor,
+    train_Yvar: Optional[Tensor] = None,
+    raise_on_fail: bool = False,
+) -> None:
+    r"""Helper function to validate input data to models.
+
+    Args:
+        train_X: A `n x d` or `batch_shape x n x d` (batch mode) tensor of
+            training features.
+        train_Y: A `n x m` or `batch_shape x n x m` (batch mode) tensor of
+            training observations.
+        train_Yvar: A `batch_shape x n x m` or `batch_shape x n x m` (batch mode)
+            tensor of observed measurement noise.
+        raise_on_fail: If True, raise an error instead of emitting a warning
+            (only for normalization/standardization checks, an error is always
+            raised if NaN values are present).
+
+    This function is typically called inside the constructor of standard BoTorch
+    models. It validates the following:
+    (i) none of the inputs contain NaN values
+    (ii) the training data (`train_X`) is normalized to the unit cube
+    (iii) the training targets (`train_Y`) are standardized (zero mean, unit var)
+    No checks (other than the NaN check) are performed for observed variances
+    (`train_Yvar`) at this point.
+    """
+    if settings.validate_input_scaling.off():
+        return
+    check_no_nans(train_X)
+    check_no_nans(train_Y)
+    if train_Yvar is not None:
+        check_no_nans(train_Yvar)
+        if torch.any(train_Yvar < 0):
+            raise InputDataError("Input data contains negative variances.")
+    check_min_max_scaling(X=train_X, raise_on_fail=raise_on_fail)
+    check_standardization(Y=train_Y, raise_on_fail=raise_on_fail)

botorch/settings.py

@@ -74,3 +74,18 @@ class debug(_Flag):
     def _set_state(cls, state: bool) -> None:
         cls._state = state
         suppress_botorch_warnings(suppress=not cls._state)
+
+
+class validate_input_scaling(_Flag):
+    r"""Flag for validating input normalization/standardization.
+
+    When set to `True`, standard botorch models will validate (up to reasonable
+    tolerance) that
+    (i) none of the inputs contain NaN values
+    (ii) the training data (`train_X`) is normalized to the unit cube
+    (iii) the training targets (`train_Y`) are standardized (zero mean, unit var)
+    No checks (other than the NaN check) are performed for observed variances
+    (`train_Y_var`) at this point.
+    """
+
+    _state: bool = True

test/models/test_utils.py

@@ -13,6 +13,7 @@
     check_no_nans,
     check_standardization,
     multioutput_to_batch_mode_transform,
+    validate_input_scaling,
 )
 from botorch.utils.testing import BotorchTestCase
 
@@ -156,3 +157,47 @@ def test_check_standardization(self):
                 self.assertTrue(any("not standardized" in str(w.message) for w in ws))
             with self.assertRaises(InputDataError):
                 check_standardization(Y=Yst * 2, raise_on_fail=True)
+
+    def test_validate_input_scaling(self):
+        train_X = 2 + torch.rand(3, 4, 3)
+        train_Y = torch.randn(3, 4, 2)
+        # check that nothing is being checked
+        with settings.validate_input_scaling(False), settings.debug(True):
+            with warnings.catch_warnings(record=True) as ws:
+                validate_input_scaling(train_X=train_X, train_Y=train_Y)
+                self.assertFalse(
+                    any(issubclass(w.category, InputDataWarning) for w in ws)
+                )
+        # check that warnings are being issued
+        with settings.debug(True), warnings.catch_warnings(record=True) as ws:
+            validate_input_scaling(train_X=train_X, train_Y=train_Y)
+            self.assertTrue(any(issubclass(w.category, InputDataWarning) for w in ws))
+        # check that errors are raised when requested
+        with settings.debug(True):
+            with self.assertRaises(InputDataError):
+                validate_input_scaling(
+                    train_X=train_X, train_Y=train_Y, raise_on_fail=True
+                )
+        # check that no errors are being raised if everything is standardized
+        train_X_min = train_X.min(dim=-1, keepdim=True)[0]
+        train_X_max = train_X.max(dim=-1, keepdim=True)[0]
+        train_X_std = (train_X - train_X_min) / (train_X_max - train_X_min)
+        train_Y_std = (train_Y - train_Y.mean(dim=-2, keepdim=True)) / train_Y.std(
+            dim=-2, keepdim=True
+        )
+        with settings.debug(True), warnings.catch_warnings(record=True) as ws:
+            validate_input_scaling(train_X=train_X_std, train_Y=train_Y_std)
+            self.assertFalse(any(issubclass(w.category, InputDataWarning) for w in ws))
+        # test that negative variances raise an error
+        train_Yvar = torch.rand_like(train_Y_std)
+        train_Yvar[0, 0, 1] = -0.5
+        with settings.debug(True):
+            with self.assertRaises(InputDataError):
+                validate_input_scaling(
+                    train_X=train_X_std, train_Y=train_Y_std, train_Yvar=train_Yvar
+                )
+        # check that NaNs raise errors
+        train_X_std[0, 0, 0] = float("nan")
+        with settings.debug(True):
+            with self.assertRaises(InputDataError):
+                validate_input_scaling(train_X=train_X_std, train_Y=train_Y_std)