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

TL;DR

This paper investigates how atom-level pretraining with quantum mechanics data enhances the robustness and generalization of graph neural network models for molecular property prediction, addressing distributional shift challenges in QSAR tasks.

Contribution

It introduces atom-level pretraining with QM data for GNNs, demonstrating improved performance and more Gaussian-like feature distributions, a novel analysis of hidden state molecular representations.

Findings

Pretraining on atom-level QM data improves model performance.

Atom-level pretraining leads to more Gaussian-like feature distributions.

The approach enhances robustness to distribution shifts in molecular data.

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.