[Paper Review] Spherical Message Passing for 3D Graph Networks

Yi Liu, Limei Wang+4 more|arXiv (Cornell University)|Feb 9, 2021

Advanced Graph Neural Networks47 references31 citations

TL;DR

Introduces spherical message passing (SMP) for 3D molecular graphs to efficiently use distance, angle, and torsion information, resulting in SphereNet which achieves state-of-the-art or competitive performance with better efficiency across several molecular benchmarks.

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).

Motivation & Objective

• Motivate and formalize complete identification of 3D molecular structures using relative 3D information in spherical coordinates.
• Develop SMP to enable efficient, scalable message passing with 1-hop geometry information.
• Introduce SphereNet to encode 3D geometries into physically meaningful representations and evaluate on standard molecular benchmarks.
• Assess completeness vs. efficiency trade-offs of SMP compared to complete 2-hop schemes.

Proposed method

• Define 3D molecular graphs in spherical coordinate system (SCS) and show completeness requires distance, angle, and torsion.
• Propose SMP where edge, node, and graph update functions (φ^e, φ^v, φ^u) integrate 3D information via a locality-driven ρ^{p→e} and ρ^{e→v} aggregations.
• Compute torsion angles using edge-based 1-hop information with a reference-plane construction to keep complexity O(nk^2).
• Argue SMP is approximately complete for distinguishing most 3D structures and unrolls into SphereNet by applying physically meaningful basis functions.
• Define SphereNet representations Ψ(d,θ,φ) via spherical Bessel and spherical harmonics to ensure translation/rotation invariance and integrate into the message passing blocks.
• Relate SMP to prior methods (SchNet, DimeNet) as special cases within the SMP framework.

Experimental results

Research questions

• RQ1 Can 3D molecular graphs be effectively identified using a complete set of relative geometric features (distance, angle, torsion) within a scalable GNN framework?
• RQ2 Does spherical message passing provide a principled, efficient alternative to edge-based 2-hop complete representations while preserving expressive power?
• RQ3 Do SphereNet representations Ψ(d,θ,φ) yield improved predictive performance across standard molecular benchmarks compared to existing baselines?
• RQ4 How does SMP/SphereNet balance accuracy, efficiency, and scalability on large-scale molecular datasets?

Key findings

• SphereNet achieves competitive or best performance on OC20 IS2RE in energy MAE and EwT across splits, with notable average improvements over baselines.
• On QM9, SphereNet sets a new state-of-the-art in std. MAE, achieving strong results across multiple properties.
• MD17 experiments show SphereNet often outperforms SchNet and DimeNet in force MAE under comparable WoFE settings, and approaches GemNet-T when using higher WoFE.
• Ablation and efficiency analyses show SMP matches complete 2-hop representations in accuracy while offering substantially lower running time (≈4x faster than complete schemes in some settings).
• SphereNet reduces training complexity to O(nk^2) while maintaining near-complete distinguishability of 3D structures, enabling large-scale molecular learning.