Add RGPE TL Mode and Identity Kernel

CHANGELOG.md

@@ -15,6 +15,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   permutation groups instead of column groups
 - `ParameterSelectorProtocol.__call__` now declares its input as positional-only
 
+### Changed
+- Transfer-learning convergence benchmarks now build the naive baseline from a
+  `TaskParameter` with `override_transfer_learning_mode="IDENTITY"` instead of dropping
+  the task parameter, keeping it alongside the default transfer-learning variant
+
 ### Fixed
 - `DiscretePermutationInvarianceConstraint` no longer erroneously removes points where
   values of invariant points are degenerate
@@ -26,6 +31,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   constraint filtering
 
 ### Added
+- `RGPESurrogate`, a rank-weighted Gaussian process ensemble for transfer learning,
+  selectable directly or by setting `override_transfer_learning_mode="RGPE"` on a
+  `TaskParameter`
+- `TransferLearningMode.IDENTITY`, a constant unit task kernel that pools all tasks into
+  a single model fit with the marginal log-likelihood, matching naive pooling
+- `IdentityKernel`, a constant unit kernel that acts as an identity element under
+  kernel multiplication
 - `simplex_coefficients` keyword argument to `SubspaceDiscrete.from_simplex` for
   weighted simplex sum constraints
 - `Symmetry` concept for expressing symmetries of the optimization problem, including

baybe/kernels/__init__.py

@@ -5,6 +5,7 @@
 """
 
 from baybe.kernels.basic import (
+    IdentityKernel,
     LinearKernel,
     MaternKernel,
     PeriodicKernel,
@@ -18,6 +19,7 @@
 
 __all__ = [
     "AdditiveKernel",
+    "IdentityKernel",
     "LinearKernel",
     "MaternKernel",
     "PeriodicKernel",

baybe/kernels/basic.py

@@ -1,7 +1,9 @@
 """Collection of basic kernels."""
 
+from __future__ import annotations
+
 import gc
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from attrs import define, field
 from attrs.converters import optional as optional_c
@@ -15,6 +17,31 @@
 from baybe.utils.conversion import fraction_to_float
 from baybe.utils.validation import finite_float
 
+if TYPE_CHECKING:
+    from baybe.searchspace.core import SearchSpace
+
+
+@define(frozen=True)
+class IdentityKernel(BasicKernel):
+    """A constant unit kernel acting as identity under kernel multiplication."""
+
+    @override
+    def _get_dimensions(
+        self, searchspace: SearchSpace
+    ) -> tuple[tuple[int, ...] | None, int | None]:
+        active_dims, _ = super()._get_dimensions(searchspace)
+        return active_dims, None
+
+    @override
+    def to_gpytorch(self, searchspace: SearchSpace):
+        from gpytorch.kernels import ConstantKernel
+
+        active_dims, _ = self._get_dimensions(searchspace)
+        kernel = ConstantKernel(active_dims=active_dims)
+        kernel.constant = kernel.constant.new_ones(kernel.constant.shape)
+        kernel.raw_constant.requires_grad_(False)
+        return kernel
+
 
 @define(frozen=True)
 class LinearKernel(BasicKernel):

baybe/parameters/categorical.py

@@ -98,22 +98,28 @@ class TaskParameter(CategoricalParameter):
     )
     """Optional override for how the task dimension is modeled.
 
