TopoMole: Topological Message Passing Meets Hyperedge Messages

Message-passing neural networks (MPNNs) rely on pairwise edges and local neighborhoods to perform molecular property prediction tasks. As the number of atoms in a system increases, higher-order structures and long-range interactions become increasingly influential. Models for predicting macroscopic material properties and molecular properties of medium-sized molecules would hence benefit from frameworks that can naturally represent the information exchange in these systems. Precisely, topological message passing (TMP) and hypergraph neural networks (HNNs) extend MPNNs to operate on complex data relations by enabling the joint representation of nodes, edges, and higher-dimensional cells. In this workshop paper, we introduce TopoMole, the first open-source JAX package for TMP that supports the generation and aggregation of messages for all adjacency relations within a cell complex together with hyperedge representation. We demonstrate its utility in two molecular property prediction tasks, highlighting its potential in AI-driven materials discovery.