MUBen: Benchmarking the Uncertainty of Molecular Representation Models

TL;DR

MUBen introduces a comprehensive benchmark to evaluate uncertainty quantification methods on state-of-the-art molecular representation backbones across a wide set of descriptors. By combining four primary backbones with diverse UQ strategies and MoleculeNet datasets under scaffold-split OOD conditions, the study reveals that Deep Ensembles most consistently improve both property predictions and calibration, albeit at a high computational cost, while post-hoc calibration like Temperature Scaling offers practical gains in many scenarios. The results also show that backbone choice matters: Uni-Mol often yields the strongest predictive performance but can exhibit overconfidence, whereas other backbones offer different trade-offs depending on task and data distribution. The work highlights the importance of selecting backbone-UQ pairs tailored to the application, and documents clear limitations such as a curated method set and coarse hyperparameter grids, pointing to future expansion of the benchmark to broaden coverage and realism in uncertainty-critical molecular tasks.

Abstract

Large molecular representation models pre-trained on massive unlabeled data have shown great success in predicting molecular properties. However, these models may tend to overfit the fine-tuning data, resulting in over-confident predictions on test data that fall outside of the training distribution. To address this issue, uncertainty quantification (UQ) methods can be used to improve the models' calibration of predictions. Although many UQ approaches exist, not all of them lead to improved performance. While some studies have included UQ to improve molecular pre-trained models, the process of selecting suitable backbone and UQ methods for reliable molecular uncertainty estimation remains underexplored. To address this gap, we present MUBen, which evaluates different UQ methods for state-of-the-art backbone molecular representation models to investigate their capabilities. By fine-tuning various backbones using different molecular descriptors as inputs with UQ methods from different categories, we assess the influence of architectural decisions and training strategies. Our study offers insights for selecting UQ for backbone models, which can facilitate research on uncertainty-critical applications in fields such as materials science and drug discovery.

Figures (10)

• Figure 1: MUBen's pipeline with datasets, backbones, UQ methods and metrics enumerated.
• Figure 2: MRR of the backbone models; each metric is macro-averaged from the reciprocal ranks of all corresponding UQ methods across all datasets. Uni-Mol stands out by consistently surpassing its counterparts in property prediction accuracy for both classification and regression tasks. Conversely, GROVER exhibits a more balanced performance in both prediction accuracy and UQ across various tasks.
• Figure 3: The calibration plot of UQ methods. Here, "true probability" refers to the empirical prediction accuracy, and "predicted probability" refers to the average predicted probability within each bin (\ref{['sec:prob.setup']}). These results reveal that Temperature Scaling typically enhances calibration effectiveness in the majority of scenarios, whereas Focal Loss tends to lead to over-calibration of the model, producing under-confident results.
• Figure 4: The absolute error between the model-predicted mean and true labels against the predicted standard deviation. We compare the performance of SGLD with the deterministic prediction on different backbones and datasets. The "$y=kx$" lines indicate whether the true labels lie within the $k$-std range of the predicted Gaussian. Also, a model is perfectly calibrated when its output points are arranged on an "$y=kx$" line for an arbitrary $k$. Notably, SGLD is observed to generate a larger variance for OOD samples, which tends to correspond more closely with the prediction errors on average.
• Figure 5: MRRs of TorchMDNet and Uni-Mol on datasets grouped by dataset property categories. MRR calculations are confined to results from these two backbones. Only relative values matter. TorchMDNet is comparible to Uni-Mol on Quantum Mechanics properties, where it is pre-trained on, but is outperformed by Uni-Mol on all other property categories.