Add RGPE TL Mode and Identity Kernel by kalama-ai · Pull Request #921 · emdgroup/baybe

kalama-ai · 2026-09-24T12:34:23Z

Transfer learning: `IDENTITY` and `RGPE` modes

Two new options for TaskParameter.override_transfer_learning_mode building on the parameter kernel-override machinery from #904. IDENTITY is backed by a new public IdentityKernel.

`IDENTITY`

Uses a new public IdentityKernel (baybe.kernels). It contributes no covariance on the task dimension, so every task collapses into a single pooled model (equivalent to the "naive" baseline where the task label is effectively ignored). Like the (POSITIVE_)INDEX_KERNEL modes, it's a plain kernel override and treated the same way in the factory.

An IDENTITY search space is a genuine single-task problem, so it is fit with the marginal log-likelihood rather than the leave-one-out pseudo-likelihood used for the true multi-task kernels.

`RGPE`

Instead of modelling task correlations with a kernel, it trains one GP per source task plus one target GP and combines their posteriors with weights derived from a leave-one-out ranking loss. With fewer than two target points there's nothing to rank, so it falls back to a uniform average. With no target data it averages the source models alone.

Each inner GP is fit over the same search space under the IDENTITY TL override (searchspace._with_task_mode(IDENTITY)), so the task dimension is inert and every inner model is a clean single-task GP over the shared features. Because the inner models are just evolved copies of the base surrogate, any parameter-specific kernel overrides carry through unchanged (supports_kernel_overrides = True).

Dispatching in the GP factory

GaussianProcessSurrogate._fit checks the task parameter's override. IDENTITY is just a kernel override and treated like (POSITIVE_)INDEX_KERNEL overrides. RGPE can't be expressed as a kernel, so the factory builds an RGPESurrogate from the current GP's own configuration (via evolve) and stashes it as a _delegate; to_botorch, _posterior, etc. then forward to it. This keeps all the usual GP settings (kernel, priors, scaling) for the inner models.

Benchmarks

The transfer-learning convergence benchmarks previously built their "naive" baseline by dropping the TaskParameter from the search space entirely. They now build it from a TaskParameter with override_transfer_learning_mode="IDENTITY" instead.

Tests

tests/test_rgpe_surrogate.py — weighting, cold-start and <2-target fallbacks,
source/target splitting, single-output/single-task guards, serialization roundtrip,
and IDENTITY-matches-naive-pooling equivalence.
tests/test_parameter_kernel_overrides.py — IDENTITY override and its interaction
with kernel overrides.
Serialization and hypothesis-strategy coverage extended for the new surrogate and
IdentityKernel.

Copilot

Copilot review overview

🟡 Changes recommended

Shallow component sharing and stale delegate state can produce incorrect models, while the delegated return type also violates its API annotation.

Review effort: Balanced
Findings: 1 · 1 · 1

Open (3)

Deep-copy evolved GP configuration for independent source models · New
Clear stale RGPE delegate when refitting in another context · New
Widen to_botorch return annotation to generic BoTorch Model · New

What changed in this PR

Adds IDENTITY and RGPE transfer-learning modes, including public kernel/surrogate APIs, GP dispatch, benchmarks, documentation, and tests.

Changes:

Implements pooled identity-kernel modeling and rank-weighted GP ensembles.
Dispatches GP surrogates to RGPE and supports composite BoTorch models.
Updates benchmarks, serialization strategies, documentation, and coverage.

