Topology-Aware Multiscale Mixture of Experts for Efficient Molecular Property Prediction

TL;DR

This work tackles the limitation of fixed interaction scales in 3D molecular GNNs by introducing MI-MoE, a topology-aware multiscale mixture of experts. The framework builds a multiscale graph filtration, defines cutoff-based experts, and uses a topology-driven gate that leverages filtration descriptors, including persistent homology, to route inputs to appropriate experts. Empirically, MI-MoE improves performance across multiple backbone architectures and tasks on MoleculeNet and polymer benchmarks, demonstrating strong gains in both regression and classification settings. The approach is plug-and-play, data-efficient, and presents a principled route to adaptively modeling molecular interactions across geometric scales with topology-informed routing.

Abstract

Many molecular properties depend on 3D geometry, where non-covalent interactions, stereochemical effects, and medium- to long-range forces are determined by spatial distances and angles that cannot be uniquely captured by a 2D bond graph. Yet most 3D molecular graph neural networks still rely on globally fixed neighborhood heuristics, typically defined by distance cutoffs and maximum neighbor limits, to define local message-passing neighborhoods, leading to rigid, data-agnostic interaction budgets. We propose Multiscale Interaction Mixture of Experts (MI-MoE) to adapt interaction modeling across geometric regimes. Our contributions are threefold: (1) we introduce a distance-cutoff expert ensemble that explicitly captures short-, mid-, and long-range interactions without committing to a single cutoff; (2) we design a topological gating encoder that routes inputs to experts using filtration-based descriptors, including persistent homology features, summarizing how connectivity evolves across radii; and (3) we show that MI-MoE is a plug-in module that consistently improves multiple strong 3D molecular backbones across diverse molecular and polymer property prediction benchmark datasets, covering both regression and classification tasks. These results highlight topology-aware multiscale routing as an effective principle for 3D molecular graph learning.

Figures (2)

• Figure 1: Overview of the MI-MoE framework. Starting from a molecular point cloud $\mathcal{P}$, a distance-based filtration induces a family of multiscale interaction graphs. From this filtration, two graph sets are derived: (i) a sparse set of cutoff-specific graphs $\{\mathcal{G}^{(c_k)}\}_{k=1}^{K}$ used as inputs to the expert models, and (ii) a denser set of graphs $\{\mathcal{G}^{(r_t)}\}_{t=1}^{T}$ used to compute topological descriptors $\widetilde{\mathrm{R}}(\mathcal{G}^{(r_t)}),\widetilde{\mathrm{W}}(\mathcal{G}^{(r_t)}),E(\mathcal{G}^{(r_t)}),\widetilde{\beta}_0(\mathcal{G}^{(r_t)}),\widetilde{\beta}_1(\mathcal{G}^{(r_t)})$ that characterize the evolution of molecular connectivity across geometric scales for routing. These descriptors are fed into a topological gating network that produces interaction-dependent routing weights $\boldsymbol{\alpha}$ over experts. Finally, the selected expert outputs are aggregated according to $\boldsymbol{\alpha}$ to produce the final molecular representation $\mathbf{h}$.
• Figure 2: MI-MoE performance across interaction radii.