Spatiotemporal Graph Learning with Direct Volumetric Information Passing and Feature Enhancement

TL;DR

This paper introduces CeFeGNN, a novel graph neural network framework that enhances spatiotemporal modeling of physical systems by embedding volumetric information and feature enhancement, outperforming existing methods.

Contribution

The paper proposes a dual-module GNN with cell embeddings and feature enhancement, upgrading local aggregation to higher order and improving modeling of volumetric data in physical systems.

Findings

CeFeGNN outperforms baseline models on PDE systems.

Embedding cell attributions improves spatial dependency capture.

Feature enhancement alleviates over-smoothness in GNNs.

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.