Improving Predictions of Molecular Properties with Graph Featurisation and Heterogeneous Ensemble Models

AI summary

Overview

• Research area: Machine learning for molecular property prediction (cheminformatics / drug discovery), specifically graph neural networks and heterogeneous ensemble modelling.
• Technical level: Intermediate. The paper assumes familiarity with standard ML concepts (descriptors, train/validation/test splits, ensembles, inductive bias) and with molecular ML tooling, but its arguments are presented in mostly plain language.
• One-sentence scope: The paper introduces a "MetaModel" that aggregates predictions from a diverse set of classical ML algorithms trained on RDKit descriptors combined with latent features extracted from a ChemProp message-passing neural network, and benchmarks it against ChemProp itself on the MoleculeNet datasets.

What This Paper Is About

Graph neural networks such as ChemProp can learn task-specific structural information from molecules, but they hand their latent representation to a simple feed-forward network head that tends to be weaker than other model classes on tabular-style inputs. The authors ask whether combining ChemProp's learned features with conventional RDKit descriptors, and then feeding them to a mixed ensemble of many different ML algorithms, can beat ChemProp on its own. They test this on the MoleculeNet benchmarks, then run ablations without external descriptors and with several hyperparameter optimisation strategies.

Key Contributions

1. The MetaModel framework — a heterogeneous ensemble of classical ML models (regression: Lasso, ridge, KNN, kernel-ridge regression, random forests, XGBoost, Gaussian processes, RBF interpolation, MLP and ResNet; classification: KNN, LDA/QDA, logistic regression, naïve Bayes, RF, XGBoost, GP, MLP and ResNet). Each sub-model receives its own randomised train/validation split; sub-models are weighted by validation performance (inverse MSE for regression; ROC-AUC or PRC for classification, with PRC used when the minority class is under 10% of data points); the ensemble is pruned to the ten best sub-models, then features are pruned to those above 2% of the importance of the most important feature, and the smaller ensemble is retrained.
2. A featurisation scheme that combines general-purpose RDKit CalcMolDescriptors features with task-specific learned "MPNN descriptors" — the latent output of a ChemProp (v2.1.2) message-passing neural network with its feed-forward head removed, after training the full model on the target.
3. A benchmark comparison against ChemProp on MoleculeNet datasets (excluding QM7b, which lacks SMILES strings, and PDBBind, which is outside scope), using MoleculeNet's recommended split strategy and metrics.
4. Two follow-on studies: an ablation where no external descriptors are provided to either model, and a hyperparameter optimisation study comparing optimising ChemProp in isolation versus optimising ChemProp for the downstream MetaModel's performance.

Main Findings

• The MetaModel beats ChemProp on all six regression datasets tested. QM7 MAE 0.181 (MetaModel with RDKit + ChemProp features) versus 0.264 (ChemProp); QM8 0.213 versus 0.235; QM9 0.019 versus 0.040; ESOL RMSE 0.302 versus 0.334; FreeSolv 0.252 versus 0.280; Lipophilicity 0.467 versus 0.479. Five of these differences were highly significant by bootstrapping and one was marginal (p = 0.05).
• The MetaModel beats ChemProp on 6 of the 9 classification datasets, but only 3 of those wins are significant. ChemProp significantly outperforms the MetaModel on ToxCast (ROC 0.763 versus 0.725 for the MetaModel with ChemProp features) and Tox21 (0.849 versus 0.824). Both are described as sparse datasets with many gaps and many target variables (617 and 12 respectively).
• Adding ChemProp's latent features helps the MetaModel on most datasets, but not all. The MetaModel with RDKit + ChemProp features outperformed the RDKit-only MetaModel significantly on 4 of the 6 regression datasets and with low significance on 2 more (p = 0.21 and p = 0.18). On ESOL and FreeSolv it did not improve the fit, and on ESOL it was significantly worse (p = 0.01). For classification, the ChemProp-feature version won on 6 of 9 datasets, with only 3 of those significant at the 90% level (P = 0.00, 0.02 and 0.09 for PCBA, ClinTox and Tox21).
• Removing external descriptors flips the advantage to ChemProp on most datasets. With MPNN features only, ChemProp outperformed the MetaModel on the majority of datasets, indicating the MetaModel relies more heavily on the RDKit descriptors. The MetaModel's performance dropped relative to ChemProp by a large margin on MUV (PRC 0.18, 3.29× its baseline, p = 1.5×10⁻⁸) and notably on QM7, where both models degraded badly (MetaModel MAE 0.42, 1.60× its baseline; ChemProp 0.44, 1.66× its baseline).
• QM7 depends on descriptors the MPNN cannot derive itself. The large error increase for both models when RDKit features are withheld implies some RDKit descriptors are crucial for that dataset and inaccessible to ChemProp.

Authors’ abstract

We explore a "best-of-both" approach to modelling molecular properties by combining learned molecular descriptors from a graph neural network (GNN) with general-purpose descriptors and a mixed ensemble of machine learning (ML) models. We introduce a MetaModel framework to aggregate predictions from a diverse set of leading ML models. We present a featurisation scheme for combining task-specific GNN-derived features with conventional molecular descriptors. We demonstrate that our framework outperforms the cutting-edge ChemProp model on all regression datasets tested and 6 of 9 classification datasets. We further show that including the GNN features derived from ChemProp boosts the ensemble model's performance on several datasets where it otherwise would have underperformed. We conclude that to achieve optimal performance across a wide set of problems, it is vital to combine general-purpose descriptors with task-specific learned features and use a diverse set of ML models to make the predictions.