Add basic `GaussianProcessSurrogate.from_prior` constructor method by kalama-ai · Pull Request #717 · emdgroup/baybe

kalama-ai · 2025-12-18T14:37:25Z

New constructor method that enables transfer learning for Gaussian Process surrogates

# Train source GP
source_gp = GaussianProcessSurrogate()
source_gp.fit(source_space, source_objective, source_data)

# Transfer mean
target_gp = GaussianProcessSurrogate.from_prior(
      prior_gp=source_gp,  # Use source GP as prior
  )

Mean Function Transfer:

Implements full mean transfer from the prior GP to the new GP
The posterior mean predictions of the pre-trained GP used as the mean module for the new GP
Frozen hyperparameters
Mean function is evaluated at source training points
New PriorMean class that wraps the prior GP as a BoTorch-compatible mean module

In the upcoming PR we will introduce a new TL surrogate that takes a search space with TaskParameter and internally falls back to the new constructor for training a source GP on the source data and a target GP on the target data using the source mean as a prior.

Further extensions:

the interface can be extended to enhanced mean transfer (initializing the HPs, evaluation at the target) and covariance transfer
this will be happening in later PRs once the other configurations are tested and the dispatching to different transfer modes is agreed on

- construct a GP by transferring knowledge from a pre-trained prior GP - basic implementation for full mean transfer - the posterior mean of the pretrained GP are used as mean module for GP - hypereparameters are frozen and mean is evaluated at source points - interface might later be extended to other mean transfers (initialize hyperparameters) or covariance transfer - new `PriorMean` class, that implements mean of prior GP as botorch module

Scienfitz · 2026-01-07T08:29:00Z

baybe/surrogates/gaussian_process/prior_modules.py

