PaperlayerView exact revision on arXiv

Source-linked AI summary

Spherical Message Passing for 3D Graph Networks

Yi Liu, Limei Wang, Meng Liu, Xuan Zhang, Bora Oztekin, Shuiwang Ji

arXiv:2102.05013v5cs.LG

TL;DR

3D molecular learning lacks a principled message passing framework that captures spatial structure efficiently. The paper analyzes spherical-coordinate representations, proposes SMP and SphereNet, and reports strong performance, efficiency, and scalability across molecular datasets.

  • Problem

    3D molecular graphs require representation learning over spatial positions, but a principled message passing framework for this setting is lacking.

  • Method

    The paper analyzes spherical-coordinate geometry and proposes spherical message passing with physically meaningful representations, implemented in SphereNet.

  • Results

    SphereNet achieves the best performance against baselines without increasing the computing budget and reduces running time by 4 times versus a complete message passing scheme.

  • Takeaways & Limitations

    SMP and SphereNet provide approximately complete, invariant, efficient, and scalable representations for 3D molecular learning.

Abstract

from arXiv · show

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

1 INTRODUCTION

The paper develops a principled message passing framework for 3D molecular graphs by representing relative geometry in spherical coordinates. SMP provides approximately complete, efficient, and invariant representations, while SphereNet applies physically meaningful representations to molecular prediction tasks.

  • 3D molecular graphs lack a principled message passing framework that directly accounts for spatial information.
  • Distance, angle, and torsion uniquely determine relative atom locations, but complete identification using edge-based 2-hop information is computationally expensive.
  • SMP uses an edge-based 1-hop torsion strategy to reduce training complexity while retaining approximately complete 3D molecular representations.
  • SMP can distinguish almost all 3D graph structures, with uncovered cases appearing rarely in practice.
  • SphereNet maps SMP geometries to physically meaningful representations and achieves the best performance against baselines without increasing the computing budget.
  • Compared with a complete 2-hop message passing scheme, SphereNet achieves comparable performance while reducing running time by 4 times.

2 COMPLETE REPRESENTATIONS OF MOLECULES

The paper analyzes 3D molecular structure identification in spherical coordinates, showing why distance, angle, and torsion are needed. These geometries capture molecular structure and enable SMP to recognize properties such as chirality.

  • Equivariant methods commonly use Cartesian coordinates and require carefully designed equivariant layers, some of whose components are expensive.
  • In spherical coordinates, each point is specified by radial distance, angle, and torsion-related geometry relative to a reference frame.
  • The H2O2 structure is uniquely defined by three bond lengths, two bond angles, and one torsion angle.
  • Changing the torsion angle changes the H2O2 structure, demonstrating that torsion is necessary for distinguishing its 3D configuration.
  • SMP systematically combines distance, angle, and torsion and can recognize molecular chirality through angle and torsion information.

3 SPHERICAL MESSAGE PASSING

SMP performs message passing in spherical coordinates, using local bond lengths, bond angles, and torsion angles to encode relative 3D structure. Its 1-hop design reduces complexity while distinguishing almost all 3D graph structures and preserving translation and rotation invariance.

  • Message Passing Scheme: SMP performs message passing in the spherical coordinate system, incorporating 3D information into edge-message updates.It extends standard graph-network updates with a position-to-edge aggregation function.
  • Message Passing Scheme: For each message, SMP defines a local coordinate system using the sender atom as origin and the message direction as the z-axis.A neighboring atom supplies the reference plane for computing the relative geometry of other neighbors.
  • Message Passing Scheme: SMP computes torsion angles by projecting neighboring atoms onto a plane perpendicular to the message direction and ordering them in a predefined direction.The torsion angles around a center sum to 2π, and summation makes the arbitrary starting reference atom irrelevant.
  • Completeness Versus Efficiency: O(nk^2) is SMP’s time complexity, reduced from O(nk^3) for complete edge-based 2-hop identification through its edge-based 1-hop construction.This reduction supports application to large molecular graphs such as OC20.
  • Completeness Versus Efficiency: SMP distinguishes chirality and almost all 3D graph structures, while the remaining indistinguishable cases are described as unlikely to occur in nature.Its relative torsion design neglects absolute neighbor order, producing the stated exceptional case.
  • Related Methods: SchNet and DimeNet can be viewed as special cases of SMP because they capture only partial 3D position information.The related methods are characterized as instances of the broader spherical message-passing scheme.

