E(n) Equivariant Topological Neural Networks
Synergizing higher-order interactions and (Euclidean) symmetries in data
Graph neural networks are limited in handling higher-order interactions. Topological Deep Learning (TDL) offers a powerful solution leveraging topological spaces. However, little is known about integrating geometric features in the TDL machinery. For this reason, we introduced E(n)-Equivariant Topological Neural Networks (ETNNs), TDL architectures equivariant to rotations, reflections, and translations, ready for heterogeneous settings, and provably expressive. ETNNs generalize most graph and TDL models, acting as a versatile matrix for equivariant architectures on complex relational structures.
This post is based on the paper by Claudio Battiloro*, Ege Karaismailoğlu*, Mauricio Tec*, George Dasoulas*, et al., “E(n) Equivariant Topological Neural Networks”, 2024, a collaboration with Michelle Audirac and Francesca Dominici. The code can be found here. *Equal contribution.
Graph Neural Networks (GNNs) have become prominent in several domains, due to their ability to merge neural networks flexibility with the knowledge of data relationships modeled through graphs. However, while many real-world complex systems comprise higher-order interactions, GNNs can explicitly model only pairwise interactions.
Topological Deep Learning (TDL)[2][3] is gaining traction for its ability to capture higher-order interactions using tools from algebraic topology, especially combinatorial topological spaces, such as simplicial or cell complexes. Leveraging this combinatorial structure, it is possible to design more sophisticated adjacency schemes among cells, i.e. single nodes or groups of nodes, than the usual adjacency in graphs.
Little is known, however, about how to integrate and leverage geometric information of the data in TDL models. In this context, nodes are embedded in a manifold or, in general, in a metric space, and have geometric features, such as positions or velocities. Motivated by this, in our recent paper [1] we introduced E(n)-Equivariant Topological Neural Networks (ETNNs), i.e. TDL architectures being equivariant w.r.t. rotations, reflections, and translations. ETNNs run over Combinatorial Complexes (CCs), being a generalization of graphs, simplicial complexes, path complexes, cell complexes, and hypergraphs, therefore ETNNs represent a legitimate framework for designing arbitrary E(n)-Equivariant Networks over combinatorial topological spaces or, in general, over combinatorial objects used to model interactions, such as hierarchical graphs.
“ETNNs represent a legitimate framework for designing arbitrary E(n)-Equivariant Networks over combinatorial topological spaces”
Why Do Higher-order Interactions Matter?
Topological Deep Learning (TDL) promotes significant advancements in understanding complex data relationships, particularly by addressing the limitations of traditional methods that focus solely on pairwise interactions. The importance of Topological Neural Networks (TNNs) lies in their ability to model higher-order interactions, which are crucial for capturing the multifaceted nature of real-world systems.
For instance, in biological networks, the interactions between proteins, genes, and other biomolecules often involve more than two entities and can have a pretty long range, leading to significant implications for understanding diseases and developing treatments. In social networks, group dynamics and community structures cannot be fully understood through pairwise relationships alone. Similarly, in recommender systems, the context and collective preferences of groups offer insights beyond individual pairwise recommendations.
By leveraging the capacity to model these complex relationships, TDL offers a more accurate and detailed representation of the data, leading to improved predictive performance and deeper insights.
The importance of Topological Deep Learning lies in its ability to model higher-order interactions, crucial for capturing the multifaceted nature of real-world systems
Message Passing Networks over Combinatorial Complexes
The foundational block of the class of TNNs we are interested in is usually message passing, this time not over graphs, but rather over a more general, higher-order object: Combinatorial Complexes (CCs).
Suppose to have a set of entities, referred to as nodes. An edge between two nodes is a subset comprising two nodes, and it is useful to model pairwise interactions among entities. A triangle between three nodes is a subset of nodes of cardinality 3, and it is useful to model triadic interactions. In general, a subset of k nodes, which we refer to as a cell, is useful to model k-adic interactions. While cells can model higher-order interactions, they don’t carry any information about how these higher-order interactions should be organized. In other words, they don’t come with a hierarchy. For this reason, a CC is a triple (S,X,r) made of a set of nodes S, a set of cells X, and a rank function r assigning a rank (or dimension) to each cell. Usually, nodes are rank 0 cells, edges are rank 1 cells. Formal mathematical treatment can be found in [2][1].
In Appendix B of our paper [1], we explained in detail how graphs, hypergraphs, simplicial, and cell complexes can be obtained as particular instances of a combinatorial complex.
In a combinatorial complex, we can leverage a way more sophisticated adjacency structure among cells than the usual node adjacency in graphs. We can do it by equipping a CC with a notion of neighborhood among cells. In particular, we can define a collection of neighborhood functions CN assigning a non-empty collection of “neighbor cells” to each cell in the complex. Usually, two main types of neighborhood functions are considered: adjacencies, in which all of the neighbors of a cell x have the same rank as x, and incidences, in which the neighbors have a higher or lower rank. Examples are the up/down incidences, defined by the containment criterion, and the up/down adjacencies, defined by the common neighbor criterion.
For instance, two nodes are up adjacent if they are connected by an edge, an edge is down incident to a node if the node is an endpoint of the edges, two edges are down adjacent if they have a common endpoint node, and they are up adjacent if they are both part of a higher rank cell. The same arguments apply to all ranks.
TNNs are then Message Passing (MP) networks over combinatorial complexes, in which cells update their embeddings by exchanging messages with neighboring cells, as induced by the neighborhood functions in CN. More formally, cell x will update its embedding h as
where the inner operator is an intra-neighborhood aggregation, and the outer operator is an inter-neighborhood aggregation. In other words, the embedding of cell x is updated by first exchanging and aggregating messages with its neighbors per each neighborhood, and then by further aggregating them across neighborhoods.
“In a combinatorial complex, we can leverage a way more sophisticated adjacency structure among cells than the usual node adjacency in graphs”
E(n)- Equivariant Topological Neural Networks
Let us now introduce the setting in which nodes (rank 0 cells) are embedded in some Euclidean space, i.e. they come with both non-geometric and geometric features, such as positions. Since many problems exhibit n-dimensional rotation, reflection, and translation symmetries, it is desirable to design models that are equivariant w.r.t. to them, i.e. w.r.t. the action of the Euclidean group E(n).
As has been done for graphs [5] and simplicial complexes [6], we represent equivariance by scalarization. Therefore, geometric features undergo an initial transformation into invariant scalars. Subsequently, they are processed in a learnable way (e.g., MLPs) before being combined along the original directions to achieve equivariance.
More formally, in an E(n)-Equivariant Topological Neural Network (ETNN) layer, the cell embeddings h and the node positions x will be updated as
where Inv() takes as input the positions of the nodes forming the involved cells, and it is an invariant function w.r.t. the action of E(n). We refer to a specific Inv() as a geometric invariant. In the paper, we propose several possible geometric invariants [1].
ETNNs are inherently able to model heterogeneous settings, as two cells can be neighbors in multiple neighborhoods, but messages among them will be dependent on the neighborhood itself. Thus, this property enables the modeling of different relation types through different (possibly overlapping) neighborhoods.
TNNs have been shown to be more expressive than vanilla GNNs in distinguishing isomorphic graphs [7]. Similarly, in our work we showed that ETNNs are at least as expressive as E(n)-Equivariant GNNs (EGNNs) [5] in distinguishing k-hop distinct geometric graphs [8]. In most cases, ETNNs are strictly more powerful than EGNNs. The quantitative gap between ETNNs and EGNNs is dependent on how the original graphs are lifted into a CC, i.e. how the set of cells X and the collection of neighborhood functions CN are defined.
In Appendix B of our work [1], we described how several known and new equivariant architectures can be derived from the ETNN framework. Therefore, ETNNs are general, heterogeneous, and more expressive.
“ETNNs are general, heterogeneous, and more expressive”
Real-World Combinatorial Complexes
To showcase the modeling power of CCs and their versatility in conjunction with ETNNs, we introduced two Real-World Combinatorial Complexes useful to model molecules and irregular multi-resolution geospatial data:
- Molecular CC. In a Molecular CC, the set of nodes is the set of atoms, while the set of cells and the rank function are such that: atoms are rank 0 cells, bonds and functional groups made of two atoms are rank 1 cells, and rings and functional groups made of more than two atoms are rank 2 cells. A Molecular CC can jointly retain all the properties of graphs, cell complexes, and hypergraphs, thus overcoming their mutual limitations. Please refer to the paper for further details [1].
- Spatial CC. In a Spatial CC, the set of cells is defined by hierarchical spatial objects. Two major examples are geometric descriptions (e.g. points, polylines, and polygons) or political partitions (e.g. zipcodes, counties, states) of a region. In the latter case, the listed objects can be modeled as cells of a CC and their rank correspond to their intrinsic dimensionality. Therefore, rank 0 cells are points, e.g. where air quality monitoring units are located, rank 1 cells are polylines, e.g. traffic roads, and rank 2 cells are polygons, e.g. census tracts.
Experiments
We first evaluated our model on molecular property prediction using the QM9 dataset, a comprehensive collection of quantum chemical calculations for small organic molecules. It includes geometric, energetic, electronic, and thermodynamic properties for 134k stable molecules with up to nine heavy atoms (C, O, N, F). We followed the experimental protocol of EGNN, and cast each molecule in the aforementioned Molecular CC model. Several details about the employed ETNN configurations and experimental settings are reported in the paper.
ETNN significantly improved the prediction performance of EGNN for 6 (reported in Table 1) out of the total 12 targets, achieving state-of-the-art performance for mu. In the remaining 6 targets, ETNN either matched or was slightly worse than EGNN. We put particular focus on making the results consistent with previous literature and reproducible (See Appendix I of the paper [2]). For instance, ETNN-graph-W is EGNN derived from the ETNN framework with our codebase and matches the number of parameters of the other tested ETNN configurations. Moreover, we made sure that EGNN derived from our codebase could perfectly reproduce all the results of the original EGNN [5] on all the properties.
Next, we introduced a novel Hyperlocal air pollution downscaling benchmark for TDL based on the geospatial CC model introduced above. The task consisted of predicting PM2.5 air pollution at a high resolution of 22m. The prediction targets of PM2.5 measurements are obtained from a publicly available dataset [9] consisting of measurements by mobile air sensors installed on cars, corresponding to the last quarter of 2021 in the Bronx, NY, USA. The targets are aggregated at the desired resolution. We mapped each point of the roads and census tracts and created a geospatial CC using custom neighborhood functions. Therefore, this was a semisupervised node regression task.
Tables 2–3 indicate that ETNN outperforms EGNN, GNN, and MLP, emphasizing the benefits of multi-level modeling combined with E(n) equivariance. We hope our benchmark will be adopted more and more by the TDL (and not only) community.
Conclusions
In this blogpost, we showcased our E(n) Equivariant Topological Neural Networks (ETNNs). ETNNs flexibly model hierarchical higher-order interactions being equivariant to rotations, reflections, and translations, proving effective in tasks of vastly different scales. ETNNs are a milestone for the TDL community, as they are a legitimate framework for designing arbitrary scalarization-based E(n) equivariant networks on arbitrary relational structures. To wrap up, the intuition that TDL proved to be true is: “Higher-order multiway interactions matter”[10]. With our ETNN, we can now reformulate it as: “synergizing higher-order multiway interactions and symmetries matters”.
The intuition that TDL proved to be true is: “Higher-order multiway interactions matter”. With our ETNNs, we can now reformulate it as: “synergyzing higher-order multiway interactions and symmetries matters”
References
[2] Hajij, Mustafa, et al. “Topological deep learning: Going beyond graph data.” 2022.
[10] Battiloro, Claudio. “Latent Topology Inference.” Blogpost, 2024.