Paper Summary

• Introduces a new model, CMPNN, for Communicative Representation Learning on Attributed Molecular Graphs
• CMPNN improves molecular embedding by:
  • Following the edge-based message passing in DMPNN
  • Introducing node-edge message communication modules such as Inner Product Kernel, Gated Graph Kernel, and Multilayer Perception
  • Updating both bond and atom embeddings during training
  • Including a message booster to enrich the message generation process

Motivation

• Adding an alternative to DMPNN (Chemprop uses this GNN architecture) that improves DMPNN by adding:
  • Node-edge message communication modules
  • Message booster
  • Updating bond embeddings
• Adding a GNN model that updates edge embeddings during training

Benchmark Results

I have benchmarked CMPNN on a subset of datasets from the TDC ADMET Benchmark Group. I have used LitGNN to perform this benchmark and results can be found in a W&B report.

Table: Prediction results of CMPNN on five chemical graph datasets. The datasets were used from the TDC ADMET Benchmark group that provides train_val/test scaffold splits. The model was trained and tested for each task for five times, and reported the mean and standard deviation of AUROC or MAE values. *Chemprop-RDKIT utilizes a hybrid approach where it combines the learned molecule embeddings with 200 global molecule features (descriptors).

Implementation Details

Below are the places where improvements have been made:

• Node-edge message communication modules
  • In the paper, the authors mention different communicators such as Inner Product Kernel, Gated Graph Kernel and Multilayer Perception. However, within their code, they don't use such different communicators. Instead, they only use an additive communicator (not mentioned in the paper; I came up with the name 'additive').
  • I have implemented 4 communicators proposed in the paper and their code:
    • Additive
    • Inner product
    • GRU
    • MLP
  • The user can choose different communicator according to their dataset.
  • These communication modules are applied during the convolution layers.
• Final communication module
  • In the paper, the authors mention applying the same communication module to the message from incoming bonds [m(v)], current atom's representation [hK(v)] and atom's initial representation [x(v)].
  • In the section 3.3 of the paper where they describe the different node-edge message communication modules, all the communication modules operate on m(v) and hK-1(v).
  • In their code, they use MLP communication module.
  • I have kept this part the same and hardcoded the MLP communication module.

Checklist

• Add torch_geometric/nn/models/cmpnn.py
• Add CMPNN to torch_geometric/nn/models/__init__.py
• Add test/nn/models/test_cmpnn.py
• Add message_booster mode to the torch_geometric/nn/aggr/multi.py:MultiAggregation class
• Add a test for message_booster mode in test/nn/aggr/test_multi.py
• Add examples/cmpnn.py
• Add CMPNN to the 'Implemented GNN Models' section of the README
• Update CHANGELOG

Thank you! Please let me know if any changes are required.