4 SPHERENET

SphereNet combines SMP with physically based representations of spherical geometries for molecular learning. The resulting representations are designed to be accurate, physically meaningful, and invariant to translations and rotations.

  • Physical Representations: SphereNet represents SMP’s distance, angle, and torsion tuple as Ψ(d, θ, ϕ) for neural molecular learning.The paper motivates this transformation because raw geometric tuples lack meaningful representations for neural networks.
  • Physical Representations: The representation uses spherical Bessel functions for distance and spherical harmonics for angular information, with simplified forms Ψ(d) and Ψ(d, θ).The highest spherical-harmonic and spherical-Bessel orders are experimental hyperparameters.
  • SphereNet: SphereNet produces accurate, physically meaningful representations while generating predictions invariant to translation and rotation of input molecules.These properties follow from combining physical representations with spherical-coordinate position information.

5 RELATED WORK

Prior 3D molecular-graph methods either use Cartesian coordinates with costly equivariant machinery or use relative geometries such as distances and angles. Related 3D graph work also commonly focuses on local distance information.

  • Equivariant Methods: Equivariant molecular GNNs typically use Cartesian coordinates and require carefully designed equivariant layers.Some components, including spherical harmonics and Clebsch–Gordan coefficients, are computationally expensive.
  • Equivariant Methods: For molecular tasks requiring invariant outputs, complicated SE(3) representations may be unnecessary because S2 representations can suffice.The cited related work also reports unsatisfactory performance on molecular tasks.
  • Relative-Geometry Methods: Relative-geometry methods are naturally invariant and use inputs such as interatomic distances, bond angles, and angles between planes.This research direction is described as being in an early stage and focused on selecting useful geometries.
  • Relative-Geometry Methods: SchNet and PhysNet incorporate distance information, while DimeNet encodes directional information as part of the relative-geometry approach.These methods differ in which geometric quantities they use during message aggregation or representation learning.
  • Other 3D Graph Data: In 3D point-cloud graphs, existing methods mainly capture distance information from local neighborhoods in 3D space.Points serve as nodes and connections between them are directed edges.

6 EXPERIMENTAL STUDIES

SphereNet is evaluated against diverse baselines on OC20, QM9, and MD17, with results emphasizing predictive performance and the contributions of spherical message-passing geometries.

  • Experimental setup: Experiments compare SphereNet with graph, geometric, and equivariant baselines across OC20, QM9, and MD17.The baseline set includes methods such as SchNet, DimeNet++, GemNet, PaiNN, NequIP, and sGDML.
  • OC20: SphereNet achieves the best OC20 energy MAE on three of four splits and in the average, while achieving the best EwT on all four splits.It reduces average energy MAE by 0.019, corresponding to 3.10% of the second-best model.
  • QM9: SphereNet improves QM9 overall mean std. MAE from 0.98 to 0.91 and achieves the best performance on five properties.It also achieves the second-best performance on three additional properties.
  • MD17: On MD17, SphereNet is evaluated for atomic-force prediction using separate models for each molecule and 1000 training samples.Baseline models use a joint force-and-energy loss, with comparisons controlled through the force-over-energy weight.
  • MD17: SphereNet achieves much better overall std. MAE than sGDML and similar results to GemNet when using WoFE=1000.This comparison is notable because GemNet-T has high complexity and carefully designed network components for performance improvement.
  • Ablation study: The ablation study shows SMP outperforming SMP w/o ϕ, which in turn outperforms SMP w/o (θ, ϕ), supporting the contributions of torsion and angle information.All three strategies use the same architecture with other network components unchanged.

7 CONCLUSIONS

