Improving Graph Neural Networks by Learning Continuous Edge Directions

TL;DR

The paper presents Continuous Edge Direction (CoED) GNN, a framework that replaces diffusion with directional information flow by assigning fuzzy, continuously tunable edge directions to graph edges. A novel complex-valued fuzzy Laplacian, with Real and Imaginary parts encoding in- and out-neighbor information, yields an expressive aggregation mechanism and, under a WL-like isomorphism framework, matches a weak form of WL for directed graphs with fuzzy edges. CoED learns edge directions jointly with GNN weights in an end-to-end manner, enabling long-range information transmission and mitigating oversmoothing, particularly in graph-ensemble settings where graph structure is fixed but node features vary. Empirical results across synthetic and real datasets—gene regulatory networks, web traffic, and power grids—show CoED outperforms baselines, validating the practical impact of learning continuous edge directions for directed and undirected graphs alike.

Abstract

Graph Neural Networks (GNNs) traditionally employ a message-passing mechanism that resembles diffusion over undirected graphs, which often leads to homogenization of node features and reduced discriminative power in tasks such as node classification. Our key insight for addressing this limitation is to assign fuzzy edge directions -- that can vary continuously from node $i$ pointing to node $j$ to vice versa -- to the edges of a graph so that features can preferentially flow in one direction between nodes to enable long-range information transmission across the graph. We also introduce a novel complex-valued Laplacian for directed graphs with fuzzy edges where the real and imaginary parts represent information flow in opposite directions. Using this Laplacian, we propose a general framework, called Continuous Edge Direction (CoED) GNN, for learning on graphs with fuzzy edges and prove its expressivity limits using a generalization of the Weisfeiler-Leman (WL) graph isomorphism test for directed graphs with fuzzy edges. Our architecture aggregates neighbor features scaled by the learned edge directions and processes the aggregated messages from in-neighbors and out-neighbors separately alongside the self-features of the nodes. Since continuous edge directions are differentiable, they can be learned jointly with the GNN weights via gradient-based optimization. CoED GNN is particularly well-suited for graph ensemble data where the graph structure remains fixed but multiple realizations of node features are available, such as in gene regulatory networks, web connectivity graphs, and power grids. We demonstrate through extensive experiments on both synthetic and real graph ensemble datasets that learning continuous edge directions significantly improves performance both for undirected and directed graphs compared with existing methods.

Figures (14)

• Figure 1: (a) When edges are undirected, information diffuses across the graph and long-range transmission of information between nodes 1 and 2 is not possible. (b) Once the optimal edge directions are learned, information can flow directly from node 1 to node 2.
• Figure 2: Schematic of training with a graph ensemble data. The input graph is undirected (left box). The graph ensemble data contains multiple realizations of node features and corresponding target values, either at the node, edge, or graph level. The phase angle formulation allows continuous edge directions to be optimized alongside the GNN parameters in an end-to-end manner (middle box). The learned edge directions (right box) enable long range information transmission across the graph.
• Figure 3: Synthetic datasets. (a-b) Triangular lattice graph with edge directions derived from the gradient of a 2d potential function $V$ (shown in a), creating long-range flows across the graph. (c-d) Gene regulatory network (GRN) represented as a directed graph where nodes are genes and edges denote interactions. Steady-state gene expression levels are obtained from GRN dynamics, with perturbations simulated by setting the expression levels of specific genes to zero.
• Figure 4: Comparison of different models on the synthetic datasets. Values are test losses reported with a common factor of $10^{-3}$ in both columns.
• Figure 5: Learned theta vs. true theta for CoED applied to directed flow on triangular lattice synthetic dataset.

Theorems & Definitions (1)

• Theorem 1