E(n) Equivariant Topological Neural Networks

Abstract: Graph neural networks excel at modeling pairwise interactions, but they cannot flexibly accommodate higher-order interactions and features. Topological deep learning (TDL) has emerged recently as a promising tool for addressing this issue. TDL enables the principled modeling of arbitrary multi-way, hierarchical higher-order interactions by operating on combinatorial topological spaces, such as simplicial or cell complexes, instead of graphs. However, little is known about how to leverage geometric features such as positions and velocities for TDL. This paper introduces E(n)-Equivariant Topological Neural Networks (ETNNs), which are E(n)-equivariant message-passing networks operating on combinatorial complexes, formal objects unifying graphs, hypergraphs, simplicial, path, and cell complexes. ETNNs incorporate geometric node features while respecting rotation, reflection, and translation equivariance. Moreover, being TDL models, ETNNs are natively ready for settings with heterogeneous interactions. We provide a theoretical analysis to show the improved expressiveness of ETNNs over architectures for geometric graphs. We also show how E(n)-equivariant variants of TDL models can be directly derived from our framework. The broad applicability of ETNNs is demonstrated through two tasks of vastly different scales: i) molecular property prediction on the QM9 benchmark and ii) land-use regression for hyper-local estimation of air pollution with multi-resolution irregular geospatial data. The results indicate that ETNNs are an effective tool for learning from diverse types of richly structured data, as they match or surpass SotA equivariant TDL models with a significantly smaller computational burden, thus highlighting the benefits of a principled geometric inductive bias. Our implementation of ETNNs can be found at https://github.com/NSAPH-Projects/topological-equivariant-networks.

• The paper proposes a novel ETNN architecture that integrates geometric features with E(n) equivariance to model higher-order interactions in combinatorial complexes.
• It demonstrates broad applicability by excelling in tasks like molecular property prediction on QM9 and fine-grained air pollution downscaling using heterogeneous geospatial data.
• ETNNs efficiently compute geometric invariants such as pairwise and Hausdorff distances, setting a new benchmark in expressive, symmetry-preserving deep learning.

E(n) Equivariant Topological Neural Networks

Graph neural networks (GNNs) have gained prominence for their ability to address various applications involving graph-structured data, including computational chemistry, physics simulations, and social networks. However, the inherent pairwise nature of graphs restricts GNNs from effectively modeling higher-order interactions and features. Topological deep learning (TDL) endeavors to overcome these limitations by introducing more sophisticated structures such as simplicial and cell complexes. Despite the progress, the integration and leveraging of geometric features in TDL models have remained largely unexplored. The paper "E(n) Equivariant Topological Neural Networks (ETNNs)" proposes a novel architecture designed to address this gap.

Overview of ETNNs

ETNNs are message-passing networks operating over combinatorial complexes. These complexes generalize various combinatorial topological spaces, including graphs, hypergraphs, and higher structures such as simplicial and cell complexes. ETNNs incorporate geometric node features while ensuring E(n) group equivariance, which encompasses rotations and translations. By doing so, ETNNs cater to scenarios requiring heterogeneous interaction modeling, providing a more expressive framework compared to traditional GNNs.

Key Contributions

1. Expressive Power: The paper demonstrates that ETNNs surpass the expressiveness of existing architectures for geometric graphs, such as EGNNs, by showing how they can model complex hierarchical interactions in a principled manner. This is achieved through the use of geometric invariants and scalarization techniques, which maintain E(n) equivariance.
2. Broad Applicability: The authors extend ETNNs to a variety of practical tasks. The experiments highlight two distinct applications:
   • Molecular property prediction using the QM9 benchmark, which underscores the capability of ETNNs to handle richly structured molecular data.
   • Land-use regression for fine-grained air pollution estimation using heterogeneous geospatial data, demonstrating ETNNs' utility in integrating multi-resolution datasets.
3. Geometric Feature Integration: ETNNs exploit various geometric invariants, including pairwise distances, Hausdorff distances, and volumes of convex hulls. These invariants are computed efficiently and used to update non-geometric and geometric features of nodes within the complexes.

Experimental Results

1. Molecular Property Prediction: The evaluation on the QM9 dataset reveals that ETNNs outperform traditional methods like EGNN on multiple molecular properties. Notably, ETNNs achieve a new state-of-the-art performance for the dipole moment ($\mu$), with a reduction in the mean absolute error over prior models.
2. Air Pollution Downscaling: The paper introduces a new benchmark for hyper-local air pollution prediction using multi-resolution irregular geospatial data. ETNNs surpass baselines like MLP, GNN, and EGNN, demonstrating their effectiveness in handling complex heterogeneous data structures. The ablation studies further underscore the contributions of various components, such as the inclusion of geometric features and the presence of virtual cells.

Implications and Future Directions

The theoretical and practical advancements presented by ETNNs open several avenues for future research:

1. Extensions to Other Equivariances: Beyond E(n) equivariance, exploring symmetries for other groups or product groups could extend the applicability of ETNNs to more complex domains.
2. Dynamic Scenarios: ETNNs are currently designed for static settings. Future work could extend them to dynamic settings, enabling modeling of time-varying data and interactions.
3. Scalability Improvements: While ETNNs show promise for moderately sized datasets, scalability remains a challenge for very large datasets. Techniques like neighbor sampling and message-passing-free architectures could be explored for efficiency.
4. Unsupervised Learning: Developing unsupervised or self-supervised versions of ETNNs could enhance their ability to learn meaningful representations from large, unlabeled datasets.

Conclusion

E(n) Equivariant Topological Neural Networks represent a significant advancement in topological deep learning by offering a more expressive and flexible way to model higher-order, hierarchical interactions while respecting geometric symmetries. Their application to both molecular data and complex geospatial data sets showcases their versatility and capability. The paper provides a solid theoretical foundation and demonstrates practical efficacy, setting the stage for future explorations in this exciting interdisciplinary field.