[Paper Review] ComENet: Towards Complete and Efficient Message Passing for 3D Molecular Graphs

Limei Wang, Yi Liu+3 more|arXiv (Cornell University)|Jun 17, 2022

Machine Learning in Materials Science27 citations

TL;DR

ComENet introduces a complete and efficient 1-hop neighborhood message passing scheme for 3D molecular graphs, achieving SE(3)-complete representations with O(nk) complexity and strong empirical performance on large datasets.

ABSTRACT

Many real-world data can be modeled as 3D graphs, but learning representations that incorporates 3D information completely and efficiently is challenging. Existing methods either use partial 3D information, or suffer from excessive computational cost. To incorporate 3D information completely and efficiently, we propose a novel message passing scheme that operates within 1-hop neighborhood. Our method guarantees full completeness of 3D information on 3D graphs by achieving global and local completeness. Notably, we propose the important rotation angles to fulfill global completeness. Additionally, we show that our method is orders of magnitude faster than prior methods. We provide rigorous proof of completeness and analysis of time complexity for our methods. As molecules are in essence quantum systems, we build the \underline{com}plete and \underline{e}fficient graph neural network (ComENet) by combing quantum inspired basis functions and the proposed message passing scheme. Experimental results demonstrate the capability and efficiency of ComENet, especially on real-world datasets that are large in both numbers and sizes of graphs. Our code is publicly available as part of the DIG library (\url{https://github.com/divelab/DIG}).

Motivation & Objective

• Motivate the need for fully incorporating 3D information in molecular graphs for accurate learning.
• Define and guarantee geometric completeness (global and local) for 3D graph representations.
• Develop an efficient 1-hop message passing scheme with rotation angles to achieve global completeness.
• Integrate quantum-inspired features within ComENet for improved 3D molecular learning.
• Demonstrate scalability and competitive accuracy on large-scale datasets (OC20, Molecule3D, QM9).

Proposed method

• Formally define completeness for 3D graphs under SE(3) and show global completeness via rotation angles on edges.
• Compute 1-hop geometric features (d, theta, phi, tau) per edge to form SE(3)-invariant representations.
• Introduce local and global convolution layers that operate within 1-hop neighborhoods to achieve O(nk) complexity.
• Use rotation angles tau (based on nearest neighbors) to uniquely determine 3D conformations and enable conformer identification.
• Incorporate quantum-inspired basis features and a neural network-based message passing framework within ComENet architecture (embedding layer, interaction layers, self-atom layer, pooling).
• Provide proofs of completeness and analyze time complexity; evaluate on OC20, Molecule3D, and QM9.

Experimental results

Research questions

• RQ1 Can a 3D molecular graph representation be globally complete (no information loss) using only 1-hop neighborhood information?
• RQ2 How can rotation angles be leveraged to achieve global completeness and enable conformer-level discrimination?
• RQ3 What is the practical time complexity of a complete 3D GNN compared to existing 2-hop methods?
• RQ4 Do quantum-inspired features combined with the proposed message passing yield competitive performance on large-scale molecular datasets?

Key findings

• Proves that a 1-hop complete geometric transformation (d, theta, phi, tau) yields global completeness for strongly connected 3D graphs.
• Achieves O(nk) time complexity by using 1-hop neighborhood and rotation angles, significantly faster than O(nk^2) or O(nk^3) methods.
• Demonstrates competitive or superior performance to state-of-the-art baselines on OC20, Molecule3D, and QM9, with substantial training/inference speedups (6-10x on some datasets).
• Shows rotation angles enable conformer-level discrimination in OC20 and identify important conformational differences.
• Proposes ComENet architecture integrating complete message passing with quantum-inspired features for scalable 3D molecular learning.