The paper presents SMP as an efficient, approximately complete message-passing scheme for 3D molecules and builds SphereNet from meaningful physical representations. Experiments across multiple datasets support its capability, efficiency, and scalability.

  • SMP provides an efficient scheme for approximately complete molecular representations without increasing the computing budget.
  • SphereNet combines SMP with meaningful physical representations for 3D molecular learning.
  • The edge-update function φe uses edge, receiver-node, sender-node, neighboring-edge, and connected-node information.
  • The node-update function φv aggregates incoming edge features and related-node information, while φu updates the global graph feature.
  • SMP updates edge messages using aggregated 3D position representations rather than absolute Cartesian coordinates.

B INFORMATION UPDATE AND ARCHITECTURE OF SPHERENET

SphereNet uses an input block, repeated interaction blocks, and output blocks to update edge messages and node features. Its interaction process incorporates distance, angle, and torsion representations.

  • Input block: The input block initializes an edge message using only the distance representation Ψ(d).
  • Interaction blocks: Each interaction block updates a center edge message using neighboring-edge messages and Ψ(d, θ, ϕ), Ψ(d, θ), and Ψ(d).
  • Architecture: SphereNet consists of an input block followed by multiple interaction blocks and output blocks.
  • Output block: The output block aggregates gated incoming messages to update the receiver-node feature.

C RELATIONS WITH PRIOR MESSAGE PASSING METHODS

Prior message-passing methods can be expressed as special cases of SMP because they encode only partial 3D information. SchNet and PhysNet use distances, whereas DimeNet adds directional angle information.

  • Existing message-passing schemes can be viewed as SMP special cases that encode partial 3D information.
  • SchNet: SchNet represents pairwise position information with a radial-basis distance embedding and uses neural-network transformations with element-wise multiplication.
  • PhysNet: PhysNet uses a smoothly cut-off radial-basis distance representation and neural-network updates modulated by element-wise multiplication.
  • DimeNet: DimeNet explicitly represents distances and directions through concatenated Ψ(d) and Ψ(d, θ) features.
  • DimeNet: DimeNet does not require the ρp→v, ρp→u, or ρe→u aggregation functions.

D EXPERIMENTAL SETUP

The experiments use common embedding dimensions and learning-rate scheduling across OC20, QM9, and MD17, with selected hyperparameters tuned by grid search.

  • All models use input embedding size 256 and output embedding size 64 for both LB2 and LB blocks.
  • Training begins with learning-rate warmup and uses either ReduceLROnPlateau or StepLR scheduling.
  • The experiments omit weight decay and dropout, while some hyperparameters are fixed and others are selected by grid search.

E OC20 DATA DESCRIPTION

The paper describes the IS2RE evaluation setting and compares SphereNet’s efficiency with other models on QM9.

  • OC20 DATA DESCRIPTION: IS2RE predicts relaxed-state structure energies, a catalysis-relevant task, using a dataset with 460,318 training structures.Performance is evaluated on a validation set because test labels are unavailable.
  • Experimental settings: The supplied tables list hyperparameter search spaces for OC20, QM9, and MD17.These are provided as Tables 6–8.
  • Efficiency comparison: SphereNet uses computational resources similar to DimeNet++ and GemNet-T on QM9 property U0.The comparison uses parameter counts and time cost per epoch on the same infrastructure.
  • Efficiency comparison: SphereNet is much more efficient than DimeNet, while GemNet-Q cannot run on QM9 using the stated infrastructure.The paper attributes SphereNet’s efficiency in part to its efficient torsion-computation method.

G SPHERENET FILTER VISUALIZATION

SphereNet filters are visualized across sampled distance, angle, and torsion locations to examine how 3D information shapes learned representations.

  • Geometric dependence: The visualization shows that changing any one of distance, angle, or torsion alters filter structure while the other two are fixed.This indicates that all three geometric variables contribute to the filters’ structural semantics.
  • Geometric dependence: SphereNet filters learn different 3D information through SMP, supporting improved representations.The interpretation follows the paper’s stated conclusion from the filter visualization.
  • Filter sampling: Figure 6 visualizes six SphereNet filters, with each row corresponding to one filter and torsion sampled at eight angles spaced by π/4.The torsion angles range from 0 to 7π/4.
Loading 2102.05013v5…