Molecular Property Prediction with Graph Neural Networks

Overview

This project explores graph neural networks for molecular property prediction using the QM9 dataset.

Molecules are represented as graphs, where atoms are nodes and chemical bonds define the graph connectivity. The goal is to predict the molecular internal energy at 0 K, U0, using PyTorch Geometric.

Two graph neural network architectures are compared:

• A GCN using atom features and graph connectivity
• A bond-aware GNN using NNConv and explicit bond features

The experiments also investigate the effect of target normalization and graph-level pooling on model performance.

---

Dataset

The project uses the QM9 dataset provided through PyTorch Geometric.

QM9 contains more than 130,000 small organic molecules with quantum-chemical properties calculated using density functional theory.

In this project:

• Target property: U0 energy
• Target index in PyTorch Geometric QM9: 7
• Training molecules: 10,000
• Validation molecules: 2,000
• Test molecules: 2,000

The remaining molecules are not used in the current experiments.

---

Molecular Graph Representation

Each molecule is represented as a graph:

• Nodes: atoms
• Node features: 11 atomic features provided by QM9
• Edges: chemical bonds
• Edge features: 4 bond features provided by QM9

The GCN model uses node features and graph connectivity.

The bond-aware model additionally uses edge attributes during message passing.

---

Model 1: GCN

The GCN architecture consists of:

• 2 GCNConv layers
• ReLU activations
• Global sum pooling
• Fully connected layer: 64 -> 128
• Output layer: 128 -> 1

Global sum pooling is used because molecular energy is an extensive property and generally scales with molecular size.

---

Model 2: Bond-Aware GNN

The second model uses NNConv, which allows bond features to influence message passing.

The architecture consists of:

• Node feature encoder
• 2 NNConv layers
• Edge-conditioned neural networks
• Global sum pooling
• Fully connected regression head

This model explicitly incorporates bond information through edge_attr.

---

Target Normalization

Direct prediction of U0 resulted in poor optimization because the target values have a large numerical scale.

Normalization statistics are calculated using only the training split:

Target mean: -11185.296875
Target std:   1081.682007

During training, the target is normalized as:

normalized_target = (target - target_mean) / target_std

Predictions are transformed back to the original scale before validation and test metrics are calculated.

---

Initial Experiment

The first GCN was trained directly on the original U0 values and used global mean pooling.

Its test performance was:

Test MSE: 746153.324
Test MAE: 653.703

Two changes were then introduced:

• Target normalization
• Global sum pooling instead of global mean pooling

These changes significantly improved model performance.

---

Results

Model	Test MSE	Test MAE
Initial GCN	746153.32	653.70
Improved GCN	243082.96	238.61
Bond-aware GNN	299921.05	350.40

The improved GCN achieved the best performance among the tested models.

The test MAE decreased from approximately 653.7 to 238.6, corresponding to a reduction of about 63.5% compared with the initial setup.

The bond-aware NNConv model did not outperform the improved GCN in the current experimental setup.

This result shows that additional model complexity does not automatically lead to better generalization.

---

GCN Training Curve

---

GCN Predictions

The following plot compares predicted and reference U0 values on the test set.

The dashed diagonal line represents perfect predictions.

---

Bond-Aware GNN Training Curve

---

Bond-Aware GNN Predictions

---

Project Structure

molecular-property-gnn/
├── src/
│   ├── train.py
│   ├── train_bond_gnn.py
│   └── evaluate.py
│
├── results/
│   ├── loss_curve.png
│   └── predicted_vs_true.png
│
├── results_bond/
│   ├── loss_curve.png
│   └── predicted_vs_true.png
│
├── README.md
├── requirements.txt
└── .gitignore

The QM9 dataset, model checkpoints, logs, and virtual environment are excluded from the repository.

---

Installation

Create a virtual environment:

python3 -m venv venv

Activate it:

source venv/bin/activate

Install the required packages:

pip install -r requirements.txt

For GPU acceleration, install a PyTorch build compatible with the local CUDA environment.

---

Training the GCN

Run:

python src/train.py --epochs 20 --batch_size 128

The best model checkpoint is saved during training.

To run the training process in the background:

nohup python -u src/train.py --epochs 20 --batch_size 128 > train.log 2>&1 &

---

Training the Bond-Aware GNN

Run:

python src/train_bond_gnn.py --epochs 30 --batch_size 128

To run it in the background:

nohup python -u src/train_bond_gnn.py --epochs 30 --batch_size 128 > bond_gnn.log 2>&1 &

---

Evaluation

To evaluate the trained GCN:

python src/evaluate.py

The script reports test MSE and MAE and generates:

results/predicted_vs_true.png

---

Technologies

• Python
• PyTorch
• PyTorch Geometric
• NumPy
• Matplotlib
• CUDA

---

Key Findings

This project demonstrates several practical aspects of molecular machine learning:

• Molecular structures can be represented naturally as graphs
• Graph neural networks can be used for quantum-chemical property prediction
• Target normalization can substantially improve optimization
• Pooling strategies should reflect the physical nature of the target property
• More complex architectures do not necessarily produce better results
• Model comparison and error analysis are important parts of the machine learning workflow

---

Future Work

Possible extensions include:

• Training on a larger fraction of QM9
• Comparing GCN with GIN, GAT, SchNet, or PaiNN
• Explicitly using 3D atomic coordinates
• Predicting additional QM9 properties
• Multi-target molecular property prediction
• Hyperparameter optimization
• Studying model performance as a function of training-set size