-    Only applies to :class:`.GaussianProcessSurrogate`. When ``None``, the surrogate's
-    kernel factory decides how the task dimension is treated. When set, the surrogate
-    attaches the requested task kernel to a task-free base kernel derived from the
-    configured factory.
+    Only applies to :class:`.GaussianProcessSurrogate`. When ``None``, the kernel
+    factory decides. The kernel modes attach the chosen task kernel to the base kernel,
+    while :attr:`~baybe.parameters.enum.TransferLearningMode.RGPE` dispatches to a
+    :class:`~baybe.surrogates.transfer_learning.rgpe.RGPESurrogate` ensemble.
     """
 
     @override
     @cached_property
     def override_kernel(self) -> Kernel | None:
         """The task kernel defined by the transfer learning mode, if any."""
-        from baybe.kernels.basic import IndexKernel, PositiveIndexKernel
+        from baybe.kernels.basic import IdentityKernel, IndexKernel, PositiveIndexKernel
 
         n_tasks, names = len(self.values), (self.name,)
         match mode := self.override_transfer_learning_mode:
             case None:
                 return None
+            case TransferLearningMode.RGPE:
+                # RGPE adds no task kernel; the surrogate dispatches to a dedicated
+                # ensemble instead (see `GaussianProcessSurrogate._fit`).
+                return None
+            case TransferLearningMode.IDENTITY:
+                return IdentityKernel(parameter_names=names)
             case TransferLearningMode.POSITIVE_INDEX_KERNEL:
                 return PositiveIndexKernel(
                     num_tasks=n_tasks, rank=n_tasks, parameter_names=names

baybe/parameters/enum.py

@@ -179,9 +179,22 @@ class SubstanceEncoding(ParameterEncoding):
 class TransferLearningMode(Enum):
     """Transfer learning modes for :class:`.TaskParameter`."""
 
+    IDENTITY = "IDENTITY"
+    """A constant task kernel that leaves the base kernel unchanged.
+
+    The task dimension adds no covariance, so all tasks are pooled into a single model.
+    """
+
     INDEX_KERNEL = "INDEX_KERNEL"
     """:class:`gpytorch.kernels.IndexKernel` for arbitrary correlations."""
 
     POSITIVE_INDEX_KERNEL = "POSITIVE_INDEX_KERNEL"
     """:class:`botorch.models.kernels.positive_index.PositiveIndexKernel` for positive
     correlations."""
+
+    RGPE = "RGPE"
+    """A rank-weighted ensemble of per-task Gaussian processes.
+
+    Adds no task kernel; instead dispatches the surrogate to a dedicated
+    :class:`~baybe.surrogates.transfer_learning.rgpe.RGPESurrogate`.
+    """

baybe/searchspace/core.py

@@ -11,12 +11,13 @@
 import numpy as np
 import numpy.typing as npt
 import pandas as pd
-from attrs import define, field
+from attrs import define, evolve, field
 from typing_extensions import Self, override
 
 from baybe.constraints import validate_constraints
 from baybe.constraints.base import Constraint
 from baybe.exceptions import (
+    IncompatibleSearchSpaceError,
     InfeasibilityError,
     _UnsupportedSearchSpaceAttributeError,
 )
@@ -36,6 +37,7 @@
 from baybe.utils.conversion import to_string
 
 if TYPE_CHECKING:
+    from baybe.parameters.enum import TransferLearningMode
     from baybe.parameters.selectors import ParameterSelectorProtocol
 
 
@@ -580,6 +582,34 @@ def _drop_parameters(self, names: Collection[str], /) -> _ReducedSearchSpace:
 
         return _ReducedSearchSpace(discrete=discrete, continuous=continuous)
 
+    def _with_task_mode(self, mode: TransferLearningMode, /) -> SearchSpace:
+        """Return a copy with the task parameter's transfer-learning mode replaced.
+
+        The mode does not affect encoding, so the computational representation is
+        unchanged.
+
+        Args:
+            mode: The transfer-learning mode to assign to the task parameter.
+
+        Raises:
+            IncompatibleSearchSpaceError: If the space has no task parameter.
+
+        Returns:
+            A copy whose task parameter carries the given mode.
+        """
+        if self._task_parameter is None:
+            raise IncompatibleSearchSpaceError(
+                "Cannot set a transfer-learning mode on a search space that does not "
+                "contain a task parameter."
+            )
+        parameters = tuple(
+            evolve(p, override_transfer_learning_mode=mode)
+            if isinstance(p, TaskParameter)
+            else p
+            for p in self.discrete.parameters
+        )
+        return evolve(self, discrete=evolve(self.discrete, parameters=parameters))
+
 
 @define(slots=False)
 class _ReducedSearchSpace(SearchSpace):

baybe/surrogates/__init__.py

@@ -8,6 +8,7 @@
 from baybe.surrogates.naive import MeanPredictionSurrogate
 from baybe.surrogates.ngboost import NGBoostSurrogate
 from baybe.surrogates.random_forest import RandomForestSurrogate
+from baybe.surrogates.transfer_learning import RGPESurrogate
 
 __all__ = [
     "BayesianLinearSurrogate",
@@ -17,5 +18,6 @@
     "GaussianProcessSurrogate",
     "MeanPredictionSurrogate",
     "NGBoostSurrogate",
+    "RGPESurrogate",
     "RandomForestSurrogate",
 ]

baybe/surrogates/composite.py

@@ -18,8 +18,6 @@
 from baybe.serialization.core import _TYPE_FIELD, add_type
 from baybe.serialization.mixin import SerialMixin
 from baybe.surrogates.base import PosteriorStatistic, SurrogateProtocol
-from baybe.surrogates.gaussian_process.core import GaussianProcessSurrogate
-from baybe.utils.basic import is_all_instance
 
 if TYPE_CHECKING:
     from botorch.models.model import ModelList
@@ -125,13 +123,17 @@ def fit(
     @override
     def to_botorch(self) -> ModelList:
         from botorch.models import ModelList
+        from botorch.models.gpytorch import GPyTorchModel
         from botorch.models.model_list_gp_regression import ModelListGP
 
-        surrogates = self._surrogates_flat
-        if is_all_instance(surrogates, GaussianProcessSurrogate):
-            return ModelListGP(*(s.to_botorch() for s in surrogates))
-        else:
-            return ModelList(*(s.to_botorch() for s in surrogates))
+        models = [s.to_botorch() for s in self._surrogates_flat]
+        # `ModelListGP` only accepts GPyTorch models (e.g. `SingleTaskGP`). If any
+        # surrogate yields a plain `Model` (e.g. an RGPE ensemble), fall back to the
+        # generic `ModelList`.
+        gp_models = [model for model in models if isinstance(model, GPyTorchModel)]
+        if len(gp_models) == len(models):
+            return ModelListGP(*gp_models)
+        return ModelList(*models)
 
     def posterior(self, candidates: pd.DataFrame, joint: bool = True) -> PosteriorList:
         """Compute the posterior for candidates in experimental representation.

