PaperlayerView exact revision on arXiv

Source-linked AI summary

ASGN: An Active Semi-supervised Graph Neural Network for Molecular Property Prediction

Zhongkai Hao, Chengqiang Lu, Zheyuan Hu, Hao Wang, Zhenya Huang, Qi Liu, Enhong Chen, Cheekong Lee

arXiv:2007.03196v1cs.LGstat.ML

TL;DR

Molecular property prediction is constrained by scarce, expensive labels and the challenges of learning molecular representations from mostly unlabeled chemical space. ASGN addresses this with an iterative teacher-student framework and diversity-based active learning, and it reports better performance and greater label efficiency than baselines when labels are limited.

  • Problem

    Molecular property prediction has few labeled molecules because obtaining properties through experiments or DFT is expensive, while semi-supervised learning must handle both molecular graph information and labeled-unlabeled imbalance.

  • Method

    ASGN alternates teacher and student graph neural networks, learning general representations from labeled and unlabeled molecules before fine-tuning property prediction and selecting diverse molecules for new labels.

  • Results

    ASGN achieves better performance than baselines with limited labels, including more than 50% reductions on several properties and about 2–3 times greater label efficiency than baselines.

  • Takeaways & Limitations

    The results support combining semi-supervised representation learning with active learning for molecular property prediction under label scarcity.

Abstract

from arXiv · show

Molecular property prediction (e.g., energy) is an essential problem in chemistry and biology. Unfortunately, many supervised learning methods usually suffer from the problem of scarce labeled molecules in the chemical space, where such property labels are generally obtained by Density Functional Theory (DFT) calculation which is extremely computational costly. An effective solution is to incorporate the unlabeled molecules in a semi-supervised fashion. However, learning semi-supervised representation for large amounts of molecules is challenging, including the joint representation issue of both molecular essence and structure, the conflict between representation and property leaning. Here we propose a novel framework called Active Semi-supervised Graph Neural Network (ASGN) by incorporating both labeled and unlabeled molecules. Specifically, ASGN adopts a teacher-student framework. In the teacher model, we propose a novel semi-supervised learning method to learn general representation that jointly exploits information from molecular structure and molecular distribution. Then in the student model, we target at property prediction task to deal with the learning loss conflict. At last, we proposed a novel active learning strategy in terms of molecular diversities to select informative data during the whole framework learning. We conduct extensive experiments on several public datasets. Experimental results show the remarkable performance of our ASGN framework.

1 INTRODUCTION

Molecular property prediction is important but limited by the high cost of DFT labels and the scarcity of labeled molecules. ASGN addresses these constraints with a teacher-student semi-supervised framework that combines labeled and unlabeled molecules with diversity-based active learning.

  • Motivation: DFT provides accurate molecular-property solutions but can require an hour for a molecule with only 20 atoms.Its computational cost limits screening across large molecular sets.
  • Motivation: Graph neural networks represent molecules as graphs of atoms and interactions, but typically require many labeled molecules for training.Labeled molecules form only a small portion of chemical space because labels require experiments or DFT calculations.
  • Challenges: Semi-supervised molecular graph learning must jointly capture node- and graph-level information while handling the severe imbalance between labeled and unlabeled molecules.Methods developed for domains such as images and text cannot be directly transferred to this setting.
  • ASGN: ASGN uses a teacher-student framework to exploit both labeled and unlabeled molecules for molecular property prediction.The teacher learns a general representation from molecular features at local and global scales.
  • ASGN: The student focuses on property prediction, while diversity-based active learning repeatedly selects informative molecules for labeling and model fine-tuning.This separates representation learning from property prediction and uses teacher embeddings to select a diversified subset.

2 RELATED WORK

Prior work spans simulation-based, data-driven, graph-based, semi-supervised, and active-learning approaches to molecular prediction. ASGN builds on these directions by using semi-supervised molecular embeddings for diversity-based active selection.

  • Molecular Property Prediction: DFT methods derive from quantum theory but are time-consuming, with complexity O(N^4) in the number of atoms.Data-driven methods reduce computational demands but handcrafted features can limit accuracy and model capacity.
  • Molecular Property Prediction: Graph neural networks model molecules as graphs and provide architectures for processing non-Euclidean molecular data.MPNN and SchNet are representative graph-based methods used for molecular property prediction.
  • Semi-supervised Representation learning: Semi-supervised representation learning incorporates unlabeled data to learn general representations, including through generation or consistency-based pseudo-labeling.These approaches have been developed across domains such as images and text.
  • Active Learning: Active learning iteratively selects batches of unlabeled data for labeling to supplement a limited labeled set.Selection strategies commonly use uncertainty or diversity criteria.
  • Active Learning: ASGN uses semi-supervised embeddings to measure molecular distances for diversity-based active selection because uncertainty-selected batches can be nearly identical.This strategy targets informative molecule selection in large datasets.

