In-Context Learning for Latent Space Bayesian Optimization
2026-06-08 • Machine Learning
Machine Learning
AI summaryⓘ
The authors study how to improve a type of AI called Bayesian optimization (BO), which helps design things like molecules using fewer tests. They focus on combining BO with special AI models trained on tables of data, called tabular foundation models, to predict outcomes. These models were initially trained on many different tasks, but their knowledge doesn’t fit well when used on the unique data from molecule design. To fix this, the authors retrain the models with tasks that better match the molecule design process, while keeping the models’ original general knowledge. This approach improved performance on tests involving molecule optimization.
Bayesian optimizationlatent-space Bayesian optimizationtabular foundation modelsTabPFNTabICLmolecular VAEpretrainingsurrogate modelsin-context learningmolecular optimization
Authors
Tuan A. Vu, Harri Lähdesmäki, Julien Martinelli
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.