Spatiotemporal Graph Learning with Direct Volumetric Information Passing and Feature Enhancement

TL;DR

CeFeGNN addresses the limitations of standard node–edge GNNs for spatiotemporal PDEs by introducing a two-level cell-embedded (CE) message-passing mechanism that leverages volumetric information, and a feature-enhanced (FE) block that expands representations via outer products and selective masking. The Encoder–Processor–Decoder architecture, with CE and FE blocks, enables higher-order spatial reasoning and mitigates over-smoothing while maintaining computational efficiency. Across Burgers, FitzHugh–Nagumo, Gray-Scott, and Black Sea datasets, CeFeGNN achieves superior generalization and accuracy, particularly in low-data regimes, outperforming baselines including MeshGraphNets, FNO, and Transolver. The work demonstrates that explicit higher-order structures and high-order feature interactions can markedly improve physics-informed spatiotemporal learning on arbitrary meshes, with potential extensions to finer meshes and richer geometric priors.

Abstract

Data-driven learning of physical systems has kindled significant attention, where many neural models have been developed. In particular, mesh-based graph neural networks (GNNs) have demonstrated significant potential in modeling spatiotemporal dynamics across arbitrary geometric domains. However, the existing node-edge message-passing and aggregation mechanism in GNNs limits the representation learning ability. In this paper, we proposed a dual-module framework, Cell-embedded and Feature-enhanced Graph Neural Network (aka, CeFeGNN), for learning spatiotemporal dynamics. Specifically, we embed learnable cell attributions to the common node-edge message passing process, which better captures the spatial dependency of regional features. Such a strategy essentially upgrades the local aggregation scheme from first order (e.g., from edge to node) to a higher order (e.g., from volume and edge to node), which takes advantage of volumetric information in message passing. Meanwhile, a novel feature-enhanced block is designed to further improve the model's performance and alleviate the over-smoothness problem. Extensive experiments on various PDE systems and one real-world dataset demonstrate that CeFeGNN achieves superior performance compared with other baselines.

Figures (10)

• Figure 1: Examples of datasets, including classic governing equations and complex real-world dataset. a, the 2D Burgers equation. b, the 2D Fitzhugh-Nagumo equation. c, the 2D Gray-Scott equation. d, the 3D Gray-Scott equation. e, the 2D Black-Sea dataset.
• Figure 2: Network architecture of CeFeGNN. a, an encoder encodes the physical variables to latent features, multiple message passing blocks process these latent features iteratively, and a decoder maps back to the physical states. b, three components in the CE block. c, the process of FE block. The definition of relative symbols refers to Section \ref{['Methodology']}.
• Figure 3: A scheme for reducing the number of parameters in FE block. Its quantitative experiment of the impact of window size and number of sub-features on CeFeGNN is shown in Table \ref{['tab:impact_of_window_size_and_number_of_sub-features']}.
• Figure 4: Cell in graphs. Green point is the centroid of cell.
• Figure 5: The test results of of all models on various datasets. a-d, the error distribution and the slices of generalization test on four grid-based datasets. The full snapshots of 3D GS dataset are displayed in Appendix Figure \ref{['fig:snapshot_3dgs']}. The results of 2D BS dataset are displayed in Figure \ref{['fig:bs_result']}. Error propagation curves refer to Appendix Section \ref{['Error propagation curve']}.

Theorems & Definitions (14)

• Lemma 3.1: Nonlinear Representation
• Definition 3.2
• Definition 3.3
• Corollary 3.4: Representation Power
• Definition 3.5: Cell in Graphs
• Corollary 3.6: Expressive Power
• Remark 3.7
• Lemma C.1: Feature Diversity
• Lemma C.2: Reduction in Ambiguity
• Corollary C.3: Improved Performance