Spherical Message Passing for 3D Graph Networks

TL;DR

3D molecular graph learning faces invariance and scalability challenges. This work analyzes complete geometric representations in the spherical coordinate system and introduces spherical message passing (SMP), an efficient 1-hop scheme that achieves approximately complete 3D learning with complexity $O(nk^2)$, versus full-edge methods at $O(nk^3)$. Building on physically meaningful representations, SphereNet encodes $(d,\theta,\varphi)$ via $\Psi(d,\theta,\varphi)$ using spherical Bessel and spherical harmonics and reports state-of-the-art results on OC20, QM9, and MD17 while maintaining similar compute budgets. The approach demonstrates strong practical impact for large molecular systems, and code is publicly available in the DIG library.

Abstract

We consider representation learning of 3D molecular graphs in which each atom is associated with a spatial position in 3D. This is an under-explored area of research, and a principled message passing framework is currently lacking. In this work, we conduct analyses in the spherical coordinate system (SCS) for the complete identification of 3D graph structures. Based on such observations, we propose the spherical message passing (SMP) as a novel and powerful scheme for 3D molecular learning. SMP dramatically reduces training complexity, enabling it to perform efficiently on large-scale molecules. In addition, SMP is capable of distinguishing almost all molecular structures, and the uncovered cases may not exist in practice. Based on meaningful physically-based representations of 3D information, we further propose the SphereNet for 3D molecular learning. Experimental results demonstrate that the use of meaningful 3D information in SphereNet leads to significant performance improvements in prediction tasks. Our results also demonstrate the advantages of SphereNet in terms of capability, efficiency, and scalability. Our code is publicly available as part of the DIG library (https://github.com/divelab/DIG).

Figures (6)

• Figure 1: The chemical structure of the $\hbox{H}_2\hbox{O}_2$.
• Figure 2: (a). The message aggregation scheme for the spherical message passing. (b). An illustration for computing torsion angles in the spherical message passing architecture.
• Figure 3: An illustration of cases that SMP can and cannot distinguish. All the neighboring nodes of $s_k$ are projected to the plate perpendicular to the message of interest. We assume all the distances and angles are fixed (the molecules can be more easily distinguished otherwise). Hence, all the angle shown are torsion angles and they are formed in the anticlockwise direction. (a) and (b) are chiral and SMP can distinguish them. This is because in (a), $q^\prime_1(90^\circ)$, $q^\prime_2(60^\circ)$, $q^\prime_3(120^\circ)$, $q^\prime_4(90^\circ)$; in (b), $q^\prime_1(60^\circ)$, $q^\prime_2(120^\circ)$, $q^\prime_3(90^\circ)$, $q^\prime_4(90^\circ)$. SMP cannot distinguish (b) and (c) but this scenario may not exist in nature. $\angle q^\prime_1s_kq^\prime_2$ in (b) and $\angle q^\prime_1s_kq^\prime_3$ in (c) usually are different as $q^\prime_2$ and $q^\prime_3$ are different atoms and the corresponding distances and angles are the same.
• Figure 4: Architecture of SphereNet. LB2 denotes a linear block with two linear layers, $\sigma$(LB) denotes a linear layer followed by an activation function, $\|$ denotes concatenation, and $\odot$ denotes element-wise multiplication. Each LB2 aims at canceling bottlenecks by performing downprojection, followed by upprojection. Hence, it is related to three hyperparameters; these are, input embedding size, intermediate size, and output embedding size. Each linear block LB is related to hyperparameters including input embedding size and output embedding size. Description of each block is in Sec. \ref{['sec:supp_B']}.
• Figure 5: Illustrations of the functions $\phi^e$ (a) and $\phi^v$ (b).