[Paper Review] A benchmark study on reliable molecular supervised learning via Bayesian learning

Doyeong Hwang, Grace Lee+3 more|arXiv (Cornell University)|Jun 12, 2020

Computational Drug Discovery Methods28 references4 citations

TL;DR

This paper benchmarks Bayesian learning methods for improving prediction reliability in molecular graph neural networks (GNNs), demonstrating that techniques like Stochastic Weight Averaging (SWA) and SWAG yield well-calibrated, uncertainty-aware predictions across multiple GNN architectures and molecular property tasks. The study shows that Bayesian methods reduce overconfidence and enhance virtual screening success by providing trustworthy predictive probabilities.

ABSTRACT

Virtual screening aims to find desirable compounds from chemical library by using computational methods. For this purpose with machine learning, model outputs that can be interpreted as predictive probability will be beneficial, in that a high prediction score corresponds to high probability of correctness. In this work, we present a study on the prediction performance and reliability of graph neural networks trained with the recently proposed Bayesian learning algorithms. Our work shows that Bayesian learning algorithms allow well-calibrated predictions for various GNN architectures and classification tasks. Also, we show the implications of reliable predictions on virtual screening, where Bayesian learning may lead to higher success in finding hit compounds.

Motivation & Objective

• To evaluate the reliability of molecular property predictions using Bayesian deep learning in graph neural networks.
• To address the problem of overconfident predictions in neural networks, especially on out-of-distribution (OOD) compounds in virtual screening.
• To assess whether Bayesian inference methods improve calibration and uncertainty estimation in GNN-based molecular classification.
• To provide a practical benchmark for reliable molecular machine learning, particularly for urgent applications like drug discovery.
• To release code for reproducibility and broader adoption in computational chemistry and virtual screening.

Proposed method

• The study employs Expected Calibration Error (ECE) to quantitatively evaluate prediction reliability, measuring the discrepancy between predicted confidence and actual correctness.
• Bayesian learning is applied via approximate inference methods: MAP, Deep Ensemble, MC-Dropout, SGLD, SWA, and SWAG, to estimate posterior distributions over model weights.
• Graph Neural Networks (GIN, GAT, GatedGCN, etc.) are used as the backbone models, with input features derived from molecular graphs and node/edge representations.
• For Bayesian marginalization, predictive distributions are computed by averaging predictions over multiple sampled weights, enabling uncertainty estimation via variance of predictive probabilities.
• The models are trained using scaffold-split data splits (80:10:10) on four MoleculeNet datasets: BACE, BBBP, HIV, and Tox21, with hyperparameters tuned per method.
• SWA and SWAG are applied with cyclic learning rate scheduling and weight averaging over 150–250 epochs, with posterior covariance scaled by 1.0 for SWAG.

Experimental results

Research questions

• RQ1 Can Bayesian learning methods improve the calibration of GNN-based molecular property predictors across diverse architectures and tasks?
• RQ2 How do different Bayesian inference techniques compare in terms of predictive reliability and performance on molecular classification benchmarks?
• RQ3 To what extent do well-calibrated predictions enhance virtual screening success by reducing false positives and false negatives?
• RQ4 What is the impact of uncertainty estimation on prediction behavior for true positives, false positives, true negatives, and false negatives?
• RQ5 How do SWA and SWAG compare to other Bayesian methods in achieving both high accuracy and low ECE in molecular property prediction?

Key findings

• SWA and its variant SWAG consistently achieved the lowest Expected Calibration Error (ECE) across all GNN architectures and molecular property prediction tasks, indicating superior calibration.
• The use of Bayesian learning methods significantly reduced overconfidence, with predictive probabilities closely matching actual correctness rates, especially on OOD and out-of-distribution samples.
• Histogram analysis showed that true positives were predominantly assigned high predictive probabilities, while false positives were often predicted with low confidence—indicating reliable uncertainty quantification.
• Models using SWA and SWAG achieved high AUROC (>0.85) and F1-scores (>0.75) on BACE and BBBP, while maintaining ECE values below 0.05, demonstrating both high performance and reliability.
• The study found that MC-Dropout and Deep Ensemble also improved calibration but were less consistent than SWA and SWAG, particularly on the Tox21 dataset with class-imbalanced tasks.
• The released codebase at https://github.com/AITRICS/mol_reliable_gnn enables reproducible benchmarking and supports urgent applications in drug discovery, such as for COVID-19.