PaperlayerView exact revision on arXiv

Source-linked AI summary

Molecular Contrastive Learning with Chemical Element Knowledge Graph

Yin Fang, Qiang Zhang, Haihong Yang, Xiang Zhuang, Shumin Deng, Wen Zhang, Ming Qin, Zhuo Chen, Xiaohui Fan, Huajun Chen

arXiv:2112.00544v2cs.LGcs.AIq-bio.QM

TL;DR

Molecular graph contrastive learning provides a self-supervised route to molecular representations, but prior methods often omit domain knowledge and relationships between unbonded atoms. KCL builds a Chemical Element KG and uses it to augment molecular graphs and encode a knowledge-aware view; it is reported to outperform state-of-the-art methods across eight datasets.

  • Problem

    Graph contrastive learning can avoid labeled-data requirements, but prior methods overlook fundamental domain knowledge and correlations between atoms that share attributes without being bonded.

    Source: §1 · p. 1

  • Method

    KCL constructs a Chemical Element KG and combines KG-guided graph augmentation, knowledge-aware message passing, and a contrastive objective across original and augmented molecular graph views.

    Source: §1 · p. 2

  • Results

    KCL is reported to outperform state-of-the-art methods on eight molecular datasets under both fine-tune and linear protocols.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The Chemical Element KG is proposed as a resource that can assist molecular learning tasks beyond those evaluated in this paper.

    Source: §1 · p. 2

    The authors plan to enrich the KG with finer-grained knowledge and more OWL2 properties and axioms, and to continue improving its quality and scale.

    Source: §5 · p. 7

Abstract

from arXiv · show

Molecular representation learning contributes to multiple downstream tasks such as molecular property prediction and drug design. To properly represent molecules, graph contrastive learning is a promising paradigm as it utilizes self-supervision signals and has no requirements for human annotations. However, prior works fail to incorporate fundamental domain knowledge into graph semantics and thus ignore the correlations between atoms that have common attributes but are not directly connected by bonds. To address these issues, we construct a Chemical Element Knowledge Graph (KG) to summarize microscopic associations between elements and propose a novel Knowledge-enhanced Contrastive Learning (KCL) framework for molecular representation learning. KCL framework consists of three modules. The first module, knowledge-guided graph augmentation, augments the original molecular graph based on the Chemical Element KG. The second module, knowledge-aware graph representation, extracts molecular representations with a common graph encoder for the original molecular graph and a Knowledge-aware Message Passing Neural Network (KMPNN) to encode complex information in the augmented molecular graph. The final module is a contrastive objective, where we maximize agreement between these two views of molecular graphs. Extensive experiments demonstrated that KCL obtained superior performances against state-of-the-art baselines on eight molecular datasets. Visualization experiments properly interpret what KCL has learned from atoms and attributes in the augmented molecular graphs. Our codes and data are available at https://github.com/ZJU-Fangyin/KCL.

1 Introduction

Molecular graph contrastive learning offers a self-supervised approach, but existing methods overlook domain knowledge and relationships between atoms without bonds. KCL uses a Chemical Element KG to add such knowledge to molecular representations.

1 Introduction

  • Existing graph contrastive methods focus on graph structure and can miss domain knowledge and correlations between unbonded atoms with common attributes.For molecular graphs, adding or removing a bond or functional group can drastically change molecular identity and properties.Source: §1 · p. 1
  • The Chemical Element KG guides augmentation to connect atoms sharing chemical attributes even when they are not directly bonded.It is based on the Periodic Table of Elements and represents element relations and attributes such as periodicity and metallicity.Source: §1 · p. 2, Figure 1 · p. 2
  • KCL uses a common graph encoder for original molecules and KMPNN for heterogeneous message passing over knowledge in augmented graphs, then contrasts the two views.Its objective maximizes agreement between positive pairs and discrepancy between hard negatives.Source: §1 · p. 2
  • The Chemical Element KG describes relations between elements and their chemical attributes and is proposed for molecular learning tasks beyond those evaluated here.The paper identifies the KG as a reusable resource for other molecular learning tasks.Source: §1 · p. 2
  • KCL is reported to outperform state-of-the-art methods on eight molecular datasets under both fine-tune and linear protocols.The framework combines knowledge-guided augmentation, knowledge-aware representation, and a contrastive objective.Source: §1 · p. 2