baybe/surrogates/gaussian_process/core.py

@@ -11,7 +11,7 @@
 from typing import TYPE_CHECKING, ClassVar
 
 import pandas as pd
-from attrs import Converter, define, field
+from attrs import Converter, define, evolve, field
 from attrs.converters import optional as optional_c
 from attrs.converters import pipe
 from attrs.validators import instance_of, is_callable, optional
@@ -20,13 +20,15 @@
 from baybe.exceptions import (
     DeprecationError,
     IncompatibleSearchSpaceError,
+    IncompatibleSurrogateError,
     ModelNotTrainedError,
     _UnsupportedSearchSpaceAttributeError,
 )
 from baybe.kernels.base import Kernel
 from baybe.objectives.base import Objective
 from baybe.parameters.base import Parameter
 from baybe.parameters.categorical import TaskParameter
+from baybe.parameters.enum import TransferLearningMode
 from baybe.searchspace.core import SearchSpace
 from baybe.surrogates.base import Surrogate
 from baybe.surrogates.gaussian_process import _override
@@ -253,6 +255,14 @@ class GaussianProcessSurrogate(Surrogate):
     _model = field(init=False, default=None, eq=False)
     """The fitted BoTorch model."""
 
+    # TODO: type should be `Surrogate | None` but is currently omitted due to:
+    #   https://github.com/python-attrs/cattrs/issues/531
+    _delegate = field(init=False, default=None, eq=False)
+    """A transfer-learning surrogate to forward fitting and prediction to.
+
+    Set when a non-kernel mode (e.g.
+    :attr:`~baybe.parameters.enum.TransferLearningMode.RGPE`) is requested."""
+
     @staticmethod
     def _make_input_transform(context: _ModelContext) -> Normalize:
         """Create the input transform for the Gaussian process."""
