[Paper Review] Advanced Graph and Sequence Neural Networks for Molecular Property Prediction and Drug Discovery

Zhengyang Wang, Meng Liu+9 more|arXiv (Cornell University)|Dec 2, 2020

Computational Drug Discovery Methods4 citations

TL;DR

This paper introduces AdvProp, a comprehensive machine learning framework for molecular property prediction and drug discovery that combines advanced graph and sequence neural networks. It proposes a multi-level message passing neural network (ML-MPNN) for graph-based learning and a contrastive-BERT model with masked embedding recovery for sequence modeling, optimized via novel AUC-focused loss functions, achieving state-of-the-art performance, including #1 ranking in ROC-AUC and PRC-AUC on the AI Cures Open Challenge for COVID-19 drug discovery.

ABSTRACT

Properties of molecules are indicative of their functions and thus are useful in many applications. With the advances of deep learning methods, computational approaches for predicting molecular properties are gaining increasing momentum. However, there lacks customized and advanced methods and comprehensive tools for this task currently. Here we develop a suite of comprehensive machine learning methods and tools spanning different computational models, molecular representations, and loss functions for molecular property prediction and drug discovery. Specifically, we represent molecules as both graphs and sequences. Built on these representations, we develop novel deep models for learning from molecular graphs and sequences. In order to learn effectively from highly imbalanced datasets, we develop advanced loss functions that optimize areas under precision-recall curves. Altogether, our work not only serves as a comprehensive tool, but also contributes towards developing novel and advanced graph and sequence learning methodologies. Results on both online and offline antibiotics discovery and molecular property prediction tasks show that our methods achieve consistent improvements over prior methods. In particular, our methods achieve #1 ranking in terms of both ROC-AUC and PRC-AUC on the AI Cures Open Challenge for drug discovery related to COVID-19. Our software is released as part of the MoleculeX library under AdvProp.

Motivation & Objective

• To develop a unified, comprehensive framework for molecular property prediction that integrates graph and sequence representations.
• To address the challenge of highly imbalanced datasets in molecular property prediction through novel AUC-optimized loss functions.
• To improve model generalization and performance by incorporating self-supervised pre-training on large-scale unlabeled molecular sequences.
• To establish new state-of-the-art results in molecular property prediction, particularly in drug discovery under real-world constraints.
• To release a publicly available, modular toolset (AdvProp) within the MoleculeX library for reproducible and scalable research.

Proposed method

• Proposes a multi-level message passing neural network (ML-MPNN) that aggregates information across nodes, edges, subgraphs, and the full molecular graph to capture hierarchical structural dependencies.
• Introduces contrastive-BERT, a sequence-based model that uses a masked embedding recovery task with contrastive learning to pre-train on large-scale unlabeled SMILES sequences.
• Employs surrogate loss functions for optimizing areas under precision-recall curves (PRC-AUC), including an AP-based surrogate loss with margin-based ranking to improve model robustness on imbalanced data.
• Uses stochastic optimization via SOAP-ADAM and PESG to directly optimize AUC metrics, replacing standard cross-entropy loss for better alignment with evaluation criteria.
• Applies a contrastive loss function based on cosine similarity between masked token predictions and their context embeddings to enhance representation learning in the sequence model.
• Integrates global molecular features (200D, extracted via RDKit) with graph-level representations through a 2-layer MLP for final prediction.

Experimental results

Research questions

• RQ1 Can a unified framework combining graph and sequence deep learning models improve molecular property prediction beyond existing state-of-the-art methods?
• RQ2 How effective are novel AUC-optimized loss functions in handling highly imbalanced molecular property datasets compared to standard cross-entropy loss?
• RQ3 To what extent does self-supervised pre-training via masked embedding recovery enhance downstream molecular property prediction performance?
• RQ4 Can multi-level message passing in graph neural networks capture richer structural information than standard message passing, leading to improved performance?
• RQ5 Does the combination of graph and sequence modeling provide complementary signals that lead to consistent performance gains across diverse molecular prediction tasks?

Key findings

• AdvProp achieved the #1 ranking in both ROC-AUC and PRC-AUC on the AI Cures Open Challenge for COVID-19 drug discovery, outperforming all other submitted methods.
• The proposed ML-MPNN model demonstrated state-of-the-art performance on molecular property prediction benchmarks, particularly in capturing complex structural dependencies across multiple graph levels.
• The contrastive-BERT model with masked embedding recovery achieved strong performance on downstream tasks, showing the effectiveness of self-supervised pre-training on large-scale unlabeled SMILES data.
• AUC-optimized loss functions, particularly the AP-based surrogate loss, significantly improved model performance on imbalanced datasets, especially in precision-recall space.
• The integration of global molecular features with graph-level representations led to consistent improvements across multiple datasets and tasks.
• The open-source release of AdvProp via the MoleculeX library enables reproducible and scalable use of the proposed models and tools in drug discovery research.