PaperlayerView exact revision on arXiv

Source-linked AI summary

ComENet: Towards Complete and Efficient Message Passing for 3D Molecular Graphs

Limei Wang, Yi Liu, Yuchao Lin, Haoran Liu, Shuiwang Ji

arXiv:2206.08515v3cs.LG

TL;DR

Existing 3D graph methods can incorporate only partial geometric information or incur high computational cost. ComENet combines rotation-angle-based global completeness with efficient 1-hop message passing. The paper reports comparable or superior performance with training and inference accelerated by 6–10 times.

  • Problem

    Existing methods either incorporate partial 3D information or have high time complexity, limiting complete and efficient learning on 3D graphs.

    Source: §1 · pp. 1–2

  • Method

    ComENet combines rotation angles for global completeness with locally complete 1-hop message passing, representing each edge with (d, θ, φ, τ).

    Source: §1 · p. 2, §2.5 · p. 5

  • Results

    ComENet achieves similar or better performance than existing methods while accelerating training and inference by 6–10 times.

    Source: §1 · p. 2, §7 · p. 10

  • Takeaways & Limitations

    The authors conclude that ComENet’s completeness and efficiency support its scalability to real-world molecule datasets large in both graph count and graph size.

    Source: §1 · p. 2, §7 · p. 10

    The completeness proof assumes a strongly connected 3D graph; the authors state that molecules in nature satisfy this condition.

    Source: §3.1 · p. 5

Abstract

from arXiv · show

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

1 Introduction

The introduction argues that 3D geometry is needed to represent molecular structures, while existing methods can miss global structure or incur high computational cost. ComENet combines rotation-angle-based global completeness with efficient message passing and is reported to perform comparably to strong methods while accelerating training and inference.

1 Introduction

  • Molecular graph learning needs 3D information because graph connections alone omit geometries such as bond lengths, while existing methods may miss global structure or scale poorly.The introduction identifies SphereNet as locally complete but unable to distinguish a wide range of structures such as conformers.Source: §1 · pp. 1–2
  • ComENet uses rotation angles for global completeness and 1-hop message passing for local completeness, with complexity O(nk).The paper reports similar or better performance and 6–10× faster training and inference than existing methods.Source: §1 · p. 2
  • Compared with SphereNet’s 2-hop message passing and O(nk2) complexity, ComENet uses 1-hop neighborhoods and rotation angles to achieve provable global completeness at O(nk).Both methods use (d, θ, φ) for local completeness; ComENet adds rotation angles to distinguish global structures such as conformers.Source: §1 · p. 2

2 The Proposed Message Passing Scheme

The proposed scheme defines completeness as distinguishing 3D graphs up to rigid transformations, then combines local geometric features with rotation angles. Its 1-hop construction is designed to encode local and global structure efficiently.

2.1 Notations & Definitions

  • A geometric transformation is complete when identical representations correspond to graphs whose coordinates differ only by an SE(3) rotation and translation.Thus, the transformation distinguishes graphs with different 3D conformations.Source: §2.1 · p. 3

2.2 Global Completeness via Rotation Angles

  • For global completeness, each shared edge between local structures contributes a rotation angle, allowing representations to extend from n-hop to (n + 1)-hop neighborhoods.The angle is measured between planes formed using the edge endpoints and their selected neighboring reference nodes.Source: §2.2 · pp. 3–4

2.3 Rotation Angles for Conformer Identification

  • Because conformers share a 2D molecular graph but differ in 3D structure, incorporating bond rotation angles supports completeness at the conformer level.The paper illustrates this with butane, whose conformers have varying conformation energies.Source: §2.3 · p. 4

2.4 Local Completeness with Improved Efficiency

  • Local completeness is computed within 1-hop neighborhoods at O(nk), using a local coordinate system whose z-axis points to the nearest neighbor and whose xz-plane uses the second nearest.The method computes (d, θ, φ) in that coordinate system, unlike SphereNet’s 2-hop approach.Source: §2.4 · pp. 4–5

2.5 Message Passing Scheme

  • The message-passing scheme represents each edge with (d, θ, φ, τ), where the first three features specify local geometry and τ determines its orientation.The resulting representation is relative and SE(3)-invariant; the scheme computes the required geometries within 1-hop neighborhoods at O(nk).Source: §2.5 · p. 5

3 Merits of Our Methods

The paper proves geometric completeness for strongly connected 3D graphs and argues that this condition holds for molecules. It also analyzes the efficiency gains of 1-hop message passing and reduced torsion-angle computation.

3.1 Geometric Completeness

  • For strongly connected 3D graphs, the four-feature transformation is complete because each added node’s relative position can be uniquely determined.The proof proceeds by induction, and the authors state that molecules in nature can be constructed as strongly connected graphs.Source: §3.1 · pp. 5–6