2 Related Works

Prior molecular representation work includes supervised graph models and self-supervised pretraining, while graph contrastive learning generates views through perturbations. The related work notes that such perturbations can harm chemical knowledge and that prior methods omit fundamental domain knowledge.

Molecular Representation Learning

  • Supervised molecular representation models use graph topology and node and edge attributes, but require costly annotations and may generalize poorly to unseen molecules.Earlier approaches also used convolutional networks for molecular fingerprints and recurrent networks over SMILES strings.Source: §2 · p. 2

Self-Supervised Learning on Graphs

  • Self-supervised molecular pretraining learns representations through vertex-attribute, neighborhood-context, node-attribute, contextual-property, and motif prediction tasks.GROVER uses a Transformer-style architecture to predict contextual properties and motif information.Source: §2 · p. 2

Contrastive Learning on Graphs

  • Graph contrastive learning forms positive and negative pairs from augmented views, but random perturbations may damage chemical knowledge and prior approaches overlook fundamental domain knowledge.MoCL uses substructure substitution with two-level knowledge, while CKGNN selects positive pairs using fingerprint similarity.Source: §2 · p. 2

3 Methodology

KCL learns molecular representations by contrasting original graphs with knowledge-augmented views. Its augmentation preserves molecular topology while adding element attributes and nonbonded atom associations, which a specialized encoder processes alongside the original graph.

3.1 Problem Formulation

  • Molecules are represented as atom-and-bond graphs, and the self-supervised objective is to learn graph encoders without labels for downstream tasks.The formulation maps each input graph to a vector representation.Source: §3.1 · pp. 2–3

3.2 Overview

  • KCL contrasts representations of original and knowledge-augmented molecular graphs, training encoders to align positives and separate hard negatives.The framework has knowledge-guided augmentation, knowledge-aware representation, and a contrastive objective.Source: §3.2 · p. 3

3.3 Knowledge-guided Graph Augmentation

  • The Chemical Element KG encodes element–attribute triples, discretizes continuous attributes, and adds one-hop attribute nodes while preserving the molecular structure.The augmented graph also links atoms with common attributes even when they are not directly bonded.Source: §3.3 · p. 3

3.4 Knowledge-aware Graph Representation

  • A GCN encodes the original graph, while KMPNN handles the augmented graph using distinct message functions for atom and attribute neighbors and attention over neighbors.KMPNN applies K rounds of message passing and a graph-level readout.Source: §3.4 · p. 4

3.5 Contrastive Objective

  • KCL projects both graph views into a contrastive space, using NT-Xent with similar molecules selected as hard negatives.Molecules are ranked by fingerprint similarity, and the most similar batch members are selected as negatives.Source: §3.5 · pp. 4–5

4 Experiments

The experiments evaluate KCL on molecular-property benchmarks, comparing its full framework and components with existing methods. Results favor KCL across fine-tune and linear protocols, while ablations and attention visualizations examine the contributions and learned patterns.

4.1 Experimental Setup

  • The setup pre-trains KCL on 250K unlabeled ZINC15 molecules and evaluates on eight MoleculeNet datasets using fine-tune and linear protocols.Fine-tuning updates the encoders and an MLP, whereas the linear protocol fixes pretrained graph embeddings and trains a linear classifier.Source: §4.1 · p. 5

