Multitask finetuning and acceleration of chemical pretrained models for small molecule drug property prediction

TL;DR

This work demonstrates that multitask finetuning of chemical foundation graph models can significantly improve small-molecule property predictions, with KERMT showing substantial gains on large, correlated ADMET datasets while KPGT excels in smaller data regimes. The authors compare multiple pretrained encoders (KERMT, KGPT, MoLFormer) against a Chemprop baseline across internal and public datasets, highlighting data-size and task-correlation effects on performance. They reveal that multitask finetuning yields the strongest benefits when downstream data are abundant, and they publish a new public multitask ADMET benchmark to standardize future evaluations. To enable industrial-scale use, the paper also reports accelerated training/inference via distributed pretraining and the cuik-molmaker library, along with public code and datasets to facilitate adoption and benchmarking.

Abstract

Chemical pretrained models, sometimes referred to as foundation models, are receiving considerable interest for drug discovery applications. The general chemical knowledge extracted from self-supervised training has the potential to improve predictions for critical drug discovery endpoints, including on-target potency and ADMET properties. Multi-task learning has previously been successfully leveraged to improve predictive models. Here, we show that enabling multitasking in finetuning of chemical pretrained graph neural network models such as Kinetic GROVER Multi-Task (KERMT), an enhanced version of the GROVER model, and Knowledge-guided Pre-training of Graph Transformer (KGPT) significantly improves performance over non-pretrained graph neural network models. Surprisingly, we find that the performance improvement from finetuning KERMT in a multitask manner is most significant at larger data sizes. Additionally, we publish two multitask ADMET data splits to enable more accurate benchmarking of multitask deep learning methods for drug property prediction. Finally, we provide an accelerated implementation of the KERMT model on GitHub, unlocking large-scale pretraining, finetuning, and inference in industrial drug discovery workflows.

Figures (7)

• Figure 1: Schematic overview of the studies.a Pretraining and finetuning scheme. The parameters of the pretrained chemical encoder are updated during finetuning with labeled data. b Pretraining algorithms investigated. The algorithms are categorized into language-based and graph-based models. c Single-task and multitask scheme. C1 - C7 represent example chemical compounds, and T1 - T4 represent example prediction tasks. For multitask model, the dataset is sparse. The datasets are divided into training and test sets using either temporal or cluster-based splits. d Summary of the investigated studies.
• Figure 2: Comparison of model performance on ADMET assay predictions.a Performance improvement of multitask (MT) models compared to single-task (ST) models. b - d Performance improvement of pretrained MT models compared to Chemprop MT model. The plots show the results on all available assays for the internal, public, and Biogen datasets, respectively. Brown solid lines represent the number of training data for each assay. A colored assay name (blue: Chemprop, yellow: KPGT, red: KERMT) represents the model which performs statistically better as indicated by the Pearson $r^2$ being higher than the upper error limit of other models. Black labeled assays indicate that two or three models performed statistically the same. $N_{\text{best}}$ represents the number of assays the model statistically outperforms. e Number of assays the model statistically outperforms ($N_{\text{best}}$) or underperforms ($N_{\text{worst}}$) for different ranges of assay training size.
• Figure 3: Model performance for on-target potency prediction.a, d Performance improvement for pretrained models compared to Chemprop. Brown line indicates number of training data for a target. a Results on single task potency models. EGFR$_{kin}$ is from the kinase multitask dataset and EGFR$_{bdb}$ is from BindingDB. d Results on multitask kinase potency model. b - c Parity plots (top) and categorical bar plots (bottom) for each model. $R^2$, MAE, and RMSE represent coefficient of determination, mean absolute error, and root mean squared error in $\log_{10}(\mathrm{nM})$ units. The categorical bar plot is annotated with the percentage of the predicted bin (x-axis) that belongs to each experimental bin (colored legend). b$EC_{50}$ predictions for Target 1. c$IC_{50}$ predictions for EGFR.
• Figure 4: Scaling of Chemprop multitask and KERMT models with training set size and number of tasks. All results are evaluated on the internal ADMET dataset with the 80-20 temporal split. a Average performance improvement by the KERMT model relative to the Chemprop MT model for different training set sizes. The average performance improvement is calculated using top 10 assays with the most performance difference between KERMT and Chemprop models. b Number of assays each model statistically outperforms for different training set sizes. c Pearson $r^2$ of each model on selected assays across a different number of tasks in a multitask model. A brown dash-dotted line with cross markers (Max. Corr. Coeff.) represents the maximum value of the absolute Spearman correlation coefficients between the selected assay and all other assays in each multitask dataset.
• Figure 5: Chemprop and KERMT multitask model performance across a different range of the Tanimoto similarity scores of the test set compounds and peptides. All results are evaluated on the internal ADMET dataset with the 80-20 temporal split. a Pearson $r^2$ of each model for selected ADMET assays. b Number of assays each model statistically outperforms for each Tanimoto similarity range. c Violin plot displaying the distributions of the performance improvement by the KERMT model relative to the Chemprop MT model. The central horizontal line in each violin plot represents the median of the distribution. In b and c, the results are evaluated using all 30 assays from the internal dataset. d Histogram of the Tanimoto similarity scores of the test set compounds. e Model performance on the test set peptides. The percentage in the assay label represents the percentage of peptides in the training data for each assay. A brown solid line represents the number of peptide training data for each assay. A colored assay name (blue: Chemprop, red: KERMT) represents the model which performs statistically better. Pearson $\overline{r^2}$ represents the average performance over all assays.