Revisiting and Enhancing Graph Neural Networks through the Lens of Amortized Flows

Model architecture design plays a central role in the study of Graph Neural Networks (GNNs), as it determines the capacity of GNNs and thus significantly impacts their practical performance. In this study, we propose a new plug-and-play module called Amortized Flow (AF) to enhance the capacity of GNNs through the lens of minimum-cost flow. In particular, for an arbitrary GNN model, we formulate its learning task in a framework based on the minimum-cost flow problem defined on the target graph. In this paradigm, the fitting of observed data/labels is achieved jointly through the message passing of the GNN and a minimum-cost flow defined on the graph edges. The amortization of the flow results in the proposed AF module, which avoids the time-consuming optimization of the flow. Adding the AF module as a residual block to the GNN leads to a new GNN architecture, whose number of additional parameters is independent with the graph size. Moreover, considering this flow-based module leads to a generalized framework for GNN design, which highlights the usefulness of edge features and graph incidence matrix. Experiments demonstrate that the AF module applies to different GNNs, and its application consistently improves model performance in various node-level classification tasks. The code is available at https://github.com/minjiecheng/GNN-AF.