+class PriorMean(gpytorch.means.Mean):
+    """GPyTorch mean module using a trained GP as prior mean.
+
+    This mean module wraps a trained Gaussian Process and uses its predictions


process in Gaussian process should not be capitalized (unless its a headline or similar)

baybe/surrogates/gaussian_process/core.py

Scienfitz · 2026-01-07T08:36:24Z

baybe/surrogates/gaussian_process/core.py

+            New GaussianProcessSurrogate instance with transfer learning
+
+        Raises:
+            ValueError: If prior_gp is not fitted


there is also a raise of the variable is not a SingleTaskGP which is not mentioned here?

Scienfitz · 2026-01-07T08:37:56Z

baybe/surrogates/gaussian_process/core.py

@@ -113,11 +115,57 @@ class GaussianProcessSurrogate(Surrogate):
    _model = field(init=False, default=None, eq=False)


given that we already have this _model attribute, can you explain why we need to introduce yet another attribute like _prior_gp? Naively I would suspect the first contains the latter. Or at least we should strive to avoid putting alot of additional attributes in this class (because they will esentially be irrelevant for non TL cases)

I see your point and I thought the same before, but unfortunately I couldn't find a good solution to this: The problem is that there is a gap between creation via from_prior and fitting the model via fit. The instance must somehow remember it should use transfer learning and its prior to be able to create the _model. I'd be happy to change this and will give it another thought. Maybe the logic could be moved to some KernelFactory or MeanFactory. Do you have any suggestions how to get rid of this attribute?

@AVHopp , do you maybe have an idea?

How about introducing a new mean factory similar to the kernel factories in BayBE?

# New default class ConstantMeanFactory(MeanFactory): def __call__(self, ..) return gpytorch.means.ConstantMean() class PriorMeanFactory(MeanFactory): def __init__(self, prior_gp: GPSurrogate): self.prior_gp = deepcopy(prior_gp) def __call__(self, batch_shape: torch.Size) : return PriorMean()

Then in from_prior I'd just replace the mean factory by the new PriorMeanFactory and could remove the attribute from the surrogate class, but this would add an entirely new factory pattern to BayBE.

yeah I would prefer that a bit

although the main reason for the factory was that search space info is needed when creating the kernels, which is not available yet when specifying the attribute here to the surrogate. that doesnt seem to bet he case here with the means, right?

So a factory is not strictly needed but I see till two advantages why I would prefer it:

it wouldnt be an emtpy unused content in no prior gp is used as it would hold the default factory

it would be more consistent to have all kinds of fatories rather than haveing a mixture of factories and other optional model-related attributes

About _model: This is supposed to hold the fitted botorch model right? So would it make any sense to only partially initialize it with the means? If no, then forget hat idea

Is this thread now still relevant, given that we agreed to a Factory approach in our meeting (iirc)?

Scienfitz · 2026-01-07T08:38:52Z

baybe/surrogates/gaussian_process/core.py

    """The actual model."""

+    # Transfer learning fields
+    _prior_gp = field(init=False, default=None, eq=False)


i see its prob the same issue as with _model so ideally you can paste the same comment thats there also here

depending ont he design this attribbute might also be removed and the commen is obsolete

Scienfitz · 2026-01-07T08:40:04Z

baybe/surrogates/gaussian_process/core.py

+        if self._prior_gp is not None and hasattr(self._prior_gp, "input_transform"):
+            # Use prior's transforms for consistency in transfer learning
+            input_transform = self._prior_gp.input_transform
+            outcome_transform = self._prior_gp.outcome_transform


since there is an explicit check for inout_transform, is it always guaranteed to have output_transform?

Why is the heck for input_transform even needed?

Scienfitz · 2026-01-07T08:49:01Z

baybe/surrogates/gaussian_process/prior_modules.py

+from torch import Tensor
+
+
+class PriorMean(gpytorch.means.Mean):


Some question for understanding:

What does this new class achieve what is not possible with the existing botorch constant mean class?

When the incoming mean is a constant mean, this class would also effectively produce a contant mean?

Afaik all our GP have constant mean, so everything woudl forever be costant mean. Is this class here then necessary? Couldnt we just use the botorch cosntant mean class for the new TL case as well, except that the number is fixed and predetermined, ie somehow "set"?

I think there might be some misunderstanding here. The incoming mean is not constant since the prior GP is fitted on some data already and we are using its posterior here. Even if the prior GP originally had a ConstantMean, once trained, its posterior mean will not be constant anymore. Or am I misunderstanding your comment?

I see thanks for clarifying, I see now the need for the class
please lets just make sure this is optimized and does not impose any computaitonal bottleneck

Is it also right that this implementaiton is the variant of completely frozen prio mean? ie the mean is not just a prior but its forever the mean for our actual GP used int he campaign?

Copilot

Pull request overview

This PR introduces a transfer learning capability for Gaussian Process surrogates through a new from_prior constructor method. The implementation enables mean function transfer from a pre-trained GP to a new GP by using the source GP's posterior mean predictions as the mean module for the target GP.

Key changes:

New PriorMean class that wraps a trained GP as a BoTorch-compatible mean module with frozen hyperparameters
New from_prior class method for constructing a GP surrogate with transfer learning capabilities
Modified _fit method to conditionally use the prior GP's mean and transforms when available

Reviewed changes

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

File	Description
baybe/surrogates/gaussian_process/prior_modules.py	Introduces `PriorMean` class to wrap a trained GP as a mean module for transfer learning
baybe/surrogates/gaussian_process/core.py	Adds `from_prior` constructor and updates `_fit` to support mean function transfer

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

Copilot · 2026-01-07T13:39:43Z

baybe/surrogates/gaussian_process/core.py

+            input_transform = Normalize(
+                train_x.shape[-1],
+                bounds=context.parameter_bounds,
+                indices=numerical_idxs,


The indices parameter expects a list but numerical_idxs is a tuple. While this may work in practice, it's inconsistent with the previous implementation that used list(numerical_idxs) on line 217 in the original code. For consistency and to match the expected type, convert the tuple to a list.

Suggested change

indices=numerical_idxs,

indices=list(numerical_idxs),

baybe/surrogates/gaussian_process/prior_modules.py

baybe/surrogates/gaussian_process/core.py

Copilot · 2026-01-07T13:39:44Z

baybe/surrogates/gaussian_process/core.py

+        from copy import deepcopy
+
+        from botorch.models import SingleTaskGP


The import statements are placed inside the method. Since deepcopy is already imported at the module level (line 5 in prior_modules.py) and SingleTaskGP is imported in the TYPE_CHECKING block (line 31), these local imports are redundant and should be removed in favor of the module-level imports.

AVHopp

Just some minor comments as @Scienfitz raised some questions that might impact the design of this code, I did not fully review everything yet.

baybe/surrogates/gaussian_process/core.py

AVHopp · 2026-01-07T13:41:36Z

baybe/surrogates/gaussian_process/core.py

+        kernel_factory: KernelFactory | None = None,
+        **kwargs,
+    ) -> GaussianProcessSurrogate:
+        """Create a GP surrogate with mean function transfer learning.


