[Paper Review] Fragment-based Pretraining and Finetuning on Molecular Graphs

Kha-Dinh Luong, Ambuj K. Singh|arXiv (Cornell University)|Oct 5, 2023

Computational Drug Discovery Methods4 citations

TL;DR

This paper proposes GraphFP, a fragment-based pretraining framework for molecular graphs that jointly trains molecular and fragment GNNs via contrastive learning to capture structural information at multiple resolutions. By mining a compact vocabulary of chemically meaningful fragments and enforcing consistency between fragment embeddings and aggregated node embeddings, GraphFP improves performance on 5 out of 8 molecular benchmarks and boosts long-range biological benchmark results by at least 11.5%.

ABSTRACT

Property prediction on molecular graphs is an important application of Graph Neural Networks. Recently, unlabeled molecular data has become abundant, which facilitates the rapid development of self-supervised learning for GNNs in the chemical domain. In this work, we propose pretraining GNNs at the fragment level, a promising middle ground to overcome the limitations of node-level and graph-level pretraining. Borrowing techniques from recent work on principal subgraph mining, we obtain a compact vocabulary of prevalent fragments from a large pretraining dataset. From the extracted vocabulary, we introduce several fragment-based contrastive and predictive pretraining tasks. The contrastive learning task jointly pretrains two different GNNs: one on molecular graphs and the other on fragment graphs, which represents higher-order connectivity within molecules. By enforcing consistency between the fragment embedding and the aggregated embedding of the corresponding atoms from the molecular graphs, we ensure that the embeddings capture structural information at multiple resolutions. The structural information of fragment graphs is further exploited to extract auxiliary labels for graph-level predictive pretraining. We employ both the pretrained molecular-based and fragment-based GNNs for downstream prediction, thus utilizing the fragment information during finetuning. Our graph fragment-based pretraining (GraphFP) advances the performances on 5 out of 8 common molecular benchmarks and improves the performances on long-range biological benchmarks by at least 11.5%. Code is available at: https://github.com/lvkd84/GraphFP.

Motivation & Objective

• Address the limitations of node-level and graph-level pretraining in GNNs for molecular property prediction.
• Overcome the issue of over-smoothing and loss of structural fidelity in existing fragment-based methods by introducing a dual-encoder framework.
• Develop a chemically faithful, self-supervised pretraining strategy that avoids data augmentation and privileged information like 3D coordinates.
• Improve downstream performance on diverse molecular and long-range biological benchmarks through fragment-level representation learning.
• Enable better generalization by learning hierarchical structural representations using a principled fragment vocabulary from principal subgraph mining.

Proposed method

• Use principal subgraph mining to extract a compact, diverse, and chemically meaningful vocabulary of prevalent fragments from a large pretraining dataset.
• Train two separate GNNs: one on molecular graphs (node-level) and one on fragment graphs (higher-order connectivity), using a contrastive learning objective.
• Enforce consistency between the fragment embedding and the aggregated embedding of its constituent atoms in the molecular graph to align multi-scale representations.
• Introduce predictive pretraining tasks that use fragment graph-derived labels (e.g., structural properties) to enhance molecular GNN understanding.
• Employ both pretrained molecular and fragment GNNs during fine-tuning to leverage fragment-level structural knowledge.
• Utilize fragment graphs not only for contrastive learning but also to generate auxiliary labels for graph-level predictive pretraining.

Figure 1: Fragment-based contrastive pretraining framework. $\textsc{GNN}_{M}$ processes molecular graphs while $\textsc{GNN}_{F}$ processes fragment graphs. The fragment-based pooling function FragPool aggregates node embeddings into a combined embedding that forms a positive contrastive pair with

Experimental results

Research questions

• RQ1 Can fragment-level pretraining improve molecular graph representation learning when compared to node-level or graph-level pretraining?
• RQ2 How does joint pretraining of molecular and fragment GNNs via contrastive learning affect downstream performance on diverse molecular benchmarks?
• RQ3 What impact does the size and quality of the fragment vocabulary have on the effectiveness of self-supervised pretraining?
• RQ4 Can fragment-based representations enhance performance on long-range biological property prediction tasks?
• RQ5 How does the absence of data augmentation or privileged information (e.g., 3D coordinates) affect the validity and generalization of the learned representations?

Key findings

• GraphFP achieves state-of-the-art or competitive performance on 5 out of 8 standard molecular property prediction benchmarks.
• The method improves performance on long-range biological benchmarks by at least 11.5% compared to baseline methods.
• The t-SNE visualization confirms that fragment embeddings capture higher-order structural patterns, with distinct clustering by molecular backbone.
• Atom embeddings within the same fragment cluster together in the latent space, indicating that pooled node representations effectively encode higher-order connectivity.
• Increasing the fragment vocabulary size beyond an optimal point (e.g., 3,200) leads to performance degradation, likely due to increased parameter count and loss of structural fidelity.
• The choice of GNN architecture significantly affects performance, with GIN showing the best average performance across benchmarks.

Figure 2: Fragment-based predictive pretraining. $\textsc{GNN}_{M}$ processes molecular graphs and produces graph-level embeddings used for prediction. The right upper box shows unique fragments that exist in the input molecule while the right lower box shows the ground-truth structural backbone. A