Preserving train inputs and targets through transforms (#3044)

Summary:

This PR preserves botorch transforms (specifically outcome_transforms, like Standardize) through state_dict loading. The fix also ensures that train_targets of a Leave-one-out model with outcome transforms will, in the default case, have the same targets as a base model, minus the point left out.


__Longer explanation:__
Transforms, and specifically learnable output transforms like Standardize, will currently:
a. Learn the parameters at initialization of the GP
b. Transform the train_Ys to the normalized space

Then, when we load a state dict, we will:
a. Impose new standardization parameters on already standardized data
b. Potentially make the transforms re-learnable, nullifying the change made by the state dict

This has undesired consequences for cross-validation, as all cross-validated models will effectively have different training data. In essence, _we don't simply leave one point out, but instead we leave one out and re-standardize_. When we have outliers in the data, this will lead to substantially different predictions when the outlier is left out, since the outlier will substantially impact the outcome transform parameters.

Notebook explaining the effect with some plots: N8342965

Reviewed By: Balandat

Differential Revision: D84571407

Author: Carl Hvarfner

botorch/models/gpytorch.py

@@ -17,7 +17,7 @@
 import warnings
 from abc import ABC
 from copy import deepcopy
-from typing import Any, TYPE_CHECKING
+from typing import Any, Mapping, TYPE_CHECKING
 
 import torch
 from botorch.acquisition.objective import PosteriorTransform
@@ -283,6 +283,111 @@ def condition_on_observations(
             ).detach()
         return fantasy_model
 
+    def _extract_targets_and_noise(self) -> tuple[Tensor, Tensor | None]:
+        r"""Extract targets and noise variance in the correct shape.
+
+        Returns a tuple of (Y, Yvar) where Y and Yvar have shape
+        [batch_shape] x n x m, with batch_shape included only if the
+        training data initially contained it.
+        """
+        if self.num_outputs > 1:
+            Y = self.train_targets.transpose(-1, -2)
+            Yvar = None
+            if isinstance(self.likelihood, FixedNoiseGaussianLikelihood):
+                Yvar = self.likelihood.noise_covar.noise.transpose(-1, -2)
+        else:
+            Y = self.train_targets.unsqueeze(-1)
+            Yvar = None
+            if isinstance(self.likelihood, FixedNoiseGaussianLikelihood):
+                Yvar = self.likelihood.noise_covar.noise.unsqueeze(-1)
+        return Y, Yvar
+
+    def _restore_targets_and_noise(
+        self, Y: Tensor, Yvar: Tensor | None, strict: bool
+    ) -> None:
+        r"""Restore targets and noise variance to the model.
+
+        Args:
+            Y: Targets tensor in shape [batch_shape] x n x m.
+            Yvar: Optional noise variance tensor in shape [batch_shape] x n x m.
+            strict: Whether to strictly enforce shape constraints.
+        """
+        if self.num_outputs > 1:
+            Y = Y.transpose(-1, -2)
+            if Yvar is not None and isinstance(
+                self.likelihood, FixedNoiseGaussianLikelihood
+            ):
+                Yvar = Yvar.transpose(-1, -2)
+                self.likelihood.noise_covar.noise = Yvar
+        else:
+            Y = Y.squeeze(-1)
+            if Yvar is not None and isinstance(
+                self.likelihood, FixedNoiseGaussianLikelihood
+            ):
+                Yvar = Yvar.squeeze(-1)
+                self.likelihood.noise_covar.noise = Yvar
+
+        self.set_train_data(targets=Y, strict=strict)
+
+    def load_state_dict(
+        self,
+        state_dict: Mapping[str, Any],
+        strict: bool = True,
+        keep_transforms: bool = True,
+    ) -> None:
+        r"""Load the model state.
+
+        Args:
+            state_dict: A dict containing the state of the model.
+            strict: A boolean indicating whether to strictly enforce that the keys.
+            keep_transforms: A boolean indicating whether to keep the input and outcome
+                transforms. Doing so is useful when loading a model that was trained on
+                a full set of data, and is later loaded with a subset of the data.
+        """
+        if not keep_transforms:
+            super().load_state_dict(state_dict, strict)
+            return
+
+        should_outcome_transform = (
+            hasattr(self, "train_targets")
+            and getattr(self, "outcome_transform", None) is not None
+        )
+
+        with torch.no_grad():
+            untransformed_Y, untransformed_Yvar = self._extract_targets_and_noise()
+            X = self.train_inputs[0]
+
+            if should_outcome_transform:
+                try:
+                    untransformed_Y, untransformed_Yvar = (
+                        self.outcome_transform.untransform(
+                            Y=untransformed_Y,
+                            Yvar=untransformed_Yvar,
+                            X=X,
+                        )
+                    )
+                except NotImplementedError:
+                    warnings.warn(
+                        "Outcome transform does not support untransforming."
+                        "Cannot load the state dict with transforms preserved."
+                        "Setting keep_transforms=False.",
+                        stacklevel=3,
+                    )
+                    super().load_state_dict(state_dict, strict)
+                    return
+
+        super().load_state_dict(state_dict, strict)
+
+        if getattr(self, "input_transform", None) is not None:
+            self.input_transform.eval()
+
+        if should_outcome_transform:
+            self.outcome_transform.eval()
+            retransformed_Y, retransformed_Yvar = self.outcome_transform(
+                Y=untransformed_Y, Yvar=untransformed_Yvar, X=X
+            )
+            self._restore_targets_and_noise(retransformed_Y, retransformed_Yvar, strict)
+
 
 # pyre-fixme[13]: uninitialized attributes _num_outputs, _input_batch_shape,
 # _aug_batch_shape
@@ -803,6 +908,38 @@ class MultiTaskGPyTorchModel(GPyTorchModel, ABC):
     "long-format" multi-task GP in the style of `MultiTaskGP`.
     """
 
+    def _extract_targets_and_noise(self) -> tuple[Tensor, Tensor | None]:
+        r"""Extract targets and noise variance for multi-task models.
+
+        Returns a tuple of (Y, Yvar) where Y and Yvar have shape
+        [batch_shape] x n x m, with batch_shape included only if the
+        training data initially contained it.
+        """
+        Y = self.train_targets.unsqueeze(-1)
+        Yvar = None
+        if isinstance(self.likelihood, FixedNoiseGaussianLikelihood):
+            Yvar = self.likelihood.noise_covar.noise.unsqueeze(-1)
+        return Y, Yvar
+
+    def _restore_targets_and_noise(
+        self, Y: Tensor, Yvar: Tensor | None, strict: bool
+    ) -> None:
+        r"""Restore targets and noise variance for multi-task models.
+
+        Args:
+            Y: Targets tensor in shape [batch_shape] x n x m.
+            Yvar: Optional noise variance tensor in shape [batch_shape] x n x m.
+            strict: Whether to strictly enforce shape constraints.
+        """
+        Y = Y.squeeze(-1)
+        if Yvar is not None and isinstance(
+            self.likelihood, FixedNoiseGaussianLikelihood
+        ):
+            Yvar = Yvar.squeeze(-1)
+            self.likelihood.noise_covar.noise = Yvar
+
+        self.set_train_data(targets=Y, strict=strict)
+
     def _apply_noise(
         self,
         X: Tensor,