3 DEFINITIONS AND NOTATIONS

The paper models each molecule as a weighted graph and partitions chemical space into labeled and unlabeled subsets. The prediction task is to learn a model that maps molecular graphs to properties using limited labels.

  • Molecular Graph: A molecular graph is a weighted graph G = (V, E) whose vertices represent atoms and whose edges connect atom pairs.Node features encode atom types, while edge weights encode coordinate distances.
  • Chemical Space: Chemical space is a set of molecules M = {G_i}, partitioned into labeled D_l and unlabeled D_u subsets.Labeled molecules have examined property vectors, whereas unlabeled molecules have unknown properties.
  • Prediction Problem: The formal task is to learn a model f(G) → y that predicts molecular properties from limited labels.The property vector y contains real-valued molecular properties.

4 ASGN: ACTIVE SEMI-SUPERVISED GRAPH NEURAL NETWORK

The ASGN section introduces the framework and its components, beginning with an iterative teacher-student architecture for semi-supervised molecular representation learning and property prediction.

  • Framework: ASGN is described as a framework whose components are presented after an overview of its overall design.The supplied passage introduces the framework section but does not specify additional component details.

4.1 Framework

ASGN combines labeled and unlabeled molecules through an iterative teacher–student graph neural network framework. Active learning selects representative molecules for labeling, while pseudo-labels provide feedback to the teacher.

  • ASGN incorporates both labeled and unlabeled molecules for molecular property prediction.
  • Teacher and student graph neural networks work iteratively, with the teacher learning general molecular representations and the student fine-tuning transferred parameters for property prediction.
  • The student assigns pseudo-labels to unlabeled molecules, which the teacher uses as feedback in subsequent training.
  • Active learning selects diversified unlabeled molecules for ground-truth labeling, adds them to the labeled set, and repeats model fine-tuning until the accuracy budget is reached.

4.2 Semi-supervised Teacher Model

The semi-supervised teacher learns molecular representations from labeled and unlabeled graphs by combining property prediction with local structural reconstruction and graph-level distribution learning. Its architecture uses message passing, pooling, and clustering-based self-supervision.

  • The teacher combines labeled property loss with unsupervised graph- and node-level losses over labeled and unlabeled molecules.
  • Message Passing Graph Neural Network: Message passing embeds nodes and updates them by aggregating neighboring representations along weighted molecular edges.
  • Message Passing Graph Neural Network: After message passing, node embeddings are pooled into a graph embedding and passed through an MLP to predict molecular properties.
  • Node Level Representation Learning: Node-level representation learning reconstructs sampled atom types and edge weights from node embeddings using separate reconstruction terms.
  • Node Level Representation Learning: Randomly sampling O(|G|) edges provides an efficiency trade-off for reconstruction because fully connected molecular graphs contain redundant information.
  • Graph Level Representation Learning: Graph-level representation learning predicts cluster distributions and optimizes their cross-entropy under a predefined prior, using an implicit clustering process with optimal transport and Sinkhorn-Knopp optimization.

4.3 Supervised Student Model

The student model resolves the teacher’s competing representation and property objectives by fine-tuning transferred weights only on labeled data. It then generates pseudo-labels that feed back into later teacher training.

  • Jointly optimizing the teacher’s objectives can weaken property prediction because representation and property losses compete, especially when labeled molecules are scarce.
  • The student transfers the teacher’s weights and fine-tunes only on the labeled dataset to focus on molecular property prediction.
  • The student predicts properties for all unlabeled molecules and uses those predictions as pseudo-labels in the next teacher iteration.
  • The teacher learns general representations while the student targets accurate property prediction, with teacher pre-training providing a warm start for the student.

4.4 Active Learning for Data Selection

