Transfer Learning Decorator by AdrianSosic · Pull Request #790 · emdgroup/baybe

AdrianSosic · 2026-05-07T07:55:32Z

DevPR, parent is #745

Last piece to the puzzle:
Presets (i.e. papers, packages, etc) can dictate certain aspects of the GP model while not saying anything about other aspects. For example, both EDBO and CHEN focus on the kernel priors but don't even consider transfer learning at all. This is a general issue, and can also cover other things like multi-fidelity etc.

For these cases, we want to follow the approach if not defined, use BayBE default mechanism/setting. However, this requires to abstract these settings/mechanism into reusable structures. This PR takes care of this step for transfer learning (which is currently the only mechanism that needs to be ported) in the form of a class decorator. Because other mechanisms will follow in the future and their extent isn't yet fully clear (e.g. multi-fidelity or transfer learning via mean injection), we keep this decorator private for now. A possible future extension of the decorator could have the form @enable_mechanism(transfer_learning=True, multi_fidelity=True) that then accepts any existing GP component and makes the necessary adjustments.

Copilot

Pull request overview

This PR introduces a class-decorator-based mechanism to enable BayBE’s default transfer learning behavior for Gaussian Process kernel factories (automatically composing a base kernel with the default task kernel when a task parameter is present), and refactors several kernel presets to use it.

Changes:

Added _enable_transfer_learning decorator in the GP kernel factory component layer.
Updated preset kernel factories (EDBO, CHEN, Smoothed EDBO, BayBE) to rely on the decorator rather than per-preset task-exclusion selectors / manual orchestration.
Adjusted kernel-factory tests and internal factory naming to reflect the new split between numerical-only and task-aware factories.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
`baybe/surrogates/gaussian_process/components/kernel.py`	Adds `_enable_transfer_learning` and updates ICM defaults to use internal numerical/task factories.
`baybe/surrogates/gaussian_process/presets/edbo.py`	Applies transfer-learning decorator and removes the preset’s default task-excluding selector.
`baybe/surrogates/gaussian_process/presets/edbo_smoothed.py`	Splits numerical kernel factory from transfer-learning-enabled alias.
`baybe/surrogates/gaussian_process/presets/chen.py`	Applies transfer-learning decorator and removes the preset’s default task-excluding selector.
`baybe/surrogates/gaussian_process/presets/baybe.py`	Refactors default kernel factory wiring to reuse Smoothed EDBO and internal task factory.
`tests/test_kernel_factories.py`	Updates imports/usages to the new internal factory names.

Comments suppressed due to low confidence (1)

baybe/surrogates/gaussian_process/presets/edbo_smoothed.py:49

_SmoothedEDBONumericalKernelFactory._make uses effective_dims = train_x.shape[-1] and passes parameter_names=self.get_parameter_names(searchspace) (which is None unless a selector is provided). Under transfer learning/multitask search spaces, this will count the task index dimension in effective_dims and make the base kernel operate on the task column as well, altering the smoothed prior interpolation. The numerical kernel used as the ICM base kernel should exclude TaskParameter by default, and dimension logic should be computed from the selected active dims/parameter names (not the raw train_x width).

@define
class _SmoothedEDBONumericalKernelFactory(_PureKernelFactory):
    """A factory providing the core numerical kernel for the smoothed EDBO preset."""

    _uses_parameter_names: ClassVar[bool] = True
    # See base class.

    @override
    def _make(
        self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
    ) -> Kernel:
        effective_dims = train_x.shape[-1]

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

CLAassistant · 2026-05-07T15:07:09Z

All committers have signed the CLA.

AdrianSosic · 2026-05-07T15:45:56Z

@copilot: please review

Copilot · 2026-05-07T15:47:14Z

@copilot: please review

I've completed a comprehensive code review of the transfer learning decorator implementation. Here are my findings:

✅ Strengths

