Fix loading state dict for SingleTaskVariationalGP/ApproximateGPyTorchModel (meta-pytorch#3251)

Summary:

**Context**: See `https://github.com/meta-pytorch/botorch/issues/3250` .

GPyTorchModel.load_state_dict previously accessed self.train_targets and self.train_inputs unconditionally, which failed for ApproximateGPyTorchModel because its training data lives on its `model` attribute (e.g. `model.model.train_targets`).

**Changes**:
- Extracted two new overridable methods on GPyTorchModel: _untransform_targets() (undo outcome transform and return raw targets) and _retransform_and_set_targets(). Note: `None` typing effectively replaces `hasattr` checks for `train_targets`
- Used these methods in `load_state_dict`
- Overrode both methods on ApproximateGPyTorchModel to read from / write to self.model.train_targets and self.model.train_inputs instead of self.

Differential Revision: D98021112

Author: esantorella

botorch/models/approximate_gp.py

@@ -128,6 +128,41 @@ def __init__(
             self.likelihood = likelihood
         self._desired_num_outputs = num_outputs
 
+    def _untransform_targets(
+        self,
+    ) -> tuple[Tensor, Tensor | None, Tensor] | None:
+        r"""Extract and untransform training targets from the inner model.
+
+        Overrides ``GPyTorchModel._untransform_targets`` because
+        ``ApproximateGPyTorchModel`` stores ``train_targets`` and
+        ``train_inputs`` on ``self.model`` (the inner ``ApproximateGP``),
+        not directly on ``self``.
+        """
+        if not hasattr(self.model, "train_targets"):
+            return None
+        if getattr(self, "outcome_transform", None) is None:
+            return None
+
+        Y = self.model.train_targets.unsqueeze(-1)
+        X = self.model.train_inputs[0]
+        Y, Yvar = self.outcome_transform.untransform(Y=Y, Yvar=None, X=X)
+        return Y, Yvar, X
+
+    def _retransform_and_set_targets(
+        self,
+        Y: Tensor,
+        Yvar: Tensor | None,
+        X: Tensor,
+    ) -> None:
+        r"""Re-apply the outcome transform and store targets on the inner model.
+
+        Overrides ``GPyTorchModel._retransform_and_set_targets`` because
+        targets must be written to ``self.model.train_targets``.
+        """
+        self.outcome_transform.eval()
+        retransformed_Y, _ = self.outcome_transform(Y=Y, Yvar=Yvar, X=X)
+        self.model.train_targets = retransformed_Y.squeeze(-1)
+
     @property
     def num_outputs(self):
         return self._desired_num_outputs

botorch/models/gpytorch.py

@@ -304,15 +304,12 @@ def _extract_targets_and_noise(self) -> tuple[Tensor, Tensor | None]:
             Y, Yvar = extract_targets_and_noise_single_output(self)
         return Y, Yvar
 
-    def _restore_targets_and_noise(
-        self, Y: Tensor, Yvar: Tensor | None, strict: bool
-    ) -> None:
+    def _restore_targets_and_noise(self, Y: Tensor, Yvar: Tensor | None) -> 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)
@@ -321,9 +318,63 @@ def _restore_targets_and_noise(
             ):
                 Yvar = Yvar.transpose(-1, -2)
                 self.likelihood.noise_covar.noise = Yvar
-            self.set_train_data(targets=Y, strict=strict)
+            self.set_train_data(targets=Y, strict=False)
         else:
-            restore_targets_and_noise_single_output(self, Y, Yvar, strict)
+            restore_targets_and_noise_single_output(
+                model=self, Y=Y, Yvar=Yvar, strict=False
+            )
+
+    def _untransform_targets(
+        self,
+    ) -> tuple[Tensor, Tensor | None, Tensor] | None:
+        r"""Extract training targets, undo the outcome transform, and return them.
+
+        Used by ``load_state_dict`` to save the untransformed targets before
+        loading new parameters, so that the outcome transform can be re-applied
+        afterward with the new transform state.
+
+        Subclasses that store training data somewhere other than ``self`` (e.g.
+        ``ApproximateGPyTorchModel`` stores it on ``self.model``) should
+        override this method.
+
+        Returns:
+            A tuple ``(Y, Yvar, X)`` of untransformed targets, noise variance,
+            and training inputs — or ``None`` if the model has no outcome
+            transform.
+        """
+        if getattr(self, "outcome_transform", None) is None:
+            return None
+
+        Y, Yvar = self._extract_targets_and_noise()
+        X = self.train_inputs[0]
+        Y, Yvar = self.outcome_transform.untransform(Y=Y, Yvar=Yvar, X=X)
+        return Y, Yvar, X
+
+    def _retransform_and_set_targets(
+        self,
+        Y: Tensor,
+        Yvar: Tensor | None,
+        X: Tensor,
+    ) -> None:
+        r"""Re-apply the outcome transform to targets and store them.
+
+        Called by ``load_state_dict`` after new parameters have been loaded,
+        to re-transform the training targets under the updated outcome
+        transform.
+
+        Subclasses that store training data somewhere other than ``self``
+        should override this method.
+
+        Args:
+            Y: Untransformed targets, shape ``batch_shape x n x m``.
+            Yvar: Untransformed noise variance, or ``None``.
+            X: Training inputs, shape ``batch_shape x n x d``.
+        """
+        self.outcome_transform.eval()
+        retransformed_Y, retransformed_Yvar = self.outcome_transform(
+            Y=Y, Yvar=Yvar, X=X
+        )
+        self._restore_targets_and_noise(Y=retransformed_Y, Yvar=retransformed_Yvar)
 
     def load_state_dict(
         self,
@@ -353,48 +404,34 @@ def load_state_dict(
             super().load_state_dict(state_dict=state_dict, strict=strict, assign=assign)
             return
 
-        should_outcome_transform = (
-            hasattr(self, "train_targets")
-            and getattr(self, "outcome_transform", None) is not None
-        )
-
+        # Before loading new parameters, untransform the current training
+        # targets so they can be re-transformed under the new outcome
+        # transform state. Returns None if no outcome transform or no
+        # training data.
         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.",
-                        BotorchWarning,
-                        stacklevel=3,
-                    )
-                    super().load_state_dict(
-                        state_dict=state_dict, strict=strict, assign=assign
-                    )
-                    return
+            try:
+                untransformed = self._untransform_targets()
+            except NotImplementedError:
+                warnings.warn(
+                    "Outcome transform does not support untransforming. "
+                    "Cannot load the state dict with transforms preserved. "
+                    "Setting keep_transforms=False.",
+                    BotorchWarning,
+                    stacklevel=3,
+                )
+                super().load_state_dict(
+                    state_dict=state_dict, strict=strict, assign=assign
+                )
+                return
 
         super().load_state_dict(state_dict=state_dict, strict=strict, assign=assign)
 
         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)
+        if untransformed is not None:
+            Y, Yvar, X = untransformed
+            self._retransform_and_set_targets(Y=Y, Yvar=Yvar, X=X)
 
 
 # pyre-fixme[13]: uninitialized attributes _num_outputs, _input_batch_shape,
@@ -935,17 +972,16 @@ def _extract_targets_and_noise(self) -> tuple[Tensor, Tensor | None]:
         """
         return extract_targets_and_noise_single_output(self)
 
-    def _restore_targets_and_noise(
-        self, Y: Tensor, Yvar: Tensor | None, strict: bool
-    ) -> None:
+    def _restore_targets_and_noise(self, Y: Tensor, Yvar: Tensor | None) -> 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.
         """
-        restore_targets_and_noise_single_output(self, Y, Yvar, strict)
+        restore_targets_and_noise_single_output(
+            model=self, Y=Y, Yvar=Yvar, strict=False
+        )
 
     def _apply_noise(
         self,

test/models/test_approximate_gp.py

@@ -17,7 +17,7 @@
     SingleTaskVariationalGP,
 )
 from botorch.models.transforms.input import Normalize
-from botorch.models.transforms.outcome import Log
+from botorch.models.transforms.outcome import Log, Standardize
 from botorch.models.utils.inducing_point_allocators import (
     GreedyImprovementReduction,
     GreedyVarianceReduction,
@@ -63,21 +63,161 @@ def test_initialization(self):
         )
         self.assertEqual(model.num_outputs, 2)
 
+    def test_load_state_dict(self) -> None:
+        test_X = torch.rand(5, 1, device=self.device)
+
+        for label, train_Y in [
+            ("with_train_Y", self.train_Y),
+            ("no_train_Y", None),
+        ]:
+            with self.subTest(label=label):
+                model = ApproximateGPyTorchModel(
+                    train_X=self.train_X,
+                    train_Y=train_Y,
+                    likelihood=BetaLikelihood().to(self.device),
+                )
+                state_dict = model.state_dict()
+
+                restored = ApproximateGPyTorchModel(
+                    train_X=self.train_X,
+                    train_Y=train_Y,
+                    likelihood=BetaLikelihood().to(self.device),
+                )
+                restored.load_state_dict(state_dict=state_dict)
+                restored_state = restored.state_dict()
+                self.assertEqual(set(state_dict.keys()), set(restored_state.keys()))
+                for key in state_dict:
+                    self.assertTrue(
+                        torch.equal(state_dict[key], restored_state[key]),
+                        msg=f"Mismatch for key {key}",
+                    )
+
+                model.eval()
+                restored.eval()
+                torch.manual_seed(0)
+                orig_posterior = model.posterior(test_X)
+                torch.manual_seed(0)
+                restored_posterior = restored.posterior(test_X)
+                self.assertAllClose(orig_posterior.mean, restored_posterior.mean)
+                self.assertAllClose(
+                    orig_posterior.variance, restored_posterior.variance
+                )
+
 
 class TestSingleTaskVariationalGP(BotorchTestCase):
     def setUp(self):
         super().setUp()
-        train_X = torch.rand(10, 1, device=self.device)
-        train_y = torch.sin(train_X) + torch.randn_like(train_X) * 0.2
+        self.train_X = torch.rand(10, 1, device=self.device)
+        self.train_y = torch.sin(self.train_X) + torch.randn_like(self.train_X) * 0.2
 
         self.model = SingleTaskVariationalGP(
-            train_X=train_X, likelihood=GaussianLikelihood()
+            train_X=self.train_X,
+            train_Y=self.train_y,
+            likelihood=GaussianLikelihood(),
+            outcome_transform=Standardize(m=1),
         ).to(self.device)
 
         mll = VariationalELBO(self.model.likelihood, self.model.model, num_data=10)
-        loss = -mll(self.model.likelihood(self.model(train_X)), train_y).sum()
+        loss = -mll(self.model.likelihood(self.model(self.train_X)), self.train_y).sum()
         loss.backward()
 
+    def test_load_state_dict(self) -> None:
+        test_X = torch.rand(5, 1, device=self.device)
+
+        for label, train_Y in [
+            ("with_train_Y", self.train_y),
+            ("no_train_Y", None),
+        ]:
+            with self.subTest(label=label):
+                model = SingleTaskVariationalGP(
+                    train_X=self.train_X,
+                    train_Y=train_Y,
+                    likelihood=BetaLikelihood(),
+                ).to(self.device)
+                state_dict = model.state_dict()
+
+                restored = SingleTaskVariationalGP(
+                    train_X=self.train_X,
+                    train_Y=train_Y,
+                    likelihood=BetaLikelihood(),
+                ).to(self.device)
+                restored.load_state_dict(state_dict=state_dict)
+                restored_state = restored.state_dict()
+                self.assertEqual(set(state_dict.keys()), set(restored_state.keys()))
+                for key in state_dict:
+                    self.assertTrue(
+                        torch.equal(state_dict[key], restored_state[key]),
+                        msg=f"Mismatch for key {key}",
+                    )
+
+                # Posterior numerical identity. manual_seed is needed because
+                # CholeskyVariationalDistribution.initialize_variational_distribution
+                # adds random noise on the first posterior call.
+                model.eval()
+                restored.eval()
+                torch.manual_seed(0)
+                orig_posterior = model.posterior(test_X)
+                torch.manual_seed(0)
+                restored_posterior = restored.posterior(test_X)
+                self.assertAllClose(orig_posterior.mean, restored_posterior.mean)
+                self.assertAllClose(
+                    orig_posterior.variance, restored_posterior.variance
+                )
+
+        # Test keep_transforms with different training data (CV-style).
+        # The restored model is built with one fewer data point, so its
+        # Standardize means/stdvs differ from the source model's.
+        with self.subTest("keep_transforms"):
+            model = SingleTaskVariationalGP(
+                train_X=self.train_X,
+                train_Y=self.train_y,
+                outcome_transform=Standardize(m=1),
+                input_transform=Normalize(d=1),
+            ).to(self.device)
+            state_dict = model.state_dict()
+            original_train_targets = model.model.train_targets.clone()
+
+            cv_X = self.train_X[:-1]
+            cv_Y = self.train_y[:-1]
+
+            for keep_transforms in [True, False]:
+                with self.subTest(keep_transforms=keep_transforms):
+                    restored = SingleTaskVariationalGP(
+                        train_X=cv_X,
+                        train_Y=cv_Y,
+                        outcome_transform=Standardize(m=1),
+                        input_transform=Normalize(d=1),
+                    ).to(self.device)
+                    restored.load_state_dict(
+                        state_dict=state_dict, keep_transforms=keep_transforms
+                    )
+
+                    if keep_transforms:
+                        # Transform params are loaded from state_dict, and
+                        # train_targets are re-standardized under the loaded
+                        # transform, so they match the original (minus the
+                        # dropped point).
+                        self.assertAllClose(
+                            restored.model.train_targets,
+                            original_train_targets[..., :-1],
+                        )
+                        self.assertTrue(
+                            torch.equal(
+                                restored.outcome_transform.means,
+                                state_dict["outcome_transform.means"],
+                            )
+                        )
+                    else:
+                        # Transform params are loaded but train_targets are
+                        # NOT re-standardized, so they still reflect the
+                        # cv_Y-based standardization and won't match.
+                        self.assertFalse(
+                            torch.allclose(
+                                restored.model.train_targets,
+                                original_train_targets[..., :-1],
+                            )
+                        )
+
     def test_posterior(self):
         # basic test of checking that the posterior works as intended
         test_x = torch.rand(30, 1, device=self.device)