BOOM: Benchmarking Out-Of-distribution Molecular Property Predictions of Machine Learning Models

TL;DR

BOOM delivers a comprehensive, open benchmark for assessing out-of-distribution extrapolation in molecular property prediction across diverse models, datasets, and tasks. It reveals that no current approach achieves universal OOD generalization, with pretrained foundation models boosting in-distribution accuracy but offering limited OOD gains, while 3D/inductive-bias architectures and data augmentation provide the most consistent OOD benefits on specific tasks. The work highlights critical factors—representation, pretraining strategy, hyperparameter choices, and data augmentation—that influence OOD performance and identifies data- and architecture-specific pathways for advancing chemical foundation models toward robust extrapolation. By providing a standardized, extensible benchmark and open-source code, BOOM aims to steer the community toward tackling the frontier of strong OOD generalization in chemical ML.

Abstract

Advances in deep learning and generative modeling have driven interest in data-driven molecule discovery pipelines, whereby machine learning (ML) models are used to filter and design novel molecules without requiring prohibitively expensive first-principles simulations. Although the discovery of novel molecules that extend the boundaries of known chemistry requires accurate out-of-distribution (OOD) predictions, ML models often struggle to generalize OOD. Furthermore, there are currently no systematic benchmarks for molecular OOD prediction tasks. We present BOOM, $\boldsymbol{b}$enchmarks for $\boldsymbol{o}$ut-$\boldsymbol{o}$f-distribution $\boldsymbol{m}$olecular property predictions -- a benchmark study of property-based out-of-distribution models for common molecular property prediction models. We evaluate more than 140 combinations of models and property prediction tasks to benchmark deep learning models on their OOD performance. Overall, we do not find any existing models that achieve strong OOD generalization across all tasks: even the top performing model exhibited an average OOD error 3x larger than in-distribution. We find that deep learning models with high inductive bias can perform well on OOD tasks with simple, specific properties. Although chemical foundation models with transfer and in-context learning offer a promising solution for limited training data scenarios, we find that current foundation models do not show strong OOD extrapolation capabilities. We perform extensive ablation experiments to highlight how OOD performance is impacted by data generation, pre-training, hyperparameter optimization, model architecture, and molecular representation. We propose that developing ML models with strong OOD generalization is a new frontier challenge in chemical ML model development. This open-source benchmark will be made available on Github.

Figures (24)

• Figure 1: Two example OOD datasets included in the BOOM benchmarking. To assess OOD performance, we split each chemical property dataset into an out-of-distribution (OOD) Test Set (blue), a in-distribution (ID) Test Set (orange) and a Train Set (green), as described in Section 2.2. (Left) 10k Dataset density OOD split and (Right) QM9 HOMO-LUMO gap OOD split.
• Figure 2: Normalized ID and OOD RMSE values for all models of interest for each property. RMSE values are normalized such that the best performing model has a value of 0 and the worst performing model has a value of 1. We omit the Regression Transformer model to remove the extreme values for better visualization. The models are sorted with respect to average RMSE, with the best model (lowest overall RMSE) appearing at the top of the heatmap and the worst model at the bottom.
• Figure 3: Representative parity plots illustrating the OOD (blue) and ID (orange) test predictions on the 10K solid heat of formation dataset. (Left) MoLFormer predictions on this task exhibit poor OOD performance with weak correlation on the OOD samples. (Right) MoLFormer displays strong OOD performance on the QM9 zero-point vibrational energy task.
• Figure 4: We present binned $R^2$ scores for OOD and standard $R^2$ scores for ID on each task for all models. The orange and blue bars indicate the performance averaged across all models for ID and OOD, respectively. Nearly all models have significant discrepancies between ID and OOD performance, but some models can reach ID-level accuracy. We observe that OOD performance is highly task-dependent.
• Figure 5: OOD Performance of chemical foundation models (ChemBERTA, MoLFormer and Regression Transformer) with and without pretraining, averaged across all tasks. We find that current pretraining strategies improve ID performance, but not OOD. The task-specific performances are provided in the Appendix (Figure \ref{['fig:pretraining_FMs_all']}).