A Variational Mean-Field Control Framework for Graph Representation Learning.

Abstract

Feature representation learning in graph neural networks (GNNs) is a dynamic process driven by progressive information exchange throughout the graph. Current GNNs typically apply pre-defined message-passing heuristics uniformly across all graph data, even when the assumed relational inductive bias conflicts with the intrinsic graph properties (e.g., heterophily). This limitation hinders a principled understanding of graph learning and motivates a shift from application-specific architectures toward a generalizable paradigm. In this work, we conceptualize GNN learning through the lens of mean-field control, where the desired learning outcome corresponds to a variational critical point, each node reaches a representation that is adaptively calibrated to both itself and the entire graph. Building on this formulation, we propose a mathematical framework based on mean-field control (MFC) to adaptively design GNNs for unseen graph data. Specifically, we jointly optimize two control patterns: diffusive mobility, governing information propagation across the graph, and reactive mobility, regulating feature transformation at individual nodes. Both controls are learned from each input instance by solving an MFC variational problem through Hamiltonian flows characterized by partial differential equations (PDE). Our variational framework unifies existing PDE-based GNNs as special mean-field games with fixed control patterns, and yields an end-to-end deep model, termed Nash-GNN. Extensive experiments show that Nash-GNN achieves state-of-the-art performance across diverse benchmarks, including heterophilic graphs and human connectomes. Beyond empirical gains, the MFC framework offers a principled lens to examine the mechanisms of graph representation learning through collective dynamics, opening promising directions for next-generation GNN design.