baybe icon indicating copy to clipboard operation
baybe copied to clipboard

Add initial version of the new SourcePriorGaussianProcessSurrogate

Open kalama-ai opened this issue 2 months ago • 1 comments

SourcePriorGaussianProcessSurrogate Implementation

The SourcePriorGaussianProcessSurrogate in baybe/surrogates/transfer_learning/source_prior.py implements transfer learning through a hierarchical approach:

  1. Training a source GP on source task data (without task dimension)
  2. Using the mean or covariance of the source GP as a prior for the target GP
  3. Training the target GP on target data

Training Flow

  1. The surrogate receives a search space and validates the transfer learning context via _validate_transfer_learning_context(), ensuring a TaskParameter is present.
  2. _identify_target_task() identifies the task parameter from the search space and determines the target task by examining the active_values of the TaskParameter.
  3. Data Extraction: The surrogate splits the training data into source and target pairs using _extract_task_data(): - Source data: (X_s, y_s) - Target data: (X_t, y_t) These pairs contain only features (no task parameter column) after filtering.
  4. Model Training (_fit()) a) A SingleTaskGP instance is created using BayBE's default configuration (mean and covariance modules). This self._source_gp receives a reduced search space (without task parameter) and is fitted to the source data (X_s, y_s). b) A new SingleTaskGP instance is created that uses the posterior mean of the source GP as a mean prior (via PriorMean class). This self._target_gp is fitted to the target data (X_t, y_t).

Making Predictions (_posterior())

  1. The method receives candidates in computational representation, extracts batch dimensions, and initializes empty tensors for mean and covariance predictions.
  2. Candidates are filtered into source and target groups using the task parameter column.
  3. Separate Predictions: - Source predictions: Made using the self._source_gp - Target predictions: Made using the self._target_gp (which has source prior)) Individual predictions are combined into unified mean and covariance tensors for the entire batch.

New GaussianProcessSurrogate.from_prior_gp class method to create a GP surrogate by transferring the mean or covariance from a pretrained prior:

  • transfer_mode "mean" or "kernel" defines wether mean or kernel from prior should be used
  • new GaussianProcessSurrogate._initialize_model method that will create the _model depending on whether a prior is given or not
  • PriorMean and PriorKernel classes in baybe/surrogates/gaussian_process/prior_modules.py extract posterior mean or covariance from a prior GP

Helper classes

  1. SourcePriorWrapperModel: Provides BoTorch compatibility by wrapping the SourcePriorGaussianProcessSurrogate for the _to_botorch() method.

Current Limitations

  1. Single Active Value: The model expects exactly one active value for the task parameter in the search space.
  2. Single Source Limitation: The surrogate is limited to one source task. Extending to multiple sources would require developing methods to combine multiple source GPs into unified prior mean (and potentially covariance) functions.
  3. Single Output: Multi-output extensions haven't been considered yet.

kalama-ai avatar Oct 24 '25 09:10 kalama-ai

First of all: Appreciate the detailed PR description :)

AVHopp avatar Oct 29 '25 09:10 AVHopp