Graph Learning with Distributional Edge Layouts

TL;DR

This work introduces Distributional Edge Layouts (DELs), a physics-inspired pre-processing technique for graph neural networks that samples a distribution of edge layouts from a Boltzmann energy surface using Langevin dynamics: ${\mathbb{P}}(\mathbf{C})=(1/Z) e^{-E(\mathbf{C})/\alpha}$. By converting steady-state layouts into edge features and embedding them into GNNs through edge-aware mechanisms, DEL captures a broader energy landscape and provides additional expressivity beyond the 1-WL test, while remaining architecture-agnostic. Empirically, DEL-enhanced backbones (GAT, Graph Transformer, GPS) achieve substantial improvements across six datasets, with analyses showing stable isomorphic layouts and discriminative power for certain non-isomorphic graphs. The approach is scalable as a pre-processing step and opens the door to leveraging global layout distributions to augment a wide range of GNN models in practice.

Abstract

Graph Neural Networks (GNNs) learn from graph-structured data by passing local messages between neighboring nodes along edges on certain topological layouts. Typically, these topological layouts in modern GNNs are deterministically computed (e.g., attention-based GNNs) or locally sampled (e.g., GraphSage) under heuristic assumptions. In this paper, we for the first time pose that these layouts can be globally sampled via Langevin dynamics following Boltzmann distribution equipped with explicit physical energy, leading to higher feasibility in the physical world. We argue that such a collection of sampled/optimized layouts can capture the wide energy distribution and bring extra expressivity on top of WL-test, therefore easing downstream tasks. As such, we propose Distributional Edge Layouts (DELs) to serve as a complement to a variety of GNNs. DEL is a pre-processing strategy independent of subsequent GNN variants, thus being highly flexible. Experimental results demonstrate that DELs consistently and substantially improve a series of GNN baselines, achieving state-of-the-art performance on multiple datasets.

Figures (10)

• Figure 1: DEL pipeline. Our framework can be summarized into the following three steps: 1) Sampling Layouts: Sampling a set of layouts with local minimum energy for a given connectivity based on the Boltzmann distribution with the help of the energy-based layout algorithms. 2) Constructing Edge Features: Once we have a set of sampled layouts, we construct edge features based on these layouts. These edge features are designed to capture the potential energy surface of the system. Models can gain insights into the potential energy landscape and identify regions of high or low energy through the edge features. 3) Edge-Aware GNNs: In this final step, we combine the constructed edge features with a GNN backbone and use edge-aware GNNs to improve the performance of downstream tasks.
• Figure 2: Comparing non-isomorphic graphs (Top: Decalin, Bottom: Bicyclopentyl) through a straightforward approach for extracting graph-level features. We sample 50 layouts to calculate graph total weights (GTW) distribution and label the minimum, maximum, and highest frequency GTWs.
• Figure 3: The layout energy change trend of the iterative process of the Fruchterman-Reingold layout algorithm for different datasets. From left to right present the average layout energy change trend of MUTAG, NCI1, PROTEINS, D&D, and IMDB datasets. In the IMDB dataset, all graphs consist of densely connected structures that tend to form layouts covering the entire plane, resulting in stability but also high potential energy.
• Figure 4: Heatmap, MDS layout and graph layout examples of MUTAG dataset.
• Figure 5: Heatmap, MDS layout and graph layout examples of PROTEINS dataset