baybe
baybe copied to clipboard
Add initial version of the new SourcePriorGaussianProcessSurrogate
SourcePriorGaussianProcessSurrogate Implementation
The SourcePriorGaussianProcessSurrogate in baybe/surrogates/transfer_learning/source_prior.py implements transfer learning through a hierarchical approach:
- Training a source GP on source task data (without task dimension)
- Using the mean or covariance of the source GP as a prior for the target GP
- Training the target GP on target data
Training Flow
- The surrogate receives a search space and validates the transfer learning context via
_validate_transfer_learning_context(), ensuring aTaskParameteris present. -
_identify_target_task()identifies the task parameter from the search space and determines the target task by examining theactive_valuesof theTaskParameter. - 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. - Model Training (
_fit()) a) ASingleTaskGPinstance is created using BayBE's default configuration (mean and covariance modules). Thisself._source_gpreceives a reduced search space (without task parameter) and is fitted to the source data(X_s, y_s). b) A newSingleTaskGPinstance is created that uses the posterior mean of the source GP as a mean prior (viaPriorMeanclass). Thisself._target_gpis fitted to the target data(X_t, y_t).
Making Predictions (_posterior())
- The method receives candidates in computational representation, extracts batch dimensions, and initializes empty tensors for mean and covariance predictions.
- Candidates are filtered into source and target groups using the task parameter column.
- Separate Predictions:
- Source predictions: Made using the
self._source_gp- Target predictions: Made using theself._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_modelmethod that will create the_modeldepending on whether a prior is given or not -
PriorMeanandPriorKernelclasses inbaybe/surrogates/gaussian_process/prior_modules.pyextract posterior mean or covariance from a prior GP
Helper classes
-
SourcePriorWrapperModel: Provides BoTorch compatibility by wrapping theSourcePriorGaussianProcessSurrogatefor the_to_botorch()method.
Current Limitations
- Single Active Value: The model expects exactly one active value for the task parameter in the search space.
- 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.
- Single Output: Multi-output extensions haven't been considered yet.
First of all: Appreciate the detailed PR description :)