Neural Atoms: Propagating Long-range Interaction in Molecular Graphs through Efficient Communication Channel

TL;DR

This work tackles the challenge of capturing long-range interactions (LRI) in molecular graphs with standard GNNs, which typically focus on short-range interactions (SRI). It introduces Neural Atoms, a learnable, compact set of $K$ pseudo-atoms that summarize groups of real atoms, enabling explicit communication among neural atoms and a back-projection that enriches original atom embeddings, effectively turning distant interactions into a single-hop operation. The method is theoretically connected to Ewald Summation by separating SRI and LRI contributions and empirically validated on four long-range benchmarks, showing substantial gains on 2D intra-molecular tasks and competitive performance against 3D Ewald-based approaches without requiring 3D coordinates. Neural Atoms are shown to be plug-in, scalable, and compatible with diverse GNN backbones, offering a practical path to enhance molecular property prediction and related tasks by better modeling LRIs.

Abstract

Graph Neural Networks (GNNs) have been widely adopted for drug discovery with molecular graphs. Nevertheless, current GNNs mainly excel in leveraging short-range interactions (SRI) but struggle to capture long-range interactions (LRI), both of which are crucial for determining molecular properties. To tackle this issue, we propose a method to abstract the collective information of atomic groups into a few $\textit{Neural Atoms}$ by implicitly projecting the atoms of a molecular. Specifically, we explicitly exchange the information among neural atoms and project them back to the atoms' representations as an enhancement. With this mechanism, neural atoms establish the communication channels among distant nodes, effectively reducing the interaction scope of arbitrary node pairs into a single hop. To provide an inspection of our method from a physical perspective, we reveal its connection to the traditional LRI calculation method, Ewald Summation. The Neural Atom can enhance GNNs to capture LRI by approximating the potential LRI of the molecular. We conduct extensive experiments on four long-range graph benchmarks, covering graph-level and link-level tasks on molecular graphs. We achieve up to a 27.32% and 38.27% improvement in the 2D and 3D scenarios, respectively. Empirically, our method can be equipped with an arbitrary GNN to help capture LRI. Code and datasets are publicly available in https://github.com/tmlr-group/NeuralAtom.

Figures (38)

• Figure 1: An exemplar molecular with the long-range interactions (dash lines) and short-range interactions (solid lines).
• Figure 2: Illustration of Neural Atoms. The mapping function $f$ is to project the original atoms to Neural Atoms, and the retrieving function $f^{-1}$ aims to inject back the information, allowing the GNN to capture LRI via the interaction between Neural Atoms.
• Figure 3: The proposed Neural Atom framework aims to obtain graph representation for different downstream tasks. The Neural Atom can enhance arbitrary by injecting LRI information via the interaction of neural atoms. We demonstrate the information exchange by the mixture of colors.
• Figure 4: Neural Atom grouping strategies.
• Figure 7: Training and validation loss curves visualizations for DimeNet++ (1) the training loss curve, (2) the validation loss curve.

Theorems & Definitions (4)

• Definition 1
• Remark 1: Expressiveness
• Remark 2: Complexity and Scalability
• Remark 3