Analysis of Atom-level pretraining with Quantum Mechanics (QM) data for Graph Neural Networks Molecular property models

TL;DR

QSAR models struggle to generalize to novel compounds due to distribution drift and activity cliffs. The authors pretrain Graphormer-based GNNs on atom-level quantum-mechanics properties, using four QM properties and multitask virtual nodes, and analyze hidden representations to connect node-level pretraining with generalization. They find that atom-level QM pretraining improves downstream ADMET tasks on the TDC dataset, yields more Gaussian-like first-layer activations, and reduces distribution shifts between training and test data. This work provides a mechanistic link between atom-level domain knowledge and robust molecular representations, informing pretraining strategy design for QSAR models.

Abstract

Despite the rapid and significant advancements in deep learning for Quantitative Structure-Activity Relationship (QSAR) models, the challenge of learning robust molecular representations that effectively generalize in real-world scenarios to novel compounds remains an elusive and unresolved task. This study examines how atom-level pretraining with quantum mechanics (QM) data can mitigate violations of assumptions regarding the distributional similarity between training and test data and therefore improve performance and generalization in downstream tasks. In the public dataset Therapeutics Data Commons (TDC), we show how pretraining on atom-level QM improves performance overall and makes the activation of the features distributes more Gaussian-like which results in a representation that is more robust to distribution shifts. To the best of our knowledge, this is the first time that hidden state molecular representations are analyzed to compare the effects of molecule-level and atom-level pretraining on QM data.

Figures (28)

• Figure 1: Different aromatic ring representations folded into a single form
• Figure 2: Distribution of first 20 features from the first layer of the Graphormer network for three different training approaches —scratch, HOMO-LUMO pretrained and atom-level pretrained— across test split of lipophilicity dataset .
• Figure 3: This heatmap illustrates the differences in distribution shifts between scratch and atom-levle pretrained networks, calculated across various feature dimensions using metrics such as Kullback-Leibler Divergence, Jensen-Shannon Divergence, Earth Mover's Distance, Total Variation Distance, and Hellinger Distance. Green hues indicate instances where the atom-levle pretrained network exhibits smaller distribution shifts compared to the scratch network, while purple hues denote the opposite. Notably, both heatmaps comparison frequently show reduced distribution shifts in atom-level pretrained networks, specially for Train-Test comparison, which likely helps to explain atom-level pretrained network's superior performance on the test data split.
• Figure 4: Distribution plots Scratch for Lipophilicity test set
• Figure 5: Distribution plots HOMO-LUMO pretrained for Lipophilicity test set