Source-linked AI summary
Motif-based Graph Self-Supervised Learning for Molecular Property Prediction
Zaixi Zhang, Qi Liu, Hao Wang, Chengqiang Lu, Chee-Kong Lee
TL;DR
Molecular GNNs need scarce, expensive labels, while existing self-supervised methods often miss semantic information in graph motifs. The paper proposes MGSSL with chemistry-informed motif extraction, motif-by-motif generation, and multi-level pre-training, and reports superior performance across downstream benchmarks. Its motif vocabulary construction is constrained by fragmentation assumptions and by BRICS producing many rare, redundant fragments.
Problem
Molecular GNNs are data-hungry because labeled molecules are scarce and expensive, while existing self-supervised tasks often fail to exploit informative graph motifs.
Method
MGSSL uses BRICS-based fragmentation with additional rules, motif-by-motif generation with breadth-first or depth-first orders, and adaptive multi-level self-supervised pre-training.
Results
MGSSL achieves the best performance on 7 of 8 downstream benchmarks, with faster training and testing convergence than selected baselines.
Takeaways & Limitations
Motif-based and multi-level pre-training enables GNNs to capture molecular information at motif and multiple structural scales for downstream prediction.
Takeaways & Limitations
BRICS alone produces more than 100k unique fragments, many appearing fewer than five times, requiring additional fragmentation rules to control vocabulary redundancy.
Abstract
from arXiv · showhide
Predicting molecular properties with data-driven methods has drawn much attention in recent years. Particularly, Graph Neural Networks (GNNs) have demonstrated remarkable success in various molecular generation and prediction tasks. In cases where labeled data is scarce, GNNs can be pre-trained on unlabeled molecular data to first learn the general semantic and structural information before being fine-tuned for specific tasks. However, most existing self-supervised pre-training frameworks for GNNs only focus on node-level or graph-level tasks. These approaches cannot capture the rich information in subgraphs or graph motifs. For example, functional groups (frequently-occurred subgraphs in molecular graphs) often carry indicative information about the molecular properties. To bridge this gap, we propose Motif-based Graph Self-supervised Learning (MGSSL) by introducing a novel self-supervised motif generation framework for GNNs. First, for motif extraction from molecular graphs, we design a molecule fragmentation method that leverages a retrosynthesis-based algorithm BRICS and additional rules for controlling the size of motif vocabulary. Second, we design a general motif-based generative pre-training framework in which GNNs are asked to make topological and label predictions. This generative framework can be implemented in two different ways, i.e., breadth-first or depth-first. Finally, to take the multi-scale information in molecular graphs into consideration, we introduce a multi-level self-supervised pre-training. Extensive experiments on various downstream benchmark tasks show that our methods outperform all state-of-the-art baselines.
1 Introduction
Molecular GNNs perform well but require scarce labeled data, motivating self-supervised pre-training. MGSSL addresses the limited use of graph motifs through motif extraction, motif generation, and multi-level pre-training.
- Motivation: GNNs require substantial labeled molecular data, but labeled molecules represent only a small portion of chemical space because obtaining labels is expensive.Wet-lab experiments and quantum chemistry calculations are time-consuming and costly.
- Motivation: Small labeled datasets make supervised GNN training prone to over-fitting and poor generalization to out-of-distribution molecules.
- Motivation: Existing GNN self-supervised tasks often overlook graph motifs, despite motifs carrying semantic information indicative of whole-graph characteristics.Functional groups such as hydroxide can indicate molecular properties including water solubility.
- Challenges: Motif-level self-supervision is challenging because molecular motif mining must preserve chemical validity, motif generation differs from node-by-node graph generation, and multi-level tasks must be unified.
- Approach: MGSSL extracts motifs with BRICS and additional fragmentation rules, generates molecules motif-by-motif using breadth-first or depth-first orders, and adaptively combines multi-level pre-training tasks.The framework makes topology and attribute predictions and adjusts task weights with the Frank-Wolfe algorithm.
2 Preliminaries and Related Work
Molecular property prediction uses graph-based machine learning to represent atoms and bonds, while graph self-supervised learning learns representations from unlabeled graphs. Existing motif-based approaches use motifs but fail to model their topology.
- Molecular Property Prediction: Molecular property prediction is important across physics, chemistry, and materials science, while density functional theory is computationally expensive with approximate complexity O(N^3).
- Molecular Property Prediction: Graph methods represent molecules through atoms and bonds, and message-passing GNNs model node and edge interactions for molecular property prediction.
- Graph Neural Networks: GNNs iteratively aggregate neighboring node and edge information, producing node embeddings whose k-th layer captures structural information within each node’s k-hop neighborhood.
- Graph Neural Networks: A graph-level embedding is obtained by applying a permutation-invariant READOUT function, such as averaging, summation, or maximization, to final node representations.
- Graph Self-Supervised Learning: Graph self-supervised learning learns intermediate representations of unlabeled graphs for unknown downstream tasks using objectives such as graph proximity or mutual-information maximization.
- Related Work: Grover and MICRO-Graph exploit graph motifs through classification or contrastive learning, but both methods fail to consider motif topology.
3 Motif-based Graph Self-supervised Learning
MGSSL combines chemistry-inspired molecule fragmentation, motif-tree generation, and multi-level self-supervised pre-training for GNNs. It predicts motif topology and labels through autoregressive generation in either DFS or BFS order.
- MGSSL has three components: molecule fragmentation, motif generation, and multi-level self-supervised pre-training for GNNs.The framework is organized around chemistry-inspired preprocessing, motif-tree modeling, and hierarchical pre-training.
- Chemistry-inspired Molecule Fragmentation: Molecules are fragmented into non-overlapping, semantically meaningful motifs whose union reconstructs the molecule graph and whose vocabulary is sufficiently frequent and moderate in size.Each motif tree represents motifs as nodes and their relative spatial relationships as edges.
- Chemistry-inspired Molecule Fragmentation: BRICS cleavage is followed by two rules that reduce ring variants and break side chains, shrinking motif-vocabulary size and increasing motif occurrence frequency.The rules break bonds between ring and non-ring atoms and split non-ring atoms with at least three neighbors.
- Motif Generation: The motif-generation objective models the likelihood of labeled and connected motif trees, using topology and motif-label prediction heads during autoregressive construction.At each step, previously generated motif attributes and structures condition the next motif and its connections.
- Motif Generation: MGSSL supports DFS and BFS generation orders: DFS recurses and backtracks through children, whereas BFS generates motif nodes layer-wise.The model first predicts whether a visited motif has children, then predicts each generated child’s motif label.
- Motif Generation: During generation, message vectors pass information among existing motifs, while topology prediction uses motif embeddings and messages and label prediction outputs a distribution over the motif vocabulary.Motif embeddings are obtained by pooling atom embeddings within each motif, and message updates use a GRU.
4 Experimental Results
MGSSL is evaluated against established self-supervised methods, supervised GIN finetuning, alternative generation orders, and multiple GNN architectures across molecular benchmarks. Results support motif-based and multi-level pre-training, while fragmentation choices influence motif vocabulary quality.
- Experimental setup: 250k unlabeled ZINC15 molecules support pre-training, followed by evaluation on 8 binary classification benchmarks from MoleculeNet.The experiments use RDKit to preprocess molecular SMILES strings.
- Experimental setup: MGSSL is compared with five self-supervised GNN pre-training methods, alongside direct supervised finetuning and BFS or DFS motif generation.The baselines include Deep Graph Infomax, attribute masking, GCC, Grover, and GPT-GNN.
- Downstream results: MGSSL achieves the best performance on 7 of 8 benchmarks, while both BFS and DFS improve downstream results and BFS has a small average edge.The authors use BFS as the default because layer-wise motif generation may help GNNs learn more motif structure.
- Convergence: MGSSL pre-trained GNNs converge faster in training and testing than baseline methods on the four benchmark datasets shown in Figure 4.Because pre-training is a one-time effort, the resulting models can be reused across downstream tasks with minimal finetuning overhead.
- Architecture robustness: Motif-based pre-training benefits all five tested GNN architectures, with GIN obtaining the largest relative gain and best post-pre-training performance.The architectures tested are GIN, GCN, RGCN, GraphSAGE, and DAGNN, averaged across 8 benchmarks.
- Ablations and fragmentation: Multi-level pre-training has larger average ROC-AUC than motif-only and sequential variants, while fragmentation controls motif-vocabulary size and frequency distribution.BRICS alone produces a vocabulary exceeding 100k motifs, with more than 90% occurring fewer than five times.
5 Conclusion and Future Works
The paper proposes MGSSL, a motif-generation pre-training method for GNNs, and combines motif-level learning with multi-level pre-training. Experiments report that MGSSL outperforms state-of-the-art baselines across downstream benchmark tasks.
- MGSSL pre-trains GNNs with a novel motif generation task to capture semantic and structural information in graph motifs.
- Retrosynthesis-based fragmentation with two additional rules derives semantically meaningful molecular motifs.
- The framework supports BFS and DFS generation orders, requiring topology and motif label predictions at each step.
- Multi-level self-supervised pre-training unifies hierarchical self-supervised tasks for GNNs.
- MGSSL outperforms all state-of-the-art baselines on various downstream benchmark tasks.