3.2 Efficiency

  • 1-hop message passing has O(nk) complexity versus O(nk2) for 2-hop DMP, with a slight reduction in expressiveness because each layer adds only one hop of context.For torsion angles, the method uses n1 + n2 + 1 for two local regions rather than Ganea et al.’s n1 × n2, reducing the count from O(nk3) to O(nk).Source: §3.2 · p. 6

4 ComENet

ComENet follows an invariant 3D GNN architecture with an embedding layer, interaction layers, a self-atom layer, and pooling. Its interaction layers use geometric features and local and global graph convolutions.

4 ComENet

  • Figure 3 shows the ComENet overview and interaction layer; TBF and SBF encode tuples (d, θ, φ) and (d, τ), respectively.The diagram also labels concatenation, down-projection, and element-wise summation.Source: Figure 3 · p. 6
  • ComENet updates node features using neighboring-node and geometric features (d, θ, φ, τ) through local and global graph convolution layers.The interaction layer follows the proposed message passing scheme.Source: §4 · pp. 6–7
  • The network embeds atom types, applies multiple interaction layers, projects each node feature to one dimension, and sum-pools node features for predictions.Source: §4 · pp. 6–7

5 Related Work

Related 3D GNNs differ in how they represent geometry and balance expressivity with computational cost. ComENet is framed as using relative information to target completeness at lower cost.

5 Related Work

  • Equivariant 3D GNNs use coordinates directly and can be efficient, but require rotation-equivariant components and often rely on restrictive type-1 spherical-harmonic bases.Higher-order bases can make convolution kernels more expressive, but introduce efficiency concerns; these methods are reported to perform worse than invariant 3D GNNs.Source: §5 · p. 7
  • SchNet uses distances, DimeNet adds bond angles, SphereNet adds torsions at O(nk2) complexity, and GemNet's node quadruplets are more expensive.The paper characterizes the first three approaches as capturing incomplete or approximate 3D information, with SphereNet described as approximately complete.Source: §5 · p. 7
  • ComENet uses relative information only and targets complete 3D representations with lower computational cost, described as comparable to that of equivariant 3D GNNs.Source: §5 · p. 7

6 Experiments

Across OC20, Molecule3D, and QM9, ComENet combines competitive or better predictive performance with lower training cost. Removing its rotation angles harms OC20 performance.

6 Experiments

  • Experiments cover OC20, Molecule3D, and QM9; Molecule3D has about 4 million graphs, and OC20 averages 77.75 nodes per graph.Source: §6 · p. 7, Table 1 · p. 7

6.1 OC20

  • 0.0135 lower average energy MAE (2.2% of the second-best model) accompanies ComENet’s better OC20 performance and at least 10 times lower training time than existing best methods.ComENet also achieves the best results on the OOD Both split for both energy MAE and EwT.Source: Table 2 · p. 8, §6.1 · p. 8

6.2 Molecule3D

  • ComENet reduces Molecule3D training time by 6–9 times versus DimeNet++ and SphereNet, with similar performance to SchNet at similar cost but slightly worse performance than DimeNet++ and SphereNet.Molecule3D contains about 4 million molecular graphs.Source: §6 · p. 7, Table 3 · p. 8, §6.2 · p. 9

6.3 QM9

  • On QM9, ComENet outperforms 1-hop SchNet, PhyNet, and MGCN, matches DimeNet, DimeNet++, PaiNN, and SphereNet across properties and overall std. MAE, and is more efficient than 2-hop methods.Source: Table 4 · p. 9, §6.3 · p. 9

6.4 Ablation Study for Identifying Conformers

  • Removing rotation angle τ from ComENet harms OC20 performance, supporting the contribution of the proposed angles to identifying conformers.The OC20 IS2RE data include 660,010 samples, combining adsorbates and catalysts with multiple conformers.Source: Table 5 · p. 10, §6.4 · p. 10

7 Conclusions, Limitations, Outlook, and Societal Impacts

ComENet is presented as a complete, efficient 3D GNN, with 6–10 times faster training and inference and superior or comparable performance. The authors identify the costly acquisition of 3D data as a limitation shared by existing 3D GNNs and outline generation and 2D–3D learning directions.

7 Conclusions, Limitations, Outlook, and Societal Impacts

  • ComENet accelerates training and inference by 6–10 times while achieving superior or comparable performance on large-scale datasets.The authors describe it as complete in incorporating 3D information, with time complexity O(nk).Source: §7 · p. 10
  • Acquiring 3D information is difficult and expensive; current approaches rely on experiments or time-consuming DFT-based computing.The authors identify this as a major limitation shared by ComENet and existing 3D GNNs.Source: §7 · p. 10
  • The authors propose generating 3D graphs from 2D graphs or from scratch, noting that 2D molecules are often unavailable in drug discovery.Source: §7 · p. 10
  • A second proposed direction uses contrastive learning to align 2D graphs with 3D views when only some training data have 3D information.The authors suggest integrating these components into end-to-end systems for applications based on 2D graph data.Source: §7 · p. 10
  • The authors identify drug and material discovery, quantum chemistry and physics, materials science, and molecular dynamics as potential application areas.They state that negative societal impacts associated with those applications and domains also apply to ComENet.Source: §7 · p. 10

