In-Context Learning for Latent Space Bayesian Optimization

2026-06-08

AI Summary

Existing tabular foundation models exhibit limited performance in latent-space Bayesian optimization due to a mismatch between their pretraining objectives and the distribution of downstream optimization tasks. To address this, this work proposes a task-specific continual pretraining strategy tailored for latent-space Bayesian optimization. Specifically, synthetic optimization tasks are constructed in the latent space of a molecular variational autoencoder (VAE), and a tabular foundation model (e.g., TabICL) is adaptively fine-tuned on these tasks. An anchoring regularizer is further introduced to preserve the model’s general-purpose regression capabilities. Evaluated on molecular optimization benchmarks, the proposed approach significantly improves optimization performance, demonstrating both the efficacy and necessity of task-customized pretraining for latent-space Bayesian optimization.

Abstract

Bayesian optimization (BO) is a central tool for sample-efficient design, and latent-space Bayesian optimization (LSBO) extends it to structured objects such as molecules and proteins. In parallel, tabular foundation models such as TabPFN and TabICL now achieve state-of-the-art regression performance and are increasingly used as BO surrogates. Because their Bayesian behavior is induced by large synthetic pretraining collections, the composition of this pretraining distribution is crucial. LSBO creates a distinctive mismatch: the induced map from latent code to objective value differs markedly from the regression tasks used to train current in-context models. We address this mismatch by complementing the pretraining stage of tabular foundation model surrogates with synthetic optimization tasks defined on the latent space of a molecular VAE. The continued-pretraining objective features a regularizer that anchors the model to the original checkpoint, preserving its broad regression prior while avoiding overspecialization to the adaptation tasks. On held-out molecular optimization benchmarks, the resulting model achieves strong performance, supporting the relevance of LSBO-specific adaptation for in-context surrogates.

Problem

• In-Context Learning
• Latent Space Bayesian Optimization
• Tabular Foundation Models
• Pretraining Mismatch
• Molecular Optimization

Innovation

• In-Context Learning
• Latent Space Bayesian Optimization
• Tabular Foundation Models
• Continual Pretraining
• Molecular Optimization