TopoMole: Topological Message Passing Meets Hyperedge Messages

Pablo Martínez Crespo, Robert S. Jordan, Marisa Gliege, Santiago Miret, Vijay Kris Narasimhan, Rocío Mercado

Keywords: topological message passing, molecular property prediction, graph neural networks, hypergraph neural networks

TL;DR: We present the first JAX-based framework for topological message passing, enabling the representation of high order molecular structures and long-range information exchange. A variety of topological models is presented and tested on the QM9 dataset.

Abstract: 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.

Submission Track: Paper Track (Full Paper)

Submission Category: AI-Guided Design

Institution Location: Gothenburg, Sweden

AI4Mat RLSF: Yes

Submission Number: 45