A.1 Algorithm for the Proposed Message Passing

The proposed message-passing algorithm computes pairwise distances, selects reference nodes, derives geometric quantities, and updates node features. Its nested loops yield O(nk) complexity.

A.1 Algorithm for the Proposed Message Passing

  • The algorithm loops over nodes and their k neighbors, computes distances, selects reference nodes, and updates node features using the proposed message passing.Reference-node selection within the inner loop has O(1) complexity, giving overall O(nk) complexity.Source: §A.1 · p. 15
  • Reference nodes are selected by minimizing distances to neighboring nodes, first for f_i and then for s_i among neighbors other than f_i.Source: §A.1 · p. 15
  • The message-passing procedure has two nested loops with O(n) outer-loop and O(k) inner-loop complexity.The authors report that ComENet is 6–10 times faster than SphereNet and has training time similar to SchNet.Source: §A.1 · p. 15

A.2 SE(3)

The section defines SE(3) as the group of rotations and translations in three dimensions and specifies its matrix representation. Its rotation component belongs to SO(3).

A.2 SE(3)

  • SE(3) comprises rotations and translations in 3D and is represented by 4 × 4 real matrices.Source: §A.2 · p. 15
  • The 4 × 4 transformation matrix contains a 3 × 3 rotation component and a translation vector.The passage writes the translation entries as t1, t2, and t3.Source: §A.2 · p. 15
  • The rotation component Rot belongs to SO(3), whose matrices satisfy Rot^T Rot = I and det(Rot) = 1.Source: §A.2 · p. 15

A.3 Proofs of Lemma 1

The proof of Lemma 1 establishes that a new node’s relative location is uniquely determined in two neighborhood cases. It uses contradiction to rule out a second location in each case.

A.3 Proofs of Lemma 1

  • The proof separates nodes with multiple neighbors from nodes with only one neighbor and seeks uniqueness of the new node’s relative location in both cases.The figure illustrates these two cases.Source: §A.3 · p. 15, Figure 6 · p. 16
  • With multiple neighbors, the proof shows that (d_ij, θ_ij, φ_ij) uniquely determines p_ij by assuming a second location and deriving a contradiction.The projection operator T is onto the plane perpendicular to p_if_i.Source: §A.3 · p. 16
  • With one neighbor, the proof uses (d_ij, θ_ij, τ_c1i) to establish that p_ij is uniquely determined, again by contradiction.The proof states that the quantities on the right side are known.Source: §A.3 · pp. 16–17

A.4 Model architecture

The architecture converts local and global 3D geometries into basis-function vectors used in separate convolutions to update node features. The section also describes the datasets, splits, and training setups used for evaluation.

A.4 Model architecture

  • ComENet transforms local (d, θ, φ) and global (d, τ) geometries into basis vectors used as edge weights in separate convolutions that update node features.Distance is included in both representations; the local and global convolution outputs are concatenated before linear layers produce the updated feature vector.Source: §A.4 · p. 17
  • Among tested basis choices, spherical Bessel functions and spherical harmonics perform best for the local and global geometry tuples.The tested alternatives include MLP, Gaussian, and sine functions.Source: §A.4 · p. 17

A.5 Data Description

  • OC20 experiments target IS2RE, the common catalysis task, using 460,328 training structures and four validation splits: ID, OOD Ads, OOD Cat, and OOD Both.The dataset models catalysts, and IS2RE predicts relaxed energy; relaxed energies are often correlated with catalyst activity and selectivity.Source: §A.5 · p. 17
  • Molecule3D contains 3,899,647 DFT-ground-state molecules and evaluates property prediction with random and scaffold 6:2:2 splits, the latter imposing distribution shift.The dataset was designed for predicting 3D geometries from molecular graphs, whereas this work predicts molecular properties from geometries.Source: §A.5 · p. 18
  • QM9 includes 134k stable small organic molecules, split into 110,000 training, 10,000 validation, and 10,831 test molecules, with 12 target properties.The listed targets include geometric, energetic, electronic, and thermodynamic properties.Source: §A.5 · p. 18

A.6 Experimental Setup

  • Molecule3D and QM9 hyperparameters are selected on validation sets, while OC20 uses 10% of training data for tuning before retraining on all training data.OC20 uses a larger cutoff; OC20 and Molecule3D use larger hidden dimensions. DimeNet++ and SphereNet baseline batches on Molecule3D are capped at 32 by GPU memory.Source: §A.6 · p. 18
Loading 2206.08515v3…