Graph Neural Network-Based\nMolecular Property Prediction\nwith Patch Aggregation

Graph neural networks (GNNs) have emerged as powerful\ntools for\nquantum chemical property prediction, leveraging the inherent graph\nstructure of molecular systems. GNNs depend on an edge-to-node aggregation\nmechanism for combining edge representations into node representations.\nUnfortunately, existing learnable edge-to-node aggregation methods\nsubstantially increase the number of parameters and, thus, the computational\ncost relative to simple sum aggregation. Worse, as we report here,\nthey often fail to improve predictive accuracy. We therefore propose\na novel learnable edge-to-node aggregation mechanism that aims to\nimprove the accuracy and parameter efficiency of GNNs in predicting\nmolecular properties. The new mechanism, called “patch aggregation”,\nis inspired by the Multi-Head Attention and Mixture of Experts machine\nlearning techniques. We have incorporated the patch aggregation method\ninto the specialized, state-of-the-art GNN models SchNet, DimeNet++,\nSphereNet, TensorNet, and VisNet and show that patch aggregation consistently\noutperforms existing learnable and nonlearnable aggregation techniques\n(sum, multilayer perceptron, softmax, and set transformer aggregation)\nin the prediction of molecular properties such as QM9 thermodynamic\nproperties and MD17 molecular dynamics trajectory energies and forces.\nWe also find that patch aggregation not only improves prediction accuracy\nbut also is parameter-efficient, making it an attractive option for\npractical applications for which computational resources are limited.\nFurther, we show that Patch aggregation can be applied across different\nGNN models. Overall, Patch aggregation is a powerful edge-to-node\naggregation mechanism that improves the accuracy of molecular property\npredictions by GNNs.