In-Context Learning Meets Small Molecule Property Prediction: Benchmarking Novel Machine Learning Approaches

Abstract

Recently, a new category of machine learning approaches for tabular data has emerged: tabular foundation models (TFM), based on in-context learning. A TFM is a neural network (usually a transformer) pretrained primarily on synthetic data. Its input is an entire data set: features and labels for training records, along with features for the target records. This study comprehensively benchmarked these methods for predicting the properties of small organic molecules. Several TFM were compared with multiple machine learning methods across 11 data sets (regression, random and structure-aware splits, up to 10,000 molecules each). The evaluation showed that TFM consistently outperform XGBoost, CatBoost, multilayer perceptrons, and other descriptor-based methods, even with careful hyperparameter selection for the baselines. TFM also demonstrate accuracy on par with or better than graph-based methods, including those pretrained on chemical data. Uni-Mol2, a pretrained deep neural network operating on 3D atomic coordinates, slightly outperforms TFM in some experiments. However, this comparison deliberately disfavors TFM, as they do not use any chemical pretraining and rely on a minimalistic set of 2D molecular descriptors without feature engineering. Some further improvement in results for relatively large data sets and random splits is achieved using retrieval: for each test molecule, the 500 closest neighbors (by Tanimoto similarity) are selected from the training set, and TFM inference is performed on this local subset. Overall, TFM (particularly the TabPFN family) are highly promising for predicting the properties of small molecules.