I think this docstring needs a bit more explanation on what exactly is done and transferred. Also, the description in the Returns: part could contain more information (but might not be needed if you add 2-3 sentences here describing what this does in more detail)

baybe/surrogates/gaussian_process/core.py

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

Scienfitz · 2026-01-09T14:50:07Z

baybe/surrogates/gaussian_process/prior_modules.py

the name x_modules is a bit inconsistent compared to our other naming
just means.py?

Scienfitz · 2026-01-09T15:03:05Z

baybe/surrogates/gaussian_process/prior_modules.py

+        Returns:
+            Mean predictions from the wrapped GP.
+        """
+        self.gp.eval()


wouldnt it make sense to move these eval statements into init because they are only needed once?

AVHopp · 2026-01-13T14:58:48Z

baybe/surrogates/gaussian_process/prior_modules.py

+    Args:
+        gp: Trained Gaussian Process to use as mean function.
+        batch_shape: Batch shape for the mean module.
+        **kwargs: Additional keyword arguments.


Is it necessary to include those in this class/the __init__? Currently they seem to be silently ignored, so I would propose to either remove them completely if possible or at least mention that they are being ignored.

AVHopp · 2026-01-13T16:05:19Z

baybe/surrogates/gaussian_process/core.py

@@ -113,11 +115,57 @@ class GaussianProcessSurrogate(Surrogate):
    _model = field(init=False, default=None, eq=False)


Is this thread now still relevant, given that we agreed to a Factory approach in our meeting (iirc)?

kalama-ai · 2026-02-02T09:20:01Z

Note: On hold until mean factory is implemented.

kalama-ai added 3 commits December 18, 2025 15:04

Make constructor method public

e88f94b

Set source GP to eval mode for predictions

cdf71a9

kalama-ai requested review from AVHopp, AdrianSosic and Scienfitz as code owners December 18, 2025 14:37

Scienfitz requested changes Jan 7, 2026

View reviewed changes

AVHopp requested a review from Copilot January 7, 2026 13:38

Copilot AI reviewed Jan 7, 2026

View reviewed changes

AVHopp reviewed Jan 7, 2026

View reviewed changes

kalama-ai and others added 5 commits January 8, 2026 12:05

Fix typo in docstring

1dff685

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

Fix type annotation

7d17a12

Create new GP from GaussianProcessSurrogate instead of botorch instance

eacdb28

Add all raises to docstring

4cc837d

Update docstring

63e385c

Scienfitz reviewed Jan 9, 2026

View reviewed changes

AVHopp reviewed Jan 13, 2026

View reviewed changes

Scienfitz assigned kalama-ai Jan 14, 2026

		@@ -113,11 +115,57 @@ class GaussianProcessSurrogate(Surrogate):
		_model = field(init=False, default=None, eq=False)

		from torch import Tensor


		class PriorMean(gpytorch.means.Mean):

		from copy import deepcopy

		from botorch.models import SingleTaskGP

Conversation

kalama-ai commented Dec 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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!

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!

kalama-ai Jan 8, 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 Jan 9, 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!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Copilot AI Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

AVHopp left a comment

Choose a reason for hiding this comment

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!

Choose a reason for hiding this comment

Uh oh!

kalama-ai commented Feb 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

kalama-ai commented Dec 18, 2025 •

edited

Loading

kalama-ai Jan 8, 2026 •

edited

Loading

Scienfitz Jan 9, 2026 •

edited

Loading