ASGN uses active learning to expand the labeled set with informative molecules selected from teacher embeddings. The selection seeks molecular diversity through a greedy k-center approximation and obtains labels using DFT or equivalent ground-truth computation.

  • Each iteration selects molecules from the unlabeled set for ground-truth labeling and adds them to the labeled set for subsequent model fine-tuning.
  • The active-learning strategy seeks a small batch of molecules that is maximally diversified in chemical space.
  • Because k-DPP subset selection is NP-hard, ASGN uses a greedy k-center approximation to choose molecules farthest from the labeled set.
  • Molecular distances are computed with the L−2 norm between teacher-model representations, avoiding handcrafted distances such as graph edit distance.

4.5 Method Summary and Discussion

ASGN iteratively combines teacher representation learning, k-center active labeling, student property prediction, and pseudo-labeling to exploit labeled and unlabeled molecules. Weight transfer and feedback connect the teacher and student models throughout training.

  • Framework workflow: The ASGN framework alternates active labeling, teacher-to-student weight transfer, student fine-tuning, and pseudo-label assignment for the remaining unlabeled molecules.Each iteration adds a newly labeled batch, transfers teacher weights, fine-tunes the student, and uses the student to assign property pseudo-labels.
  • Discussion: Separating teacher and student models is intended to alleviate loss conflict, while weight transfer provides a warm start that avoids overfitting and accelerates training.The models also communicate through transferred weights and pseudo-label feedback.
  • Teacher model: The teacher is trained with representation objectives L = Lr + Lc + Lp to produce graph embeddings for the unlabeled dataset.The resulting embeddings are used as the basis for subsequent active-data selection.
  • Active learning: k-center active learning selects a new batch using graph embeddings, adding diverse molecules to the labeled set.The selected batch has size b and is queried for labels before student fine-tuning.
  • Student model: The student is fine-tuned by minimizing the property-prediction loss Lp on labeled data, then predicts properties for the remaining unlabeled molecules.The student model is ultimately applied to molecular property prediction, while its pseudo-label feedback supports continued framework learning.

5 EXPERIMENTS

Experiments evaluate ASGN on QM9 and OPV under limited-label effectiveness, label-efficiency, ablation, and representation-visualization settings. ASGN improves molecular property prediction, reduces labeling needs, and benefits from its teacher-student design, weight transfer, and learned graph representations.

  • Experimental setup: Experiments use QM9 and OPV to evaluate ASGN with limited labeled molecules and unlabeled data.QM9 contains approximately 130,000 molecules, while OPV contains roughly 100,000 medium-sized molecules.
  • Effectiveness experiment: With a 5,000-label budget, ASGN reduces MAE by more than 50% for several QM9 and OPV properties versus the state-of-the-art method.The reported properties include U, U0, α, and Cv.
  • Effectiveness experiment: Semi-supervised reconstruction outperforms supervised MPGNN and Mean-Teachers, while graph-level representation learning surpasses InfoGraph.The authors attribute these results to reconstruction capturing molecular domain knowledge and global representation learning benefiting property prediction.
  • Efficiency experiment: At fixed label numbers, ASGN achieves lower MAE than baselines and performs better than a fully supervised model trained on all labeled data.The authors connect this result to the active-learning strategy and semi-supervised loss used as regularization.
  • Efficiency experiment: At fixed error, ASGN is about 2 ∼3 times label efficient than baselines, using 50% of labels on QM9 and 40% on OPV to reach full accuracy.The study measures label-rate–MAE curves while progressively adding selected unlabeled molecules to the labeled set.
  • Ablation and visualization experiments: Ablations show that the student improves property prediction, the teacher-student framework is necessary, and weight transfer makes training and testing MAE faster, more stable, and ultimately better.t-SNE visualization also shows clustered molecular embeddings with visible distances between clusters.

6 CONCLUSIONS

ASGN improves molecular property prediction with limited labels by incorporating unlabeled molecules through a teacher-student framework and semi-supervised representation learning. Diversity-based active learning selects new molecules for labeling, while ablations examine the necessity of the framework’s components.

  • The framework uses two graph neural networks in an iterative teacher-student design, with semi-supervised representation learning for molecular graphs.
  • Weight transfer and pseudo labeling optimize the two models to balance their loss functions.
  • Diversity-based active learning selects new molecules for labeling during the framework’s learning process.
  • ASGN achieves much better performance than baselines when labels are limited, and ablation experiments assess the necessity of its components.
Loading 2007.03196v1…