4.2 Performance Comparison (Q1 & Q2)

  • KCL achieves the best performance across all fine-tune datasets, with 7.1% overall relative improvement over state-of-the-art baselines.The reported relative improvement is 2.6% on classification tasks and 20.4% on regression tasks; on FreeSolv, with 642 labeled molecules, the improvement is 16.8%.Source: §4.2 · p. 6, Table 2 · p. 6
  • Under the linear protocol, both KCL versions outperform alternative graph augmentations, and KCL improves 7.0% over prior contrastive-learning methods.The linear-protocol comparison uses classification datasets; regression results are deferred to Appendix E.1.Source: §4.2 · p. 6, Table 3 · p. 6

4.3 Ablation Study (Q3 & Q4)

  • Knowledge-feature initialization and hard-negative mining together yield the best ablation performance; removing either component decreases performance.The model without both components almost always performs worst; the authors associate initialization with capturing structural triple information and hard-negative mining with more discriminative representations.Source: Figure 3 · p. 6, §4.3 · p. 7
  • KMPNN outperforms RGCN and MPNN on augmented graphs, while changing the encoder for original molecular graphs has no significant performance impact.The authors attribute KMPNN's advantage to heterogeneous attentive message passing for different types of knowledge.Source: §4.3 · p. 7
  • Contrastive learning increases average performance by 8.5% on classification and 56.9% on regression compared with KCL without contrastive learning.The comparison is under the fine-tune protocol; the linear-protocol counterpart is deferred to Appendix E.4.Source: §4.3 · p. 7, Table 4 · p. 7

4.4 Chemical Interpretability Analysis (Q5)

  • In a BBBP example, KMPNN attends to electron-related attributes and more lively atomic neighbors, while coarse-grained attributes receive more attention than fine-grained ones.The authors connect electron affinity, electronegativity, metallicity, and ionization with electron loss and molecular polarity; they also note higher attention to Cl and hydroxyl neighbors in the example.Source: §4.4 · p. 7, Figure 4 · p. 7

5 Conclusion and Future Work

The paper concludes that Chemical Element KG and KCL incorporate domain knowledge into molecular graph representation learning, with experiments showing effectiveness under both fine-tune and linear protocols. Future work targets richer knowledge and a larger, improved KG.

5 Conclusion and Future Work

  • The authors report that KCL excels previous methods under both fine-tune and linear protocols, with better interpretation and representation capability.The framework uses Element KG to build microscopic connections between elements.Source: §5 · p. 7
  • Planned KG extensions include finer-grained domain knowledge, more OWL2 properties and axioms, and continued quality and scale improvements.The authors also plan to open-source the Chemical Element KG.Source: §5 · p. 7

A.1 Feature Initialization

  • Feature initialization uses random atom and bond vectors and RotatE-trained embeddings for attributes and relations.Atom and attribute dimensions are 128, while bond and relation dimensions are 64; the authors report that tested atom/bond dimension combinations had no significant performance impact.Source: §A.1 · p. 10

A.2 Architecture of KMPNN

KMPNN processes atom and attribute neighbors in augmented molecular graphs using heterogeneous message passing and attention. The original molecular graph is encoded with a GNN, for which the implementation uses GCN and pooled graph readout.

A.2 Architecture of KMPNN

  • KMPNN uses different message-passing parameters for neighbor types and attention parameters for atom–attribute versus atom–atom interactions.The architecture shares parameters among nodes corresponding to the same knowledge type and assigns different attention values to neighbor nodes.Source: §A.2 · p. 10, Figure 5 · p. 10

A.3 GNN-based Encoder

  • For the original molecular graph, the encoder aggregates neighboring messages and converts node representations into a graph embedding through a readout function.The implementation adopts GCN with weighted-sum and max pooling; the authors report that different GNN architectures have little impact on performance.Source: §A.3 · p. 10

B.1 Dataset Description

The benchmarks cover six binary classification datasets and two binary regression datasets, with scaffold and random splits. Baselines span supervised, pretrained, and graph-contrastive methods.

