SE3Set: Harnessing equivariant hypergraph neural networks for molecular representation learning

TL;DR

SE3Set addresses the challenge of learning molecular representations that respect 3D rotational symmetry while capturing higher-order, many-body interactions. It introduces a fragmentation-based hypergraph construction that fuses 2D chemical topology with 3D geometry, and an $SE(3)$-equivariant hypergraph attention architecture built on AllSet and Equiformer. The method achieves competitive results on QM9 and MD17 and yields about a $20\%$ MAE improvement on MD22, signaling the significance of higher-order interactions in larger molecules. These capabilities provide a robust, physically grounded framework for accurate molecular property prediction and simulation.

Abstract

In this paper, we develop SE3Set, an SE(3) equivariant hypergraph neural network architecture tailored for advanced molecular representation learning. Hypergraphs are not merely an extension of traditional graphs; they are pivotal for modeling high-order relationships, a capability that conventional equivariant graph-based methods lack due to their inherent limitations in representing intricate many-body interactions. To achieve this, we first construct hypergraphs via proposing a new fragmentation method that considers both chemical and three-dimensional spatial information of molecular system. We then design SE3Set, which incorporates equivariance into the hypergragh neural network. This ensures that the learned molecular representations are invariant to spatial transformations, thereby providing robustness essential for accurate prediction of molecular properties. SE3Set has shown performance on par with state-of-the-art (SOTA) models for small molecule datasets like QM9 and MD17. It excels on the MD22 dataset, achieving a notable improvement of approximately 20% in accuracy across all molecules, which highlights the prevalence of complex many-body interactions in larger molecules. This exceptional performance of SE3Set across diverse molecular structures underscores its transformative potential in computational chemistry, offering a route to more accurate and physically nuanced modeling.

Figures (4)

• Figure 1: Folic acid fragmentation illustrated with CID 135398658 from PubChem. (a) Preprocessing to identify cleavable bonds for fragmentation. (b) Initial fragments formed using BFS, color-coded by functional groups (blue), rings (orange), and single atoms (green). (c) Fragments merged to satisfy atom count criteria, detailed in \ref{['appendix:fragment_rules']}. (d) Expansion of fragments shown with directional arrows.
• Figure 2: Overall architecture of SE3Set. (a) SE3Set begins with node and hyperedge embeddings, cycles through V2E and E2V attention modules for iterative updates, and concludes with normalization and a feed-forward block for output. (b) Embedding. Atomic numbers and position vectors are transformed into initial embeddings for nodes and hyperedges. (c) Attention Block. Merges feature sets with positional or hyperedge data for feature processing. (d) Feed-Forward Block. Enhances feature sets through a streamlined network. (e) V2E Module. Utilizes node features and their relative positions to update hyperedge features. (f) E2V Module. Employs hyperedge features to refresh node features, using tensor products (left) or summation (right) for updates. Symbols $\otimes$, $\oplus$, and $\odot$ in figures denote depth-wise tensor product, summation, and Hadamard multiplication, respectively. $h_i^{\alpha}$ represents hyperedge features, $x_i$ is for node features, superscript $n$ indicates the number of updates, and $\vec{r}_{ij}$ is the relative position vector between nodes $i$ and $j$.
• Figure 3: Ablation studies on the QM9 dataset's HOMO task (units: $\rm{meV}$). The variable $c_w$ represents the threshold for expansion in the fourth step of fragmentation, guided by the fragment bond order defined in Eq. \ref{['eq:bo_lendvay']}. The term BRICS denotes another fragmentation method implemented in RDKits. Additionally, the E2V summation refers to the architectural framework specified from Eq. \ref{['eq:E2V_sum_begin']} to Eq. \ref{['eq:E2V_sum_end']}.
• Figure 4: Distribution of fragments in QM9 dataset. (a) Fragment Count Distribution. The distribution remains consistent regardless of the value of $c_w$ or the bond order calculation method employed. (b) Molecule Size vs. Fragment Count Distribution. Generally, the more atoms molecule has, the more fragments will generate. It is also invariant for $c_w$ or bond order calculation scheme. Average Atom Count per Fragment Distribution (c) $c_w=0.1$, (d) $c_w=0.05$, (e) $c_w=0.01$ for Lendvay bond order and (f) $c_w=0.4$ and (g) $c_w=0.2$ for exponential bond order, respectively. (h) BRICS Fragment Count Distribution. (i) BRICS Molecule Size vs. Fragment Count Distribution (j) BRICS Average Atom Count per Fragment.

Theorems & Definitions (2)

• Definition 4.1
• Definition 4.2