Message Passing on the Edge: Towards Scalable and Expressive GNNs
Abstract
Graph neural networks (GNNs) are widely used in graph learning and most architectures propagate information by passing messages between vertices. In this work, we shift our attention to GNNs that perform message passing on edges and introduce EB-1WL, an edge-based color-refinement test, and a corresponding architecture, EB-GNN. Our EB-GNN architecture is inspired by the classic triangle-counting algorithm of Chiba and Nishizeki and passes messages along edges and triangles. Our contributions are as follows: 1. Theoretically, we show that EB-1WL is significantly more expressive than 1WL. We provide a complete logical characterization of EB-1WL in first-order logic, along with distinguishability results via homomorphism counting. To the best of our knowledge, EB-GNN has the strongest theoretical expressivity guarantees among edge-based message-passing GNNs in the literature. 2. Unlike many GNN architectures that are more expressive than 1WL, we prove that EB-1WL and EB-GNN admit near-linear time and memory usage on practical graph learning workloads. 3. We show in experiments that EB-GNN is a highly efficient general-purpose architecture: it substantially outperforms simple MPNNs and remains competitive with task-specialized state-of-the-art GNNs at substantially lower computational cost.
Lay Summary
Graphs are mathematical structures that model objects and their relationships, used to represent real-world phenomena like molecules, social networks, or road maps. Graph Neural Networks (GNNs) are the dominant approach for learning from graph data. Most GNNs pass messages between the nodes (the objects in the network), but in many applications like predicting chemical reactions it is more natural to focus on the edges (the relationships between objects). We introduce EB-GNN, a new GNN that passes messages along edges and triangles. EB-GNN can detect structural patterns in graphs that standard GNNs fundamentally cannot. We use the fact that real-world graphs are sparse, to make our GNN run in near-linear time, meaning it scales to large real-world graphs without becoming prohibitively expensive. We also provide two rigorous mathematical characterizations of exactly what EB-GNN can and cannot distinguish. In experiments, EB-GNN outperforms standard GNNs and remains competitive with highly specialized models, at a fraction of the computational cost.