File	Description
`baybe/kernels/basic.py`	Adds `IdentityKernel`.
`baybe/kernels/__init__.py`	Exports the identity kernel.
`baybe/parameters/enum.py`	Adds `IDENTITY` and `RGPE` modes.
`baybe/parameters/categorical.py`	Resolves task-mode overrides.
`baybe/searchspace/core.py`	Adds task-mode cloning.
`baybe/surrogates/transfer_learning/rgpe.py`	Implements RGPE fitting and prediction.
`baybe/surrogates/transfer_learning/__init__.py`	Exports RGPE.
`baybe/surrogates/gaussian_process/core.py`	Adds RGPE delegation.
`baybe/surrogates/gaussian_process/presets/baybe.py`	Selects MLL for identity mode.
`baybe/surrogates/composite.py`	Supports generic BoTorch model lists.
`baybe/surrogates/__init__.py`	Publishes RGPE.
`tests/test_rgpe_surrogate.py`	Tests RGPE behavior and dispatch.
`tests/test_parameter_kernel_overrides.py`	Tests identity overrides.
`tests/test_iterations.py`	Adjusts iteration coverage.
`tests/serialization/test_surrogate_serialization.py`	Tests RGPE serialization.
`tests/hypothesis_strategies/surrogates.py`	Adds RGPE generation.
`tests/hypothesis_strategies/kernels.py`	Adds identity-kernel generation.
`benchmarks/domains/michalewicz/convergence_tl.py`	Uses identity-mode baseline.
`benchmarks/domains/hartmann/convergence_tl.py`	Uses identity-mode baseline.
`benchmarks/domains/easom/convergence_tl.py`	Uses identity-mode baseline.
`benchmarks/domains/direct_arylation/convergence_tl.py`	Uses identity-mode baseline.
`benchmarks/domains/aryl_halides/core.py`	Uses identity-mode baseline.
`docs/concepts/transfer_learning.md`	Documents ensemble transfer learning.
`docs/references.bib`	Adds the RGPE citation.
`CHANGELOG.md`	Records the new APIs and benchmarks.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

+        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)


+        task_param = self._searchspace._task_parameter
+        if (
+            task_param is not None
+            and task_param.override_transfer_learning_mode is TransferLearningMode.RGPE
+        ):


    def to_botorch(self) -> GPyTorchModel:
+        if self._delegate is not None:
+            return self._delegate.to_botorch()


AVHopp

Initial look into some of the architecture - did in particular not look at the surrogate itself so far, because I feel that we first need to talk about the design regarding _model and _delegate

AVHopp · 2026-10-05T11:35:21Z

+            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."
+            )


Is this not possible or not implemented? Because depending on this, I'd say that a NotImplementedError might be more suitable.

AVHopp · 2026-10-05T11:38:46Z

+        # 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


DO we need that check here or is this already guaranteed since we are in the else branch? In the second case, I would prefer to rather have a full "if-else" here with an explicit assert on task_param before that check, as this would read a bit nicer imo (but only my preference, so feel free to just resolve if you disagree)

AVHopp · 2026-10-05T11:40:12Z

    """:class:`botorch.models.kernels.positive_index.PositiveIndexKernel` for positive
    correlations."""
+
+    RGPE = "RGPE"


Maybe RGPE_SURROGATE to be more in the same style as ..._KERNEL?

AVHopp · 2026-10-05T11:40:36Z

 class TransferLearningMode(Enum):
    """Transfer learning modes for :class:`.TaskParameter`."""

+    IDENTITY = "IDENTITY"


Should this maybe be renamed into IDENTITY_KERNEL or CONSTANT_KERNEL such that it fits the other members?

AVHopp · 2026-10-05T14:10:29Z

+            delegate = RGPESurrogate(base_surrogate=evolve(self))
+            delegate.fit(self._searchspace, self._objective, self._measurements)
+            self._delegate = delegate
+            return


Should this already return at this point? Don't we then loose some of the errors and warning regarding deprecations, misconfigurations and so on silently?

AVHopp · 2026-10-05T14:16:25Z


+    # 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)


I am sure that you investigated this, but why can't we simply use _model? From what I see, there is basically never a situation where we actually need both _model and _delegate, and also semantically, the RGPE model is the model - so what is the reason for this separation?

AVHopp · 2026-10-05T14:21:08Z

