Flow-Attentional Graph Neural Networks

TL;DR

FlowA tten tion redefines attention in GNNs to normalize across outgoing edges, aligning message passing with Kirchhoff's first law $ \sum_{u \in \mathcal{N}_{in}(v)} \psi(u,v) = \sum_{u \in \mathcal{N}_{out}(v)} \psi(v,u)$. This simple change enhances expressivity by distinguishing node multisets that standard attention cannot and enables FlowDAGNN to distinguish DAGs from their computation trees. The authors validate the approach on flow-graph tasks (power grids and electronic circuits), showing consistent improvements in graph-level classification and regression over strong baselines, and demonstrate applicability to DAGs with a reversed-forward flow architecture. They also discuss limitations on non-flow graphs and highlight potential integrations with Graph Transformer models for broader impact on flow-structured domains like utilities and logistics.

Abstract

Graph Neural Networks (GNNs) have become essential for learning from graph-structured data. However, existing GNNs do not consider the conservation law inherent in graphs associated with a flow of physical resources, such as electrical current in power grids or traffic in transportation networks, which can lead to reduced model performance. To address this, we propose flow attention, which adapts existing graph attention mechanisms to satisfy Kirchhoff$\text{'}$s first law. Furthermore, we discuss how this modification influences the expressivity and identify sets of non-isomorphic graphs that can be discriminated by flow attention but not by standard attention. Through extensive experiments on two flow graph datasets (electronic circuits and power grids) we demonstrate that flow attention enhances the performance of attention-based GNNs on both graph-level classification and regression tasks.

Figures (8)

• Figure 1: Two non-isomorphic graphs that are equivalent as informational graphs, but different as resource flow graphs. a The two different directed graph structures represent the same computation (example adapted from zhang2019dvae). b The same graph structures as above represent different electronic circuits.
• Figure 2: a Standard graph attention mechanism as it is applied in attentional GNNs. The attention weights associated with edges of the same color sum to 1. b The proposed flow attention mechanism applied in FlowGNNs. The flow attention weights associated with edges of the same color sum to 1. c Two snapshots during the reverse and forward pass of FlowDAGNN. Nodes marked in green have already been updated.
• Figure 3: Two non-isomorphic DAGs together with their corresponding computation trees, which are equivalent. Distinct node features are visualized by different colors. The middle and right columns show some example standard and flow attention weights. While the standard attention weights are always the same for both DAGs, the flow attention weights are different.
• Figure 4: A simple example circuit described as a DAG, which explains why the reverse pass is necessary in FlowDAGNN. Without the reverse pass, the flow attention weights from the input node to the $R_1$ nodes would be identical, whereas the electrical current flow is different due to $R_1 \ll R_2$.
• Figure 5: Distribution of the number of nodes (left) and number of edges (right) within the Ckt-Bench101 dataset dong2023cktgnn.

Theorems & Definitions (12)

• Lemma 4.1
• Corollary 4.2
• Lemma 5.1
• Lemma 5.2
• Corollary 5.3
• Theorem 5.4
• proof
• proof : Proof of Lemma \ref{['lem:attention_expressivity']}
• proof : Proof of Corollary \ref{['cor:attention_dags']}
• proof : Proof of Lemma \ref{['lem:absolute_flow']}