Support `PositiveIndexKernel` and dispatching via `TaskParameter` by kalama-ai · Pull Request #728 · emdgroup/baybe

kalama-ai · 2026-01-16T16:07:46Z

Dispatching between BoTorch's PositiveIndexKernel and GPyTorch's IndexKernel for transfer learning.

New task_correlation Parameter on TaskParameter

Can specify correlation mode when creating a TaskParameter:

# Use PositiveIndexKernel (default)
task_param = TaskParameter(
    name="task",
    values=["source_task_1", "source_task_2", "target_task"],
    active_values=["target_task"],
    task_correlation=TaskCorrelation.POSITIVE, 
)

# Use IndexKernel
task_param = TaskParameter(
    name="task",
    values=["source_task_1", "source_task_2", "target_task"],
    active_values=["target_task"],
    task_correlation=TaskCorrelation.UNKNOWN,

Kernel Dispatching

The GaussianProcessSurrogate selects the kernel based on the task_correlation:

TaskCorrelation.POSITIVE → uses botorch.models.kernels.PositiveIndexKernel
TaskCorrelation.UNKNOWN → uses gpytorch.kernels.IndexKernel

Integrated both modes into benchmarks.

- PositiveIndexKernel will normalize the diagonal elements onf the target task - requires integer value of target task to identify index - only single index supported

…with PositiveIndexKernel - new property transfer_mode returning the TL mode of the searchpace, if a TaskParameter is provided (required for dispatching between two kernels) - new property target_task_idxs retunrning the indices of the active_values of the TaskParameter in its computational representation (this is required for normalization in PositiveIndexKernel)

…GaussianProcessSurrogate - add required properties (transfer_mode and target_task_idxs) to _ModelContext - dispatch between two kernels given transfer_mode - if transfer_mode is `joint_pos` (use PositiveIndexKernel) we implicitely assume only one active_value for identifying the target_task (a wrong configuration will raise an error in TaskParameter)

- `TransferMode` was replaced by `TaskCorrelation` because it was hard to understand without knowing about the kernels - positive correlation -> use PositiveIndexKernel - unknown correlation -> use IndexKernel since it might be more robust

Scienfitz · 2026-02-04T16:46:02Z

baybe/parameters/categorical.py

    # See base class.

+    task_correlation: TaskCorrelation = field(default=TaskCorrelation.POSITIVE)
+    """Task correlation. Defaults to positive correlation via PositiveIndexKernel."""


Suggested change

"""Task correlation. Defaults to positive correlation via PositiveIndexKernel."""

"""Task correlation influencing which kernel will be used du default for task parameters."""

Scienfitz · 2026-02-04T16:48:30Z

baybe/parameters/categorical.py

+        """
+        # Check POSITIVE constraint: must have exactly one active value
+        # Note: _active_values is the internal field, could be None
+        if value == TaskCorrelation.POSITIVE and self._active_values is not None:


use value is TaskCorrelation.POSITIVE (always check sentinels via is and never ==)

Why are you suing ._active_values and not .active_values? The latter takes care of defaulting to values if user does not specify anything and can never become None

Scienfitz · 2026-02-04T16:50:28Z

baybe/parameters/categorical.py

+            if len(self._active_values) > 1:
+                raise ValueError(
+                    f"Task correlation '{TaskCorrelation.POSITIVE.value}' requires "
+                    f"one active value, but {len(self._active_values)} were provided: "


Suggested change

f"one active value, but {len(self._active_values)} were provided: "

f"exactly one active value, but {len(self._active_values)} were provided: "

Scienfitz · 2026-02-04T16:58:03Z

baybe/searchspace/core.py

            return 1

+    @property
+    def target_task_idxs(self) -> list[int] | None:


would always prefer returning tuples such cases unless there is a limitation that it really must be a list

Scienfitz · 2026-02-04T17:00:51Z

baybe/surrogates/gaussian_process/core.py

+    @property
+    def task_correlation(self) -> TaskCorrelation | None:
+        """Get the task correlation mode of the task parameter, if available."""
+        return self.searchspace.task_correlation
+
+    @property
+    def target_task_idxs(self) -> list[int] | None:
+        """Determine target task index for PositiveIndexKernel normalization."""
+        return self.searchspace.target_task_idxs
+


how necessary are these helpers? I can get them just via gp.searchspace.x which is not tremendously worse than just gp.x

Scienfitz · 2026-02-04T17:02:20Z

baybe/surrogates/gaussian_process/core.py

+        elif context.task_correlation == TaskCorrelation.POSITIVE:
+            task_covar_module = (
+                botorch.models.kernels.positive_index.PositiveIndexKernel(
+                    num_tasks=context.n_tasks,
+                    active_dims=context.task_idx,
+                    rank=context.n_tasks,  # TODO: make controllable
+                    target_task_index=context.target_task_idxs[0],
+                )
+            )
+            covar_module = base_covar_module * task_covar_module
+        elif context.task_correlation == TaskCorrelation.UNKNOWN:


just for our common understanding: these parts will eventually have to e outsourced to a default_task_kernel_factory or similar (not needed in this PR)

Scienfitz · 2026-02-04T17:05:53Z

benchmarks/definition/regression/core.py

+                    "source_data_seed": settings.random_seed + mc_iter,
+                }
+                result.update(metrics)
+                results.append(result)


since you expanded the benchmarks: are they still feasible or are they now timing out due to the longer runtime?

Scienfitz · 2026-02-04T17:07:09Z

benchmarks/domains/aryl_halides/core.py

+        data: The benchmark data.
+        target_tasks: The target tasks for transfer learning.
+        source_tasks: The source tasks for transfer learning.
+        task_correlation: The task correlation mode (UNKNOWN or POSITIVE).


please do not hardcode possible enum values here in such comments (not well maintainable) - maybe link the actual enum

Scienfitz · 2026-02-04T17:10:50Z

benchmarks/domains/easom/convergence_tl.py

+        active_values=["Target_Function"],
+        task_correlation=TaskCorrelation.POSITIVE,
+    )
+    params_tl_index = params + [task_param_index]


your approach ehre to expand the benchmakrs is to copy the code and make 1 entry for POSITIVE and one entry for UNKNOWN

Instead you could loop over the possible values of TaskCorreclation and automatically create as many searchspaces with autogenerated names etc.

That would have the major advanatge that youd never have to touch this code again to incorporate choices that might be added in the future

Scienfitz · 2026-02-04T17:15:55Z

baybe/parameters/categorical.py

    encoding: CategoricalEncoding = field(default=CategoricalEncoding.INT, init=False)
    # See base class.

+    task_correlation: TaskCorrelation = field(default=TaskCorrelation.POSITIVE)


the only big potential problem with this PR I cans pot is the naming of this attribute

In isolation the name is totally accurate and fine. But we already have plans to expand this attribute so have potentially more choices, like eg RGPE, MeanTransfer, CovarTransfer etc (with names yet to be decided) -> the name correlation is then not appropriate anymore. Instead this attribute embodies something like TL_MODE or TL_METHOD or TL_ALGORITHM.

Now of course we could change the name of the attribute later, but since this is merged to main and potentially released before we have the other choices, we would introduce a breaking change that has tobe deprecated. So it would be beneficial if we would avoid that situation.

Here two proposals how to do that:

make this attribute private for now indicating to users that its not fully public and can change at any moment

already now decide on the attribute name, which should be doable because it will have to be a rather generic one (see proposals above)

@AdrianSosic do you agree with this issue of the attribute name?

AVHopp · 2026-03-05T08:00:53Z

baybe/parameters/categorical.py

@kalama-ai @AdrianSosic can you quickly comment on the state of this PR? If I remember correctly, this was one of the PRs that are somewhat depending on the current refactoring work of Adrian. Has this code here already been rebased and is thus ready to review? Or do I misremember?

kalama-ai added 7 commits January 16, 2026 12:00

Add transfer learning mode to TaskParameter

cc9803d

Raise errro if PositiveIndexKernel is used with multiple sctive_values

545727d

- PositiveIndexKernel will normalize the diagonal elements onf the target task - requires integer value of target task to identify index - only single index supported

Add PositiveIndexKernel baseline to benchmarks

32c5167

Improve comments in benchmarks

45ece28

kalama-ai requested review from AVHopp, AdrianSosic and Scienfitz as code owners January 16, 2026 16:07

Scienfitz assigned kalama-ai Jan 21, 2026

Scienfitz requested changes Feb 4, 2026

View reviewed changes

AVHopp reviewed Mar 5, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support `PositiveIndexKernel` and dispatching via `TaskParameter`#728

Support `PositiveIndexKernel` and dispatching via `TaskParameter`#728
kalama-ai wants to merge 7 commits intomainfrom
feat/support-positive-index-kernel

kalama-ai commented Jan 16, 2026

Uh oh!

Scienfitz Feb 4, 2026

Uh oh!

Scienfitz Feb 4, 2026

Uh oh!

Scienfitz Feb 4, 2026

Uh oh!

Scienfitz Feb 4, 2026

Uh oh!

Scienfitz Feb 4, 2026

Uh oh!

Scienfitz Feb 4, 2026

Uh oh!

Scienfitz Feb 4, 2026

Uh oh!

Scienfitz Feb 4, 2026

Uh oh!

Scienfitz Feb 4, 2026 •

edited

Loading

Uh oh!

Scienfitz Feb 4, 2026 •

edited

Loading

Uh oh!

AVHopp Mar 5, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

	"""Task correlation. Defaults to positive correlation via PositiveIndexKernel."""
	"""Task correlation influencing which kernel will be used du default for task parameters."""

	f"one active value, but {len(self._active_values)} were provided: "
	f"exactly one active value, but {len(self._active_values)} were provided: "

Conversation

kalama-ai commented Jan 16, 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!

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!

Scienfitz Feb 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Scienfitz Feb 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Scienfitz Feb 4, 2026 •

edited

Loading

Scienfitz Feb 4, 2026 •

edited

Loading