Improved decorator design: The decorator now creates a proper subclass instead of modifying the class in-place, which is cleaner and avoids mutation issues.
Correct parameter scoping: The decorator temporarily modifies parameter_selector to exclude task parameters before calling the base factory, ensuring dimension-dependent logic operates only on numerical features.
Proper dimensionality calculation: The new _get_effective_dimensionality method correctly computes dimensions based on selected parameters, accounting for the parameter selector.
Clean factory definitions: SmoothedEDBOKernelFactory and BayBEKernelFactory are now properly defined with correct class names for serialization.

🔍 Observations

Decorator signature: The name parameter allows explicit control over the created class name, which is useful for the SmoothedEDBOKernelFactory case.
Sanity check mechanism: The temporary narrowing of _supported_parameter_kinds (lines 168-169, 182) acts as a safety check to ensure the base factory never sees task parameters.
Selector composition: The decorator correctly composes the original selector with the task-exclusion selector (lines 172-178).

✨ Code Quality

The implementation is well-structured, properly documented, and follows BayBE's patterns. The decorator successfully abstracts transfer learning while maintaining backward compatibility and correct serialization behavior.

No issues found. The implementation is production-ready.

AdrianSosic · 2026-05-07T18:39:23Z

        """Construct the kernel."""


+def _enable_transfer_learning(


@Scienfitz: in principle ready and working. However, I have to admit that this was significantly more painful than anticipated, with many footguns along the way. So I'm open to a very harsh review and a complete change of direction, if you prefer and have an alternative/simpler idea.

But I hope that you get my intent for this: I think we need some mechanism that lets us say fill this preset with our default approach for a certain aspect that the preset does not specify, and the filling should be very much done without copying code since the BayBE defaults are expected to move. So we need something like a single source of truth. That said: maybe you have some smarter idea.

only annoying ting is the _NAME / NAME thing but I dont know a better alternative

Scienfitz · 2026-05-18T08:48:41Z

@AdrianSosic after our chat you were not sure whether this PR is still relevant, now seeing it still there I assume its still relevant and should be reviewed?

kalama-ai · 2026-05-20T10:18:50Z

+        if original_selector is None:
+            self.parameter_selector = _task_exclude_selector
+        else:
+            self.parameter_selector = lambda p: (


The TaskParameter will be excluded here if the orgininal parameter_selector contained it. Should we maybe raise a warning?

Hm, maybe I'm looking at this from the wrong perspective, but isn't that what we want to happen?

Example: My search space contains 3 parameters (2 regular and 1 task): (p1, p2, p_task). If I now call a factory that is TL-compatible (which the factories that get decorated de facto are), then two things can happen:

I pass a selector that excludes the task parameter, e.g. filters down to just p1. In this case, I want that the inner logic gets executed, i.e. the execution internally dispatches to the original non-decorated factory.

I pass a selector that includes the task parameter, e.g. selects p1 and task. In this case, task should be split off, the regular factory should be called on p1, and then we assemble everything via the ICM mechanism.

Or am I overlooking something? If you see a problem, a minimum example would be helpful

kalama-ai · 2026-05-20T10:24:14Z

+            target_cls._supported_parameter_kinds = broadened_kinds  # type: ignore[attr-defined]
+            self.parameter_selector = original_selector
+
+        if searchspace.task_idx is not None:


I can't oversee 100% where we would put the logic for TL modes that affect other components than the kernel later. But, the current version will definetely work well with the dispatching between IndexKernel and PositiveIndexKernel, which will happen within the ICMKernelFactory.

In my current naive vision, the decorator would then not just be applied to the kernel factor but to all components that require changes. It would process all of them and apply the necessary changes, e.g.

for kernel-TL, it would do exactly what we do no, i.e. only modify the kernel factory

for mean-injection-TL, it would not affect the kernel but alter the mean-factory to do the prior-mean stuff

...

However, this is exactly what I meant with I can't yet foresee exactly, so let's keep the decorator private for now :D But yeah, potentially the decorator logic might become quite a beast, since it not only needs to capture TL stuff across all components, but maybe also multi-fidelity stuff etc in the end. My current fear is that the complexity might get out of hands...

Any action items at this point I need to take care of, or can I close?

Scienfitz · 2026-05-20T17:16:46Z

-
-BayBENumericalKernelFactory = SmoothedEDBOKernelFactory
-"""The factory providing the default numerical kernel for Gaussian process surrogates."""  # noqa: E501
+class BayBEKernelFactory(SmoothedEDBOKernelFactory):  # type: ignore[valid-type, misc]


I think I roughly get the need for this (horrendously feeling) _NAME / NAME split for some of the objects (to which I couldnt come up witha better idea either)

But could you sum tis up in 1-2 sentences in the PR description? It seems the commits also have explanation but it seemingly changed during the course of them and they don't really get the point across very nicely

* Provides a single source of truth for defining the TL logic * Enables TL for non-TL presets by applying the decorator

…ator

`_enable_transfer_learning` now accepts an optional `name` parameter so that the dynamically created class can have the correct `__name__` when the function is called directly (rather than used as a decorator). This fixes serialization for `SmoothedEDBOKernelFactory`, which was previously serialized as `_SmoothedEDBONumericalKernelFactory`.

Simple aliases like `BayBEKernelFactory = SmoothedEDBOKernelFactory` cause the serialized type name to be that of the underlying class, which means the identity is lost on deserialization. Using thin subclasses ensures each factory has its own stable `__name__`.

@Define-

When used as a decorator (@_enable_transfer_learning), modify the class in-place instead of creating a subclass with the same __name__. The previous approach left two concrete classes with identical names in the subclass registry, causing find_subclass to resolve to the @Define- processed intermediate (without the TL wrapper) during deserialization. When called with an explicit name argument (for cases like SmoothedEDBOKernelFactory where the original class is reused elsewhere), the subclass approach is preserved since the distinct name avoids any collision.

The Protocol metaclass (_ProtocolMeta) defaults __module__ to 'abc' when creating classes via 3-arg type(). Set it explicitly from the parent class so that SmoothedEDBOKernelFactory correctly reports its module as baybe.surrogates.gaussian_process.presets.edbo_smoothed.

…learning

Scienfitz · 2026-05-21T15:31:51Z

+        # __module__ must be set explicitly because the Protocol metaclass
+        # would otherwise default it to "abc".
+        target_cls = type(
+            name, (cls,), {"__doc__": cls.__doc__, "__module__": cls.__module__}


isnt it strange that it is made a subclass? in a sense it should be an equivalent copy, not a subclass

Fully agree. It's just that the subclassing mechanism is a bit less works since it automatically populates the attributes and stuff. But if you prefer the conceptually cleaner approach: 1ccad46

Shall I keep it?

cool, if it works I'd prefer it 👍

Scienfitz · 2026-05-21T15:37:19Z

        )

-        return BayBETaskKernelFactory(TypeSelector((TaskParameter,)))
+        assert _BayBETaskKernelFactory._supported_parameter_kinds is _ParameterKind.TASK


are these new asserts in the defaults for mypy or for another purpose?

No, not for mypy, but to guarantee that the defaults that are injected here actually are compatible with the contract that the ICM mechanism requires. However, this is probably from the earlier drafting days and should actually become a proper validation step. Would it be fine for you if I turned this into actual validators?

I was commenting due to this simple principle: asserts for anything other than pytest or mypy type-narrowing should be avoided, in the latter case ideally they'd get a comment so we dont have to always redo this kind of thread here

…orProtocol Extends the method to accept either a parameter name (existing behavior) or a selector, returning the combined comp-rep indices of all matching parameters. Uses this to simplify _get_effective_dimensionality in _PureKernelFactory and the inline dimensionality sums in the EDBO likelihood factories.

Parameter name uniqueness is already enforced as a searchspace invariant, making the multi-match case impossible. The no-match case now returns an empty tuple, consistent with the selector path and the existing behavior for parameters absent from the comp-rep.

Co-authored-by: Martin Fitzner <martin.fitzner@merckgroup.com>

Scienfitz

fix CI please

AdrianSosic self-assigned this May 7, 2026

AdrianSosic added refactor dev labels May 7, 2026

AdrianSosic changed the base branch from main to dev/gp May 7, 2026 07:55

AdrianSosic marked this pull request as ready for review May 7, 2026 07:56

AdrianSosic requested a review from Scienfitz as a code owner May 7, 2026 07:56

Copilot AI review requested due to automatic review settings May 7, 2026 07:56

AdrianSosic requested a review from AVHopp as a code owner May 7, 2026 07:56

AdrianSosic marked this pull request as draft May 7, 2026 07:56

Copilot started reviewing on behalf of AdrianSosic May 7, 2026 07:56 View session

Copilot AI reviewed May 7, 2026

View reviewed changes

Copilot started work on behalf of AdrianSosic May 7, 2026 14:48 View session

Copilot finished work on behalf of AdrianSosic May 7, 2026 15:09

Copilot started work on behalf of AdrianSosic May 7, 2026 15:10 View session

Copilot finished work on behalf of AdrianSosic May 7, 2026 15:12

Copilot started work on behalf of AdrianSosic May 7, 2026 15:43 View session

AdrianSosic force-pushed the feature/tl_decorator branch from 6b3abe7 to a713bfe Compare May 7, 2026 15:44

emdgroup deleted a comment from Copilot AI May 7, 2026

Copilot stopped work on behalf of AdrianSosic due to an error May 7, 2026 15:45
The session was cancelled by the user.

Copilot started work on behalf of AdrianSosic May 7, 2026 15:46 View session

Copilot finished work on behalf of AdrianSosic May 7, 2026 15:48

AdrianSosic marked this pull request as ready for review May 7, 2026 18:36

AdrianSosic commented May 7, 2026

View reviewed changes

Scienfitz added this to the 0.15.0 milestone May 19, 2026

kalama-ai reviewed May 20, 2026

View reviewed changes

Scienfitz requested changes May 20, 2026

View reviewed changes

AdrianSosic added 9 commits May 21, 2026 09:05

Extract transfer learning mechanism into a reusable decorator

0611ff6

* Provides a single source of truth for defining the TL logic * Enables TL for non-TL presets by applying the decorator

Scope inner factory to non-task parameters in transfer learning decor…

43d1e8e

…ator

Add serialization roundtrip tests for kernel factories

ded217f

Add SerialMixin to _PureKernelFactory and move serialization test

21ed537

Suppress mypy errors from dynamic class creation in _enable_transfer_…

41a1fc7

…learning

AdrianSosic force-pushed the feature/tl_decorator branch from 5229f7f to 41a1fc7 Compare May 21, 2026 07:06

Scienfitz reviewed May 21, 2026

View reviewed changes

AdrianSosic and others added 4 commits May 22, 2026 16:00

Add blank line for readability

4d75d5a

Co-authored-by: Martin Fitzner <martin.fitzner@merckgroup.com>

Create sibling class instead of subclass in _enable_transfer_learning

1ccad46

Scienfitz approved these changes May 22, 2026

View reviewed changes

Conversation

AdrianSosic commented May 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

CLAassistant commented May 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

AdrianSosic commented May 7, 2026

Uh oh!

Copilot AI commented May 7, 2026

✅ Strengths

🔍 Observations

✨ Code Quality

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Scienfitz commented May 18, 2026

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!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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!

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!

Scienfitz left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

AdrianSosic commented May 7, 2026 •

edited

Loading

CLAassistant commented May 7, 2026 •

edited

Loading