B.1 Dataset Description

  • The benchmark comprises 6 binary classification datasets and 2 binary regression datasets, evaluated with scaffold and random splits.Scaffold splitting separates molecules with distinct two-dimensional frameworks, making test molecules structurally different from training molecules.Source: §B.1 · p. 10
  • Classification datasets cover permeability, adverse drug reactions, clinical toxicity, BACE-1 inhibition, and toxicity labels.ToxCast labels come from high-throughput screening tests; Tox21 measures compound toxicity.Source: §B.1 · p. 10
  • The regression datasets are ESOL, documenting compound solubility, and FreeSolv, containing small-molecule hydration free energies.FreeSolv includes values from experiments and alchemical free-energy calculations.Source: §B.1 · p. 10

B.2 Baselines

  • Supervised baselines include GCN, Weave, MPNN, DMPNN, CMPNN, and CoMPT.Their approaches include node-neighbor convolution, atom-pair features, node-and-edge messages, and communicative message passing.Source: §B.2 · p. 11
  • Pretraining baselines include N-GRAM, Hu et al.'s node and graph tasks, and GROVER's contextual-property and motif prediction.These methods use vertex attributes, neighborhood context, node attributes, or contextual and motif information as pretraining signals.Source: §B.2 · p. 11
  • Graph-contrastive baselines include InfoGraph, MICRO-Graph, GraphCL, JOAO, and MoCL.They use node-graph mutual information, motif-based contrast, selected or automated augmentation, or domain knowledge.Source: §B.2 · p. 11

C Implementation and Pre-training Details

The authors implemented KCL with PyTorch and Deep Graph Library on a four-GPU server. Downstream fine-tuning used validation-set early stopping and random hyperparameter search.

C Implementation and Pre-training Details

D Downstream Details

  • Downstream fine-tuning used validation-set early stopping and random search over hyperparameter combinations.Fine-tuning tasks ran on a single GPU; Table 7 lists the downstream hyperparameters.Source: D Downstream Details · p. 11, Table 7 · p. 12

E.1 Regression Results under Linear Protocol

Contrastive learning is associated with average gains over KMPNN without contrastive learning, while ablations and encoder comparisons assess the framework's components. Regression performance is reported under both linear and fine-tune protocols.

E.1 Regression Results under Linear Protocol

  • Regression performance is reported under the linear protocol as well as the fine-tune protocol.The linear protocol is used for classification in the main content; Table 8 gives additional regression results.Source: §E.1 · p. 12, Table 8 · p. 12

E.2 Effect of Different Settings

  • Removing either knowledge feature initialization or negative sampling decreases performance; omitting both almost always performs worst.KCL with both components shows the best performance among the tested architectures.Source: §E.2 · p. 12, Table 9 · p. 12

E.3 Effect of Different Encoders

  • Changing the original-graph GNN encoder has no significant performance impact, while KMPNN performs better on augmented graphs than the compared heterogeneous GNN and general message-passing framework.Table 11 reports the encoder comparisons.Source: §E.3 · p. 12

E.4 Effect of Contrastive Learning

  • The cited comparison reports that contrastive learning produces average performance boosts of 0.7% on classification and 76.5% on regression over the model without contrastive learning.Source: §E.4 · p. 12

E.5 Chemical Interpretability Analysis

Attention visualizations show which element attributes and neighboring atoms receive attention in BBBP, BACE, and ToxCast examples. The reported patterns favor coarse-grained attributes over fine-grained ones.

E.5 Chemical Interpretability Analysis

  • In BBBP, atoms attend to electron affinity, electronegativity, metallicity, and ionization, attributes related to losing electrons.The passage connects electron gain or loss to molecular polarity and, in turn, permeability.Source: §E.5 · p. 12, Figure 6 · p. 12
  • In the BBBP visualization, more lively atomic neighbors receive greater attention.Source: §E.5 · p. 12, Figure 6 · p. 12
  • In BACE and ToxCast examples, coarse attributes receive more attention than fine-grained attributes such as weight and radius.The coarse attributes named include electronegativity, conductivity, density, ionization, metallicity, state, and family.Source: §E.5 · p. 13, Figure 7 · p. 13
Loading 2112.00544v2…