diff --git a/CHANGELOG.md b/CHANGELOG.md index b143763a6c..c623592d67 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/baybe/kernels/__init__.py b/baybe/kernels/__init__.py index 9323a2b631..1f9d36ccee 100644 --- a/baybe/kernels/__init__.py +++ b/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", diff --git a/baybe/kernels/basic.py b/baybe/kernels/basic.py index b740235e63..b3bcd764ae 100644 --- a/baybe/kernels/basic.py +++ b/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): diff --git a/baybe/parameters/categorical.py b/baybe/parameters/categorical.py index 47b642f71e..af7eb8af95 100644 --- a/baybe/parameters/categorical.py +++ b/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 diff --git a/baybe/parameters/enum.py b/baybe/parameters/enum.py index 45a32bb2c1..5fae023b40 100644 --- a/baybe/parameters/enum.py +++ b/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`. + """ diff --git a/baybe/searchspace/core.py b/baybe/searchspace/core.py index 90093bbee9..7e2a00d783 100644 --- a/baybe/searchspace/core.py +++ b/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): diff --git a/baybe/surrogates/__init__.py b/baybe/surrogates/__init__.py index 60910e4be7..10189fa4fe 100644 --- a/baybe/surrogates/__init__.py +++ b/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", ] diff --git a/baybe/surrogates/composite.py b/baybe/surrogates/composite.py index acf9149861..3478997860 100644 --- a/baybe/surrogates/composite.py +++ b/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. diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py index 0a76661518..4560a55380 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/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,6 +20,7 @@ from baybe.exceptions import ( DeprecationError, IncompatibleSearchSpaceError, + IncompatibleSurrogateError, ModelNotTrainedError, _UnsupportedSearchSpaceAttributeError, ) @@ -27,6 +28,7 @@ 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 ) diff --git a/baybe/surrogates/gaussian_process/presets/baybe.py b/baybe/surrogates/gaussian_process/presets/baybe.py index f68717f6dd..d190481dbe 100644 --- a/baybe/surrogates/gaussian_process/presets/baybe.py +++ b/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` diff --git a/baybe/surrogates/transfer_learning/__init__.py b/baybe/surrogates/transfer_learning/__init__.py new file mode 100644 index 0000000000..6c5e375987 --- /dev/null +++ b/baybe/surrogates/transfer_learning/__init__.py @@ -0,0 +1,12 @@ +"""Transfer learning surrogates. + +This subpackage collects surrogates that build a target model from one or more source +models, dispatched via +:attr:`~baybe.parameters.categorical.TaskParameter.override_transfer_learning_mode`. +""" + +from baybe.surrogates.transfer_learning.rgpe import RGPESurrogate + +__all__ = [ + "RGPESurrogate", +] diff --git a/baybe/surrogates/transfer_learning/rgpe.py b/baybe/surrogates/transfer_learning/rgpe.py new file mode 100644 index 0000000000..15d0120bb4 --- /dev/null +++ b/baybe/surrogates/transfer_learning/rgpe.py @@ -0,0 +1,431 @@ +"""Rank-weighted Gaussian process ensemble (RGPE) surrogate for transfer learning. + +The implementation follows the BoTorch RGPE tutorial: +https://archive.botorch.org/v/latest/tutorials/meta_learning_with_rgpe +""" + +from __future__ import annotations + +import gc +from typing import TYPE_CHECKING, Any, ClassVar, cast + +from attrs import define, evolve, field +from attrs.validators import ge, instance_of +from typing_extensions import override + +from baybe.exceptions import IncompatibleSearchSpaceError +from baybe.surrogates.base import Surrogate +from baybe.surrogates.gaussian_process.core import GaussianProcessSurrogate +from baybe.utils.dataframe import to_tensor + +if TYPE_CHECKING: + import pandas as pd + from botorch.models.gpytorch import GPyTorchModel + from botorch.models.transforms.input import InputTransform + from botorch.models.transforms.outcome import OutcomeTransform + from botorch.posteriors import GPyTorchPosterior, Posterior + from gpytorch.means import Mean + from gpytorch.module import Module + from torch import Tensor + + from baybe.parameters.base import Parameter + from baybe.searchspace.core import SearchSpace + + +def _roll_col(x: Tensor, shift: int) -> Tensor: + """Rotate the columns of a tensor to the right by ``shift``. + + Args: + x: The tensor whose last dimension is rotated. + shift: The number of positions to rotate by. + + Returns: + The column-rotated tensor. + """ + import torch + + return torch.cat((x[..., -shift:], x[..., :-shift]), dim=-1) + + +def _compute_ranking_loss(f_samps: Tensor, target_y: Tensor) -> Tensor: + """Compute the ranking loss of each posterior sample against the target points. + + The loss counts discordant pairwise orderings between the sampled predictions and + the observed targets. A 3D ``f_samps`` holds leave-one-out samples of the target + model (its diagonal are the out-of-sample predictions); a 2D ``f_samps`` holds + source-model samples. + + Args: + f_samps: An ``n_samples x n x n`` tensor (target LOOCV samples) or an + ``n_samples x n`` tensor (source-model samples). + target_y: An ``n x 1`` tensor of observed target values. + + Returns: + An ``n_samples`` tensor with the ranking loss per sample. + """ + import torch + + n = target_y.shape[0] + if f_samps.ndim == 3: + # Target model: compare each LOO out-of-sample prediction (the diagonal) to + # every in-sample prediction. + cartesian_y = torch.cartesian_prod( + target_y.squeeze(-1), + target_y.squeeze(-1), + ).view(n, n, 2) + rank_loss = ( + ( + (f_samps.diagonal(dim1=1, dim2=2).unsqueeze(-1) < f_samps) + ^ (cartesian_y[..., 0] < cartesian_y[..., 1]) + ) + .sum(dim=-1) + .sum(dim=-1) + ) + else: + rank_loss = torch.zeros( + f_samps.shape[0], dtype=torch.long, device=target_y.device + ) + y_stack = target_y.squeeze(-1).expand(f_samps.shape) + for i in range(1, n): + rank_loss += ( + (_roll_col(f_samps, i) < f_samps) ^ (_roll_col(y_stack, i) < y_stack) + ).sum(dim=-1) + return rank_loss + + +def _loocv_sample_preds( + train_x: Tensor, train_y: Tensor, target_model: GPyTorchModel, n_samples: int +) -> Tensor: + """Draw leave-one-out cross-validation samples of the target model. + + Builds a batch GP in the target model's transformed input/output space, reusing its + fitted hyperparameters, with each batch element trained on all but one target point. + Working in the transformed space is fine because the ranking loss only depends on + the ordering of predictions, which output standardization preserves. + + Args: + train_x: The target inputs in the model's transformed input space (``n x d``). + train_y: The target values in the model's transformed output space (``n x 1``). + target_model: The fitted target Gaussian process model. + n_samples: The number of Monte Carlo samples to draw. + + Returns: + An ``n_samples x n x n`` tensor, indexing the ``n`` leave-one-out models along + dimension 1 and the ``n`` target points along dimension 2. + """ + from copy import deepcopy + + import torch + from botorch.models import SingleTaskGP + from botorch.sampling.normal import SobolQMCNormalSampler + + n = len(train_x) + masks = torch.eye(n, dtype=torch.bool, device=train_x.device) + train_x_cv = torch.stack([train_x[~m] for m in masks]) + train_y_cv = torch.stack([train_y[~m] for m in masks]) + + # Transform-free batch GP reusing the target model's fitted modules, which broadcast + # across the LOO batch. Transforms are off since the data is already in the model's + # transformed space (this also avoids re-standardizing the standardized targets). + model = SingleTaskGP( + train_x_cv, + train_y_cv, + covar_module=cast("Module", deepcopy(target_model.covar_module)), + mean_module=cast("Mean", deepcopy(target_model.mean_module)), + likelihood=deepcopy(target_model.likelihood), + input_transform=None, + outcome_transform=None, + ) + model.eval() + with torch.no_grad(): + posterior = model.posterior(train_x) + sampler = SobolQMCNormalSampler(sample_shape=torch.Size([n_samples])) + return sampler(posterior).squeeze(-1) + + +def _compute_rank_weights( + train_x: Tensor, + train_y: Tensor, + source_models: tuple[GPyTorchModel, ...], + target_model: GPyTorchModel, + n_samples: int, +) -> Tensor: + """Compute the RGPE rank weights for the source and target models. + + Each model is weighted by the Monte Carlo probability that it has the lowest ranking + loss on the target points, following Feurer, Letham and Bakshy (ICML 2018 AutoML + Workshop). Mirrors the BoTorch RGPE tutorial without weight-dilution regularization. + + Args: + train_x: The target training inputs (raw comp-rep, ``n x d``). + train_y: The target training targets (``n x 1``). + source_models: The fitted source Gaussian process models. + target_model: The fitted target Gaussian process model. + n_samples: The number of Monte Carlo samples used to estimate the weights. + + Returns: + A tensor of length ``len(source_models) + 1`` with the model weights, ordered as + ``(sources..., target)``. + """ + import torch + from botorch.sampling.normal import SobolQMCNormalSampler + + ranking_losses = [] + sampler = SobolQMCNormalSampler(sample_shape=torch.Size([n_samples])) + for model in source_models: + posterior = model.posterior(train_x) + f_samps = sampler(posterior).squeeze(-1).squeeze(-1) + ranking_losses.append(_compute_ranking_loss(f_samps, train_y)) + + # The target model's LOOCV samples are drawn in its transformed space. The output + # transform is only order-preserving, which is all the ranking loss depends on. + target_model.eval() + transformed_x = target_model.transform_inputs(train_x) + # `outcome_transform` is always present on a fitted SingleTaskGP, but mypy sees it + # as an optional, non-callable attribute. + transformed_y, _ = target_model.outcome_transform(train_y) # type: ignore[operator] + target_f_samps = _loocv_sample_preds( + transformed_x, transformed_y, target_model, n_samples + ) + ranking_losses.append(_compute_ranking_loss(target_f_samps, train_y)) + + ranking_loss_tensor = torch.stack(ranking_losses) + best_models = torch.argmin(ranking_loss_tensor, dim=0) + rank_weights = ( + best_models.bincount(minlength=len(ranking_losses)).type_as(train_x) / n_samples + ) + return rank_weights + + +@define +class RGPESurrogate(Surrogate): + """A rank-weighted Gaussian process ensemble (RGPE) for transfer learning. + + Implements the ensemble of Feurer, Letham and Bakshy (ICML 2018 AutoML Workshop, + https://arxiv.org/abs/1802.02219): one Gaussian process is fitted per source task + and, once enough target data exists, one on the target task. The posterior is a + rank-weighted sum of the individual posteriors (``mean = Σ wᵢ μᵢ``, + ``cov = Σ wᵢ² Σᵢ``), with weights measuring how well each model ranks the target + points (via leave-one-out cross-validation for the target model). + + Select it directly or by setting + :attr:`~baybe.parameters.categorical.TaskParameter.override_transfer_learning_mode` + to :attr:`~baybe.parameters.enum.TransferLearningMode.RGPE`. + + Computing the weights needs at least two target points; with fewer, they fall back + to a uniform average over the available models. + + Note: + Only single-output objectives and a single target task are supported, and no + weight-dilution regularization is applied. The ``to_botorch`` representation + does not support ``fantasize`` (and hence acquisition functions that require it, + such as the knowledge gradient). + """ + + supports_kernel_overrides: ClassVar[bool] = True + # See base class. The inner GPs are fitted on the identity-mode space, which still + # carries the parameter overrides (only the task dimension is made inert). + + supports_transfer_learning: ClassVar[bool] = True + # See base class. + + base_surrogate: GaussianProcessSurrogate = field( + factory=GaussianProcessSurrogate, + validator=instance_of(GaussianProcessSurrogate), + ) + """The Gaussian process configuration used for the source and target models.""" + + n_mc_samples: int = field(default=256, validator=[instance_of(int), ge(1)]) + """The number of Monte Carlo samples used to estimate the ranking weights.""" + + _source_gps: tuple[GaussianProcessSurrogate, ...] = field( + init=False, factory=tuple, eq=False, repr=False + ) + """The GPs trained on the source data, one per source task that has measurements. + + Available after fitting. + """ + + _target_gp: GaussianProcessSurrogate | None = field( + init=False, default=None, eq=False, repr=False + ) + """The GP trained on the target data. + + ``None`` before fitting or when the target task has no measurements yet. + """ + + # TODO: type should be `Tensor | None` but is currently omitted due to: + # https://github.com/python-attrs/cattrs/issues/531 + _weights = field(init=False, default=None, eq=False, repr=False) + """The ensemble weights, ordered as ``(sources..., target)``. Available after + fitting. The target entry is present only when a target GP was fitted.""" + + @override + def _fit(self, train_x: Tensor, train_y: Tensor) -> None: + # The base-class inputs are ignored: the inner GPs are refitted from the stored + # measurements over the identity-mode space (task dimension inert). + import torch + + assert self._searchspace is not None # ensured by base class + assert self._objective is not None # ensured by base class + assert self._measurements is not None # ensured by base class + + identity_searchspace, sources, target_measurements = self._split_measurements() + + # Fit one source GP per source task that has data. + source_gps = [] + for _, source_measurements in sources: + source_gp = evolve(self.base_surrogate) + source_gp.fit(identity_searchspace, self._objective, source_measurements) + source_gps.append(source_gp) + self._source_gps = tuple(source_gps) + + source_models = tuple(gp._model for gp in self._source_gps) + n_target = len(target_measurements) + + if n_target == 0: + # Cold start: no target data yet, average over the source models only. + self._target_gp = None + n_models = len(source_models) + self._weights = torch.full((n_models,), 1.0 / n_models) + return + + self._target_gp = evolve(self.base_surrogate) + self._target_gp.fit(identity_searchspace, self._objective, target_measurements) + + n_models = len(source_models) + 1 + if n_target < 2: + # Not enough target points to rank, average over all available models. + self._weights = torch.full((n_models,), 1.0 / n_models) + return + + train_x_target = to_tensor( + identity_searchspace.transform(target_measurements, allow_extra=True) + ) + train_y_target = to_tensor( + self._objective._pre_transform(target_measurements, allow_extra=True) + ) + if train_y_target.ndim == 1: + train_y_target = train_y_target.unsqueeze(-1) + self._weights = _compute_rank_weights( + train_x_target, + train_y_target, + source_models, + self._target_gp._model, + self.n_mc_samples, + ) + + @override + @staticmethod + def _make_parameter_scaler_factory( + parameter: Parameter, / + ) -> type[InputTransform] | None: + # The inner Gaussian processes handle input scaling themselves. + return None + + @override + @staticmethod + def _make_target_scaler_factory() -> type[OutcomeTransform] | None: + # The inner Gaussian processes handle output scaling themselves. + return None + + @override + def _posterior(self, candidates_comp_scaled: Tensor, /) -> Posterior: + import torch + from botorch.posteriors import GPyTorchPosterior + from gpytorch.distributions import MultivariateNormal + + assert self._weights is not None # set during fitting + + models: list[GaussianProcessSurrogate] = list(self._source_gps) + if self._target_gp is not None: + models.append(self._target_gp) + + # Restrict the combination to models with non-zero weight and renormalize. + nonzero = self._weights > 0 + weights = self._weights[nonzero] / self._weights[nonzero].sum() + selected = [m for m, keep in zip(models, nonzero.tolist()) if keep] + + weighted_means = [] + weighted_covars = [] + for weight, model in zip(weights, selected): + # Inner GPs share the full comp-rep layout, so candidates pass through + # unchanged. + posterior = cast( + "GPyTorchPosterior", model._posterior(candidates_comp_scaled) + ) + weighted_means.append(weight * posterior.mean.squeeze(-1)) + weighted_covars.append( + posterior.distribution.lazy_covariance_matrix * weight**2 + ) + + mean = torch.stack(weighted_means).sum(dim=0) + covariance = weighted_covars[0] + for covar in weighted_covars[1:]: + covariance = covariance + covar + return GPyTorchPosterior(MultivariateNormal(mean, covariance)) + + def _split_measurements( + self, + ) -> tuple[SearchSpace, list[tuple[Any, pd.DataFrame]], pd.DataFrame]: + """Validate the task configuration and split the measurements by task. + + Returns: + The search space switched to identity mode (task dimension inert but still + present in the layout), an ordered list of ``(task_value, measurements)`` + pairs for the source tasks that have data, and the target-task measurements + (which may be empty). + + Raises: + IncompatibleSearchSpaceError: If the space has no task parameter, does not + describe exactly one active (target) task, or if no source task has + measurements. + """ + from baybe.parameters.enum import TransferLearningMode + + assert self._searchspace is not None + assert self._measurements is not None + + searchspace = self._searchspace + measurements = self._measurements + + task_param = searchspace._task_parameter + if task_param is None: + raise IncompatibleSearchSpaceError( + f"'{self.__class__.__name__}' requires a search space that contains a " + f"task parameter." + ) + + active_values = set(task_param.active_values) + source_values = set(task_param.values) - active_values + if len(active_values) != 1: + raise IncompatibleSearchSpaceError( + f"'{self.__class__.__name__}' requires exactly one active (target) " + f"task value, but the task parameter describes {len(active_values)}." + ) + # A task parameter has at least two values, so one active value leaves at least + # one source value. + (target_value,) = active_values + + task_name = task_param.name + sources = [ + (value, subset) + for value in task_param.values + if value in source_values + if not (subset := measurements[measurements[task_name] == value]).empty + ] + if not sources: + raise IncompatibleSearchSpaceError( + f"'{self.__class__.__name__}' requires measurements for at least one " + f"source task, but none were provided." + ) + target_measurements = measurements[measurements[task_name] == target_value] + + identity_searchspace = searchspace._with_task_mode( + TransferLearningMode.IDENTITY + ) + return identity_searchspace, sources, target_measurements + + +# Collect leftover original slotted classes processed by `attrs.define` +gc.collect() diff --git a/benchmarks/domains/aryl_halides/core.py b/benchmarks/domains/aryl_halides/core.py index 049d46d1b8..04b2a26e14 100644 --- a/benchmarks/domains/aryl_halides/core.py +++ b/benchmarks/domains/aryl_halides/core.py @@ -18,6 +18,7 @@ from baybe.objectives import SingleTargetObjective from baybe.parameters import SubstanceParameter, TaskParameter from baybe.parameters.base import DiscreteParameter +from baybe.parameters.enum import TransferLearningMode from baybe.searchspace import SearchSpace from baybe.settings import Settings from baybe.simulation import simulate_scenarios @@ -48,6 +49,7 @@ def make_searchspace( data: pd.DataFrame, target_tasks: Sequence[str] | None = None, source_tasks: Sequence[str] | None = None, + transfer_learning_mode: TransferLearningMode | None = None, ) -> SearchSpace: """Create the search space for the benchmark.""" params: list[DiscreteParameter] = [ @@ -65,6 +67,7 @@ def make_searchspace( name="aryl_halide", values=all_tasks, active_values=target_tasks, + override_transfer_learning_mode=transfer_learning_mode, ) ) return SearchSpace.from_product(parameters=params) @@ -113,7 +116,12 @@ def aryl_halide_tl_substance_benchmark( source_tasks=source_tasks, target_tasks=target_tasks, ) - searchspace_nontl = make_searchspace(data=data) + searchspace_nontl = make_searchspace( + data=data, + source_tasks=source_tasks, + target_tasks=target_tasks, + transfer_learning_mode=TransferLearningMode.IDENTITY, + ) lookup = make_lookup(data, target_tasks) initial_data = make_initial_data(data, source_tasks) diff --git a/benchmarks/domains/direct_arylation/convergence_tl.py b/benchmarks/domains/direct_arylation/convergence_tl.py index a5d2bc24eb..268d7ac724 100644 --- a/benchmarks/domains/direct_arylation/convergence_tl.py +++ b/benchmarks/domains/direct_arylation/convergence_tl.py @@ -15,6 +15,7 @@ TaskParameter, ) from baybe.parameters.base import DiscreteParameter +from baybe.parameters.enum import TransferLearningMode from baybe.searchspace import SearchSpace from baybe.settings import Settings from baybe.simulation import simulate_scenarios @@ -40,7 +41,8 @@ def load_data() -> pd.DataFrame: def make_searchspace( data: pd.DataFrame, - use_task_parameter: bool, + use_task_parameter: bool = True, + transfer_learning_mode: TransferLearningMode | None = None, ) -> SearchSpace: """Create the search space for the benchmark.""" params: list[DiscreteParameter] = [ @@ -62,6 +64,7 @@ def make_searchspace( name="Temp_C", values=["90", "105", "120"], active_values=["105"], + override_transfer_learning_mode=transfer_learning_mode, ) ) return SearchSpace.from_product(parameters=params) @@ -120,11 +123,10 @@ def direct_arylation_tl_temperature( searchspace = make_searchspace( data=data, - use_task_parameter=True, ) searchspace_nontl = make_searchspace( data=data, - use_task_parameter=False, + transfer_learning_mode=TransferLearningMode.IDENTITY, ) lookup = make_lookup(data) diff --git a/benchmarks/domains/easom/convergence_tl.py b/benchmarks/domains/easom/convergence_tl.py index 6debc5ca2d..26e50d7e92 100644 --- a/benchmarks/domains/easom/convergence_tl.py +++ b/benchmarks/domains/easom/convergence_tl.py @@ -11,6 +11,7 @@ from baybe.objectives import SingleTargetObjective from baybe.parameters import NumericalDiscreteParameter, TaskParameter from baybe.parameters.base import DiscreteParameter +from baybe.parameters.enum import TransferLearningMode from baybe.searchspace import SearchSpace from baybe.settings import Settings from baybe.simulation import simulate_scenarios @@ -90,9 +91,15 @@ def easom_tl_47_negate_noise5(settings: ConvergenceBenchmarkSettings) -> pd.Data values=["Target_Function", "Source_Function"], active_values=["Target_Function"], ) + task_param_naive = TaskParameter( + name="Function", + values=["Target_Function", "Source_Function"], + active_values=["Target_Function"], + override_transfer_learning_mode=TransferLearningMode.IDENTITY, + ) params_tl = params + [task_param] - searchspace_nontl = SearchSpace.from_product(parameters=params) + searchspace_nontl = SearchSpace.from_product(parameters=params + [task_param_naive]) searchspace_tl = SearchSpace.from_product(parameters=params_tl) objective = SingleTargetObjective( diff --git a/benchmarks/domains/hartmann/convergence_tl.py b/benchmarks/domains/hartmann/convergence_tl.py index 37646b81c2..9f3fe9b64a 100755 --- a/benchmarks/domains/hartmann/convergence_tl.py +++ b/benchmarks/domains/hartmann/convergence_tl.py @@ -13,6 +13,7 @@ from baybe.objectives import SingleTargetObjective from baybe.parameters import NumericalDiscreteParameter, TaskParameter from baybe.parameters.base import DiscreteParameter +from baybe.parameters.enum import TransferLearningMode from baybe.searchspace import SearchSpace from baybe.settings import Settings from baybe.simulation import simulate_scenarios @@ -105,9 +106,17 @@ def benchmark_fn(settings: ConvergenceBenchmarkSettings) -> pd.DataFrame: values=("Target_Function", "Source_Function"), active_values=("Target_Function",), ) + task_param_naive = TaskParameter( + name="Function", + values=("Target_Function", "Source_Function"), + active_values=("Target_Function",), + override_transfer_learning_mode=TransferLearningMode.IDENTITY, + ) params_tl = params + [task_param] - searchspace_nontl = SearchSpace.from_product(parameters=params) + searchspace_nontl = SearchSpace.from_product( + parameters=params + [task_param_naive] + ) searchspace_tl = SearchSpace.from_product(parameters=params_tl) objective = SingleTargetObjective( diff --git a/benchmarks/domains/michalewicz/convergence_tl.py b/benchmarks/domains/michalewicz/convergence_tl.py index ae4e58845b..932167547b 100644 --- a/benchmarks/domains/michalewicz/convergence_tl.py +++ b/benchmarks/domains/michalewicz/convergence_tl.py @@ -18,6 +18,7 @@ from baybe.objectives import SingleTargetObjective from baybe.parameters import NumericalContinuousParameter, TaskParameter from baybe.parameters.base import Parameter +from baybe.parameters.enum import TransferLearningMode from baybe.searchspace import SearchSpace from baybe.settings import Settings from baybe.simulation import simulate_scenarios @@ -26,7 +27,9 @@ from benchmarks.definition.base import RunMode -def make_searchspace(use_task_parameter: bool) -> SearchSpace: +def make_searchspace( + transfer_learning_mode: TransferLearningMode | None = None, +) -> SearchSpace: """Create search space for the benchmark.""" params: list[Parameter] = [ NumericalContinuousParameter( @@ -35,14 +38,14 @@ def make_searchspace(use_task_parameter: bool) -> SearchSpace: ) for k in range(5) ] - if use_task_parameter: - params.append( - TaskParameter( - name="Function", - values=["Target_Function", "Source_Function"], - active_values=["Target_Function"], - ) + params.append( + TaskParameter( + name="Function", + values=["Target_Function", "Source_Function"], + active_values=["Target_Function"], + override_transfer_learning_mode=transfer_learning_mode, ) + ) return SearchSpace.from_product(parameters=params) @@ -117,7 +120,7 @@ def michalewicz_tl_continuous(settings: ConvergenceBenchmarkSettings) -> pd.Data • Includes baseline with no transfer learning (0 points) • Creates two campaign types: - Transfer learning: with task parameter to distinguish data sources - - Non-transfer learning (naive): without task parameter + - Non-transfer learning (naive): task parameter with an inert identity kernel Args: settings: Configuration settings for the convergence benchmark @@ -131,8 +134,10 @@ def michalewicz_tl_continuous(settings: ConvergenceBenchmarkSettings) -> pd.Data "Source_Function": Michalewicz(dim=5, negate=True, noise_std=0.15), "Target_Function": Michalewicz(dim=5, negate=True), } - searchspace_nontl = make_searchspace(use_task_parameter=False) - searchspace_tl = make_searchspace(use_task_parameter=True) + searchspace_nontl = make_searchspace( + transfer_learning_mode=TransferLearningMode.IDENTITY + ) + searchspace_tl = make_searchspace() objective = make_objective() campaign_tl = Campaign( diff --git a/docs/concepts/transfer_learning.md b/docs/concepts/transfer_learning.md index 3273e20853..c371e53972 100644 --- a/docs/concepts/transfer_learning.md +++ b/docs/concepts/transfer_learning.md @@ -156,9 +156,11 @@ TaskParameter( Note that the transfer learning mechanism can be implement in many different ways, for example, using multiple models, shared architectures, special kernels in a single model, etc. -Encoding the context information via a [TaskParameter] is not restricted to any of -these methods, even though BayBE currently only implements a variant via kernels, -specifically: the Intrinsic Coregionalization Model {cite:p}`NIPS2007_66368270`. +Encoding the context information via a [`TaskParameter`] is not restricted to any of +these methods. BayBE currently implements a kernel-based variant (the Intrinsic +Coregionalization Model {cite:p}`NIPS2007_66368270`) and an ensemble-based variant +(the rank-weighted Gaussian process ensemble {cite:p}`Feurer2018`), selectable via +the [`TaskParameter`]'s `override_transfer_learning_mode` argument. ``` ## Seeing Transfer Learning in Action diff --git a/docs/references.bib b/docs/references.bib index 714fdb3e3b..d13979754d 100644 --- a/docs/references.bib +++ b/docs/references.bib @@ -9,6 +9,13 @@ @inproceedings{NIPS2007_66368270 year = {2007} } +@inproceedings{Feurer2018, + author = {Feurer, Matthias and Letham, Benjamin and Bakshy, Eytan}, + booktitle = {ICML 2018 AutoML Workshop}, + title = {Scalable Meta-Learning for {Bayesian} Optimization using Ranking-Weighted {Gaussian} Process Ensembles}, + year = {2018} +} + @article{Chen2026, author = {Chen, Guanming and Fleck, Maximilian and Stuyver, Thijs}, title = {Leveraging Hidden-Space Representations Effectively in Bayesian Optimization for Experiment Design through Dimension-Aware Hyperpriors}, diff --git a/tests/hypothesis_strategies/kernels.py b/tests/hypothesis_strategies/kernels.py index 341b14f7fe..bddfb01e28 100644 --- a/tests/hypothesis_strategies/kernels.py +++ b/tests/hypothesis_strategies/kernels.py @@ -6,6 +6,7 @@ import hypothesis.strategies as st from baybe.kernels.basic import ( + IdentityKernel, IndexKernel, LinearKernel, MaternKernel, @@ -40,6 +41,14 @@ def active_parameter_names(names: Sequence[str] | None = None): return st.just(names) +def identity_kernels(parameter_names: Sequence[str] | None = None): + """A strategy that generates identity kernels.""" + return st.builds( + IdentityKernel, + parameter_names=active_parameter_names(parameter_names), + ) + + def linear_kernels(parameter_names: Sequence[str] | None = None): """A strategy that generates linear kernels.""" return st.builds( @@ -144,6 +153,7 @@ def base_kernels(parameter_names: Sequence[str] | None = None): return st.one_of( [ matern_kernels(parameter_names), # on top because it is the default + identity_kernels(parameter_names), linear_kernels(parameter_names), rbf_kernels(parameter_names), rq_kernels(parameter_names), diff --git a/tests/hypothesis_strategies/surrogates.py b/tests/hypothesis_strategies/surrogates.py index 45fee48f07..d2529c981e 100644 --- a/tests/hypothesis_strategies/surrogates.py +++ b/tests/hypothesis_strategies/surrogates.py @@ -24,6 +24,7 @@ from baybe.surrogates.gaussian_process.presets.edbo_smoothed import ( SmoothedEDBOMeanFactory, ) +from baybe.surrogates.transfer_learning.rgpe import RGPESurrogate from tests.hypothesis_strategies.kernels import kernels _MEAN_FACTORIES = [ @@ -65,3 +66,16 @@ def gaussian_process_surrogates(): st.builds(BayBEFitCriterionFactory), ), ) + + +def rgpe_surrogates(): + """A strategy generating RGPE surrogate instances with varying configurations. + + The inner Gaussian process configuration is drawn from the full space of GP + surrogate configurations and the number of Monte Carlo samples is varied. + """ + return st.builds( + RGPESurrogate, + base_surrogate=gaussian_process_surrogates(), + n_mc_samples=st.integers(min_value=1, max_value=512), + ) diff --git a/tests/serialization/test_surrogate_serialization.py b/tests/serialization/test_surrogate_serialization.py index c3899558eb..d0bc96c8f6 100644 --- a/tests/serialization/test_surrogate_serialization.py +++ b/tests/serialization/test_surrogate_serialization.py @@ -10,8 +10,12 @@ from baybe.surrogates.gaussian_process.core import GaussianProcessSurrogate from baybe.surrogates.ngboost import NGBoostSurrogate from baybe.surrogates.random_forest import RandomForestSurrogate +from baybe.surrogates.transfer_learning.rgpe import RGPESurrogate from baybe.utils.basic import get_subclasses -from tests.hypothesis_strategies.surrogates import gaussian_process_surrogates +from tests.hypothesis_strategies.surrogates import ( + gaussian_process_surrogates, + rgpe_surrogates, +) from tests.serialization.utils import assert_roundtrip_consistency @@ -20,7 +24,7 @@ [ c for c in get_subclasses(Surrogate) - if not issubclass(c, GaussianProcessSurrogate) + if not issubclass(c, (GaussianProcessSurrogate, RGPESurrogate)) ], ) def test_surrogate_roundtrip(request, surrogate_cls: type[Surrogate]): @@ -43,6 +47,12 @@ def test_gaussian_process_surrogate_roundtrip(surrogate: GaussianProcessSurrogat assert_roundtrip_consistency(surrogate) +@given(rgpe_surrogates()) +def test_rgpe_surrogate_roundtrip(surrogate: RGPESurrogate): + """A serialization roundtrip yields an equivalent RGPE surrogate for all configs.""" + assert_roundtrip_consistency(surrogate) + + @pytest.mark.parametrize( "surrogate", [ diff --git a/tests/test_iterations.py b/tests/test_iterations.py index 1c6bd86fec..9e149a3b65 100644 --- a/tests/test_iterations.py +++ b/tests/test_iterations.py @@ -59,6 +59,7 @@ EDBOKernelFactory, ) from baybe.surrogates.linear import BayesianLinearSurrogate +from baybe.surrogates.transfer_learning.rgpe import RGPESurrogate from baybe.targets.numerical import NumericalTarget from baybe.utils.basic import get_subclasses from baybe.utils.dataframe import create_fake_input @@ -83,6 +84,9 @@ def ongoing_campaign(campaign): cls, BetaBernoulliMultiArmedBanditSurrogate ): continue + # RGPE needs a transfer-learning space (a task parameter); see test_rgpe_surrogate. + if issubclass(cls, RGPESurrogate): + continue p = param( cls(), id=cls.__name__, @@ -397,9 +401,19 @@ def test_kernels(ongoing_campaign, n_iterations, batch_size): ] ), ] - for mode in [None, *TransferLearningMode] + for mode in [ + None, + *(m for m in TransferLearningMode if m is not TransferLearningMode.RGPE), + ] + ], + ids=[ + "parameter_kernel_override", + *( + mode.name + for mode in TransferLearningMode + if mode is not TransferLearningMode.RGPE + ), ], - ids=["parameter_kernel_override", *(mode.name for mode in TransferLearningMode)], ) def test_parameter_kernel_override_iteration( ongoing_campaign, n_iterations, batch_size diff --git a/tests/test_parameter_kernel_overrides.py b/tests/test_parameter_kernel_overrides.py index 8fb01d28b7..c32ca753e7 100644 --- a/tests/test_parameter_kernel_overrides.py +++ b/tests/test_parameter_kernel_overrides.py @@ -18,7 +18,7 @@ from baybe.exceptions import IncompatibleOverrideError from baybe.kernels import MaternKernel, RBFKernel -from baybe.kernels.basic import IndexKernel, PositiveIndexKernel +from baybe.kernels.basic import IdentityKernel, IndexKernel, PositiveIndexKernel from baybe.kernels.composite import AdditiveKernel, ProductKernel, ScaleKernel from baybe.parameters import ( CategoricalParameter, @@ -66,7 +66,7 @@ def _rbf_kernel(lengthscale, *, frozen=False): @pytest.mark.parametrize( "mode", - [None, *TransferLearningMode], + [None, *(m for m in TransferLearningMode if m is not TransferLearningMode.RGPE)], ids=lambda mode: mode.name if mode else "no-task", ) def test_selector_does_not_exclude_overridden_parameters(mode): @@ -155,6 +155,7 @@ def test_nondefault_residual_indices(kernel_or_factory, expected_residual_dims): TaskParameter("task", ["a", "b"], override_transfer_learning_mode=mode), ) for mode in TransferLearningMode + if mode is not TransferLearningMode.RGPE ], ], ids=[ @@ -162,7 +163,11 @@ def test_nondefault_residual_indices(kernel_or_factory, expected_residual_dims): "multi-dim", "all-overridden", "task-without-tl-override", - *(mode.name for mode in TransferLearningMode), + *( + mode.name + for mode in TransferLearningMode + if mode is not TransferLearningMode.RGPE + ), ], ) def test_fitted_model_uses_parameter_kernel_overrides( @@ -176,12 +181,12 @@ def test_fitted_model_uses_parameter_kernel_overrides( expected = {"base": gk.MaternKernel, "override": gk.RBFKernel} if task_parameter is not None: parameters.append(task_parameter) - expected["task"] = ( - gk.IndexKernel - if task_parameter.override_transfer_learning_mode - == TransferLearningMode.INDEX_KERNEL - else BoPositiveIndexKernel - ) + expected["task"] = { + None: BoPositiveIndexKernel, + TransferLearningMode.INDEX_KERNEL: gk.IndexKernel, + TransferLearningMode.POSITIVE_INDEX_KERNEL: BoPositiveIndexKernel, + TransferLearningMode.IDENTITY: gk.ConstantKernel, + }[task_parameter.override_transfer_learning_mode] searchspace = SearchSpace.from_product(parameters) measurements = pd.DataFrame( [ @@ -201,7 +206,12 @@ def test_fitted_model_uses_parameter_kernel_overrides( searchspace.get_comp_rep_parameter_indices(name): cls for name, cls in expected.items() } - assert all(k.ard_num_dims == len(k.active_dims) for k in leaves) + # The constant (identity) task kernel carries no lengthscales, hence no ARD dims. + assert all( + k.ard_num_dims == len(k.active_dims) + for k in leaves + if k.ard_num_dims is not None + ) @pytest.mark.parametrize( @@ -330,6 +340,11 @@ def test_task_parameter_equivalence(values, mode, expected): ("mode", "expected"), [ param(None, None, id="none"), + param( + TransferLearningMode.IDENTITY, + IdentityKernel(parameter_names=("task",)), + id="identity", + ), param( TransferLearningMode.INDEX_KERNEL, IndexKernel(num_tasks=3, rank=3, parameter_names=("task",)), diff --git a/tests/test_rgpe_surrogate.py b/tests/test_rgpe_surrogate.py new file mode 100644 index 0000000000..b3134eba28 --- /dev/null +++ b/tests/test_rgpe_surrogate.py @@ -0,0 +1,265 @@ +"""Tests for the rank-weighted Gaussian process ensemble (RGPE) surrogate.""" + +from typing import Literal + +import pandas as pd +import pytest +import torch + +from baybe import Campaign +from baybe.exceptions import ( + IncompatibleSearchSpaceError, + IncompatibleSurrogateError, +) +from baybe.objectives import SingleTargetObjective +from baybe.parameters import NumericalContinuousParameter, TaskParameter +from baybe.parameters.enum import TransferLearningMode +from baybe.recommenders import BotorchRecommender, TwoPhaseMetaRecommender +from baybe.searchspace import SearchSpace +from baybe.surrogates import GaussianProcessSurrogate, RGPESurrogate +from baybe.targets import NumericalTarget + +_TARGET = "A" +_SOURCE = "B" + + +def _make_searchspace( + *, + mode: TransferLearningMode | None = None, + values=(_TARGET, _SOURCE), + active=(_TARGET,), +) -> SearchSpace: + """Build a single-task-parameter transfer-learning search space.""" + return SearchSpace.from_product( + [ + NumericalContinuousParameter("x", (0, 5)), + TaskParameter( + "task", + values=values, + active_values=active, + override_transfer_learning_mode=mode, + ), + ] + ) + + +def _make_measurements( + training_data: Literal["source", "target", "both"] = "both", +) -> pd.DataFrame: + """Build measurements for the source and/or target task. + + Uses a fixed lookup rather than the random fake-data utilities because the ranking + weights need controlled values and exact per-task point counts. + """ + lookup = pd.DataFrame( + { + "x": [1.0, 2.0, 3.0, 4.0, 0.5, 1.5, 2.5, 3.5], + "y": [1.0, 2.0, 3.0, 4.0, 1.2, 2.1, 2.9, 4.2], + "task": [_TARGET] * 4 + [_SOURCE] * 4, + } + ) + if training_data == "source": + return lookup[lookup["task"] == _SOURCE].reset_index(drop=True) + if training_data == "target": + return lookup[lookup["task"] == _TARGET].reset_index(drop=True) + return lookup + + +@pytest.fixture(name="objective") +def fixture_objective() -> SingleTargetObjective: + """A single-target maximization objective.""" + return NumericalTarget("y").to_objective() + + +@pytest.fixture(name="candidates") +def fixture_candidates() -> pd.DataFrame: + """Candidate points in experimental representation (target task).""" + return pd.DataFrame({"x": [0.75, 2.25, 4.5], "task": [_TARGET] * 3}) + + +def test_fit_and_posterior(objective, candidates): + """The fitted ensemble yields a valid posterior with normalized weights.""" + surrogate = RGPESurrogate(n_mc_samples=32) + surrogate.fit(_make_searchspace(), objective, _make_measurements("both")) + + # One source GP, one target GP, and normalized weights over both. + assert len(surrogate._source_gps) == 1 + assert surrogate._target_gp is not None + assert surrogate._weights.shape == (2,) + assert torch.all(surrogate._weights >= 0) + assert surrogate._weights.sum().item() == pytest.approx(1.0) + + posterior = surrogate.posterior(candidates) + assert posterior.mean.numel() == len(candidates) + assert torch.all(posterior.variance > 0) + assert torch.isfinite(posterior.mean).all() + + +def test_inner_gps_use_identity_mode_and_do_not_redispatch(objective): + """Inner GPs fit on the identity-mode space instead of re-dispatching to RGPE. + + The inner GPs run on the identity-mode space, whose task parameter no longer carries + the RGPE mode. This is what prevents infinite recursion in the factory dispatch. + """ + surrogate = RGPESurrogate(n_mc_samples=32) + surrogate.fit(_make_searchspace(), objective, _make_measurements("both")) + + for gp in (*surrogate._source_gps, surrogate._target_gp): + assert gp is not None + # A concrete GP was fitted rather than dispatched to another RGPE ensemble. + assert gp._delegate is None + assert gp._model is not None + # The inner space carries the identity mode, which is the recursion guard. + task_param = gp._searchspace._task_parameter + assert task_param is not None + assert ( + task_param.override_transfer_learning_mode is TransferLearningMode.IDENTITY + ) + + +def test_override_dispatch_matches_direct(objective, candidates): + """Selecting RGPE via the task override dispatches to an RGPE delegate.""" + measurements = _make_measurements("both") + + gp = GaussianProcessSurrogate() + gp.fit(_make_searchspace(mode=TransferLearningMode.RGPE), objective, measurements) + assert isinstance(gp._delegate, RGPESurrogate) + # The undelegated GP model is never built along the dispatch path. + assert gp._model is None + + # The dispatched delegate produces the same posterior as a direct RGPE surrogate + # (up to Monte Carlo noise in the weights). + direct = RGPESurrogate() + direct.fit(_make_searchspace(), objective, measurements) + + delegated_mean = gp.posterior(candidates).mean + direct_mean = direct.posterior(candidates).mean + assert torch.allclose(delegated_mean, direct_mean, atol=0.1) + + +def test_identity_mode_matches_naive_pooling(objective, candidates): + """IDENTITY mode predicts like a task-free GP trained on the pooled data.""" + measurements = _make_measurements("both") + + # IDENTITY mode: keep the task parameter but render it inert. + gp_identity = GaussianProcessSurrogate() + gp_identity.fit( + _make_searchspace(mode=TransferLearningMode.IDENTITY), objective, measurements + ) + + # Naive pooling: no task parameter, task label dropped from the data. + gp_naive = GaussianProcessSurrogate() + gp_naive.fit( + NumericalContinuousParameter("x", (0, 5)).to_searchspace(), + objective, + measurements.drop(columns="task"), + ) + + identity_post = gp_identity.posterior(candidates) + naive_post = gp_naive.posterior(candidates.drop(columns="task")) + + assert torch.allclose(identity_post.mean, naive_post.mean, atol=1e-5) + assert torch.allclose(identity_post.variance, naive_post.variance, atol=1e-5) + + +@pytest.mark.parametrize( + ("measurements", "target_fitted", "expected_weights"), + [ + pytest.param( + _make_measurements("source"), + False, + [1.0], + id="cold-start", + ), + pytest.param( + pd.concat( + [ + _make_measurements("source"), + pd.DataFrame({"x": [2.0], "y": [2.0], "task": [_TARGET]}), + ], + ignore_index=True, + ), + True, + [0.5, 0.5], + id="single-target-point", + ), + ], +) +def test_insufficient_target_data_uses_uniform_weights( + objective, candidates, measurements, target_fitted, expected_weights +): + """Without enough target points to rank, the ensemble weights all models uniformly. + + This covers both the cold start (no target data, averaging over the source models) + and the single-target-point case (too few points for the ranking loss). + """ + surrogate = RGPESurrogate(n_mc_samples=32) + surrogate.fit(_make_searchspace(), objective, measurements) + + assert (surrogate._target_gp is not None) == target_fitted + assert torch.allclose(surrogate._weights, torch.tensor(expected_weights)) + assert torch.isfinite(surrogate.posterior(candidates).mean).all() + + +@pytest.mark.parametrize( + ("searchspace", "training_data", "match"), + [ + pytest.param( + SearchSpace.from_product([NumericalContinuousParameter("x", (0, 5))]), + "both", + "task parameter", + id="no-task-parameter", + ), + pytest.param( + _make_searchspace(active=(_TARGET, _SOURCE)), + "both", + "exactly one active", + id="multiple-active-tasks", + ), + pytest.param( + _make_searchspace(), + "target", + "source task", + id="missing-source-measurements", + ), + ], +) +def test_invalid_transfer_learning_setup_is_rejected( + searchspace, training_data, match, objective +): + """RGPE rejects an invalid task configuration or missing source measurements.""" + with pytest.raises(IncompatibleSearchSpaceError, match=match): + RGPESurrogate().fit(searchspace, objective, _make_measurements(training_data)) + + +@pytest.mark.parametrize("via_override", [False, True], ids=["direct", "override"]) +def test_campaign_recommendation(objective, via_override): + """A campaign recommends target-task points using the RGPE surrogate.""" + searchspace = _make_searchspace( + mode=TransferLearningMode.RGPE if via_override else None + ) + surrogate = ( + GaussianProcessSurrogate() if via_override else RGPESurrogate(n_mc_samples=32) + ) + recommender = TwoPhaseMetaRecommender( + recommender=BotorchRecommender(surrogate_model=surrogate) + ) + campaign = Campaign(searchspace, objective, recommender) + campaign.add_measurements(_make_measurements("both")) + + recommendations = campaign.recommend(batch_size=2) + assert (recommendations["task"] == _TARGET).all() + + +def test_posterior_mean_function_rejects_delegate(objective): + """A dispatched GP has no single mean module and fails loudly if asked for one.""" + gp = GaussianProcessSurrogate() + gp.fit( + _make_searchspace(mode=TransferLearningMode.RGPE), + objective, + _make_measurements("both"), + ) + with pytest.raises(IncompatibleSurrogateError, match="not implemented"): + gp.posterior_mean_function( + _make_searchspace(), objective, _make_measurements("both") + )