@@ -381,6 +404,8 @@ def posterior_mean_function(

    @override
    def to_botorch(self) -> GPyTorchModel:


This holds for a lot of the functions here: If I would use the same GaussianProcessSurrogate object for different spaces, it could happen that both _delegate and _model are being set and used. However, as soon as _delegate has been set, the object will always return with respect to this model. Questionable if this is a misuse of the GaussianProcessSurrogate and whether or not it is our responsibility to do something regarding this, but I at least want to flag it.

If we are able to somehow get rid of this duality regarding _model and _delegate, this issue might also just resolve itself.

AVHopp · 2026-10-05T14:22:27Z

+        source_tasks=source_tasks,
+        target_tasks=target_tasks,
+        transfer_learning_mode=TransferLearningMode.IDENTITY,


Why did you explicitly add source_tasks and target_tasks here, but not in some of the other benchmarks?

AVHopp · 2026-10-05T14:24:31Z

+    return st.builds(
+        RGPESurrogate,
+        base_surrogate=gaussian_process_surrogates(),
+        n_mc_samples=st.integers(min_value=1, max_value=512),


Why such a high max_value? Is there a special reason for such a variety of values or would 1 and >1 already suffice?

- constant task kernel that leaves the base kernel untouched, pooling all tasks

- rank-weighted ensemble of per-task GPs, usable directly or via the RGPE task override - inner GPs fit on the identity-mode space so the task dimension stays inert

kalama-ai force-pushed the feature/rgpe-with-kernel-override branch 5 times, most recently from ba44509 to a7d73c2 Compare September 24, 2026 18:50

Scienfitz assigned kalama-ai Sep 25, 2026

kalama-ai force-pushed the feature/rgpe-with-kernel-override branch from a7d73c2 to abcd67d Compare September 28, 2026 10:11

kalama-ai changed the base branch from main to feature/param_kernel_override September 28, 2026 10:35

kalama-ai added this pull request to stack #905 September 28, 2026 10:36

kalama-ai force-pushed the feature/rgpe-with-kernel-override branch from d9a7e30 to 584a25d Compare September 28, 2026 17:51

kalama-ai marked this pull request as ready for review September 29, 2026 10:05

kalama-ai requested review from AVHopp, AdrianSosic and Scienfitz as code owners September 29, 2026 10:05

Copilot AI balanced review requested due to automatic review settings September 29, 2026 10:05

kalama-ai marked this pull request as draft September 29, 2026 10:05

Copilot started reviewing on behalf of kalama-ai September 29, 2026 10:05 View session

kalama-ai marked this pull request as ready for review September 29, 2026 10:06

Copilot AI reviewed Sep 29, 2026

View reviewed changes

Scienfitz added the new feature New functionality label Sep 29, 2026

Scienfitz changed the title ~~Feature/rgpe with kernel override~~ Add RGPE TL Mode and Identity Kernel Sep 29, 2026

AVHopp requested changes Oct 5, 2026

View reviewed changes

Scienfitz added this to the 0.16.0 milestone Oct 7, 2026

kalama-ai added 6 commits October 7, 2026 20:04

Add IDENTITY transfer-learning mode

ab9258d

- constant task kernel that leaves the base kernel untouched, pooling all tasks

Add RGPE transfer-learning surrogate

b9095db

- rank-weighted ensemble of per-task GPs, usable directly or via the RGPE task override - inner GPs fit on the identity-mode space so the task dimension stays inert

Document transfer-learning modes

1ce60ca

Update CHANGELOG

3c08bc7

Fix IDENTITY transfer learning mode to use MLL instead of LOOCV

35443bd

Use IDENTITY task kernel for naive transfer-learning benchmark baseline

fb6338a

Scienfitz force-pushed the feature/rgpe-with-kernel-override branch from 584a25d to fb6338a Compare October 7, 2026 18:04

Navigation Menu

Repository navigation

Add RGPE TL Mode and Identity Kernel - #921

Add RGPE TL Mode and Identity Kernel#921
kalama-ai wants to merge 6 commits into
feature/param_kernel_overridefrom
feature/rgpe-with-kernel-override

kalama-ai commented Sep 24, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

AVHopp left a comment

Uh oh!

AVHopp Oct 5, 2026

Uh oh!

AVHopp Oct 5, 2026

Uh oh!

AVHopp Oct 5, 2026

Uh oh!

AVHopp Oct 5, 2026

Uh oh!

AVHopp Oct 5, 2026

Uh oh!

AVHopp Oct 5, 2026

Uh oh!

AVHopp Oct 5, 2026

Uh oh!

AVHopp Oct 5, 2026

Uh oh!

AVHopp Oct 5, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		@@ -381,6 +404,8 @@ def posterior_mean_function(

		@override
		def to_botorch(self) -> GPyTorchModel:

Conversation

kalama-ai commented Sep 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Transfer learning: IDENTITY and RGPE modes

IDENTITY

RGPE

Dispatching in the GP factory

Benchmarks

Tests

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

AVHopp left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

kalama-ai commented Sep 24, 2026 •

edited

Loading

Transfer learning: `IDENTITY` and `RGPE` modes

`IDENTITY`

`RGPE`