@@ -352,7 +362,20 @@ def posterior_mean_function(
         Returns:
             A mean module ready to be used as the mean of a new
             :class:`GaussianProcessSurrogate`.
+
+        Raises:
+            IncompatibleSurrogateError: If the surrogate dispatches to a transfer
+                learning ensemble, for which posterior mean functions are not
+                implemented.
         """
+        if self._delegate is not None:
+            raise IncompatibleSurrogateError(
+                f"Providing a posterior mean function for a "
+                f"'{self.__class__.__name__}' that dispatches to a "
+                f"'{type(self._delegate).__name__}' transfer learning ensemble is not "
+                f"implemented."
+            )
+
         if self._model is None:
             warnings.warn(
                 f"'{self.__class__.__name__}' has not been fitted yet. "
@@ -381,6 +404,8 @@ def posterior_mean_function(
 
     @override
     def to_botorch(self) -> GPyTorchModel:
+        if self._delegate is not None:
+            return self._delegate.to_botorch()
         if self._model is None:
             raise ModelNotTrainedError(
                 "The surrogate must be trained before a BoTorch model can be created."
@@ -401,6 +426,9 @@ def _make_target_scaler_factory() -> type[OutcomeTransform] | None:
 
     @override
     def _posterior(self, candidates_comp_scaled: Tensor, /) -> Posterior:
+        # Forward to the transfer-learning delegate if one was set up during fitting.
+        if self._delegate is not None:
+            return self._delegate._posterior(candidates_comp_scaled)
         # Model being fit is guaranteed by the call in `posterior`
         assert self._model is not None
         return self._model.posterior(candidates_comp_scaled)
@@ -590,10 +618,23 @@ def _fit(self, train_x: Tensor, train_y: Tensor) -> None:
 
         context = _ModelContext(self._searchspace, self._objective, self._measurements)
 
+        # RGPE is handled by a dedicated ensemble surrogate rather than a task kernel.
+        # Dispatch to it, reusing this GP's configuration for the inner models.
+        task_param = self._searchspace._task_parameter
+        if (
+            task_param is not None
+            and task_param.override_transfer_learning_mode is TransferLearningMode.RGPE
+        ):
+            from baybe.surrogates.transfer_learning.rgpe import RGPESurrogate
+
+            delegate = RGPESurrogate(base_surrogate=evolve(self))
+            delegate.fit(self._searchspace, self._objective, self._measurements)
+            self._delegate = delegate
+            return
+
         # Check for custom kernel + multi-task clash (only relevant when the task
         # parameter has no kernel override, since the override mechanism handles task
         # kernel attachment explicitly).
-        task_param = self._searchspace._task_parameter
         has_task_override = (
             task_param is not None and task_param.override_kernel is not None
         )

baybe/surrogates/gaussian_process/presets/baybe.py

@@ -15,7 +15,7 @@
 from baybe.kernels.basic import PositiveIndexKernel
 from baybe.objectives.base import Objective
 from baybe.parameters.categorical import TaskParameter
-from baybe.parameters.enum import _ParameterKind
+from baybe.parameters.enum import TransferLearningMode, _ParameterKind
 from baybe.parameters.selectors import (
     ParameterSelectorProtocol,
     TypeSelector,
@@ -274,11 +274,22 @@ class BayBEFitCriterionFactory(FitCriterionFactoryProtocol):
     def __call__(
         self, searchspace: SearchSpace, objective: Objective, measurements: pd.DataFrame
     ) -> FitCriterion:
-        return (
-            FitCriterion.MARGINAL_LOG_LIKELIHOOD
-            if searchspace.n_tasks == 1
-            else FitCriterion.LEAVE_ONE_OUT_PSEUDOLIKELIHOOD
-        )
+        task_param = searchspace._task_parameter
+
+        # Without a (multi-valued) task parameter, this is an ordinary single-task GP.
+        if searchspace.n_tasks == 1:
+            return FitCriterion.MARGINAL_LOG_LIKELIHOOD
+        # IDENTITY mode keeps the task dimension but makes it inert (constant task
+        # kernel), so the model is effectively single-task and is fit like one.
+        elif (
+            task_param is not None
+            and task_param.override_transfer_learning_mode
+            is TransferLearningMode.IDENTITY
+        ):
+            return FitCriterion.MARGINAL_LOG_LIKELIHOOD
+        # Genuine multi-task model: use LOO cross-validation pseudo-likelihood.
+        else:
+            return FitCriterion.LEAVE_ONE_OUT_PSEUDOLIKELIHOOD
 
 
 # Collect leftover original slotted classes processed by `attrs.define`