Chi-Geometry: A Library for Benchmarking Chirality Prediction of GNNs

TL;DR

Chi-Geometry is a library for generating synthetic graph datasets to benchmark GNNs' ability to predict molecular chirality, enabling more interpretable assessments and guiding the development of improved GNN architectures.

Contribution

The paper introduces Chi-Geometry, a novel library for creating controlled synthetic datasets for benchmarking GNNs on chirality prediction, and demonstrates its use in designing two new GNN architectures.

Findings

All-to-all connected GNN achieves high accuracy but with quadratic computational cost.

Virtual node GNN maintains linear scaling with competitive accuracy.

Benchmarking guided the design of architectures balancing accuracy and efficiency.

Abstract

We introduce Chi-Geometry - a library that generates graph data for testing and benchmarking GNNs' ability to predict chirality. Chi-Geometry generates synthetic graph samples with (i) user-specified geometric and topological traits to isolate certain types of samples and (ii) randomized node positions and species to minimize extraneous correlations. Each generated graph contains exactly one chiral center labeled either R or S, while all other nodes are labeled N/A (non-chiral). The generated samples are then combined into a cohesive dataset that can be used to assess a GNN's ability to predict chirality as a node classification task. Chi-Geometry allows more interpretable and less confounding benchmarking of GNNs for prediction of chirality in the graph samples which can guide the design of new GNN architectures with improved predictive performance. We illustrate Chi-Geometry's efficacy by using it to generate synthetic datasets for benchmarking various state-of-the-art (SOTA) GNN architectures. The conclusions of these benchmarking results guided our design of two new GNN architectures. The first GNN architecture established all-to-all connections in the graph to accurately predict chirality across all challenging configurations where previously tested SOTA models failed, but at a computational cost (both for training and inference) that grows quadratically with the number of graph nodes. The second GNN architecture avoids all-to-all connections by introducing a virtual node in the original graph structure of the data, which restores the linear scaling of training and inference computational cost with respect to the number of nodes in the graph, while still ensuring competitive accuracy in detecting chirality with respect to SOTA GNN architectures.