[Paper Review] GrADE: A graph based data-driven solver for time-dependent nonlinear partial differential equations

Y V Ravi Kumar, Souvik Chakraborty|arXiv (Cornell University)|Aug 24, 2021

Model Reduction and Neural Networks4 citations

TL;DR

This paper proposes GrADE, a novel graph-based data-driven solver for time-dependent nonlinear PDEs that combines graph neural networks (GNNs) for spatial modeling, neural ODEs for temporal evolution, and an attention mechanism to enhance feature weighting. The framework achieves accurate, scalable solutions for 1D and 2D Burgers’ equations with reduced training time and improved generalization using depth refinement, outperforming traditional FEM benchmarks with larger time steps and minimal data.

ABSTRACT

The physical world is governed by the laws of physics, often represented in form of nonlinear partial differential equations (PDEs). Unfortunately, solution of PDEs is non-trivial and often involves significant computational time. With recent developments in the field of artificial intelligence and machine learning, the solution of PDEs using neural network has emerged as a domain with huge potential. However, most of the developments in this field are based on either fully connected neural networks (FNN) or convolutional neural networks (CNN). While FNN is computationally inefficient as the number of network parameters can be potentially huge, CNN necessitates regular grid and simpler domain. In this work, we propose a novel framework referred to as the Graph Attention Differential Equation (GrADE) for solving time dependent nonlinear PDEs. The proposed approach couples FNN, graph neural network, and recently developed Neural ODE framework. The primary idea is to use graph neural network for modeling the spatial domain, and Neural ODE for modeling the temporal domain. The attention mechanism identifies important inputs/features and assign more weightage to the same; this enhances the performance of the proposed framework. Neural ODE, on the other hand, results in constant memory cost and allows trading of numerical precision for speed. We also propose depth refinement as an effective technique for training the proposed architecture in lesser time with better accuracy. The effectiveness of the proposed framework is illustrated using 1D and 2D Burgers' equations. Results obtained illustrate the capability of the proposed framework in modeling PDE and its scalability to larger domains without the need for retraining.

Motivation & Objective

• To address the computational inefficiency and limitations of traditional numerical solvers for nonlinear PDEs on irregular domains.
• To overcome the shortcomings of fully connected and convolutional neural networks in handling unstructured data and irregular geometries.
• To develop a differentiable, end-to-end framework that models spatial dynamics via graph attention and temporal dynamics via neural ODEs.
• To improve training efficiency and solution accuracy through a novel depth refinement strategy.
• To validate the framework on benchmark PDEs like Burgers’ equation with minimal training data and large time steps.

Proposed method

• The framework uses graph neural networks (GAT) to model spatial dependencies by representing the domain as a graph where nodes correspond to spatial points and edges to neighborhood relationships.
• An attention mechanism, implemented via feed-forward networks, dynamically assigns importance to neighboring nodes, enhancing feature representation and model expressivity.
• Temporal evolution is modeled using a Neural ODE, which parameterizes the time derivative as a continuous function, enabling constant memory cost and flexible time integration.
• The architecture employs a two-stage graph network: the first layer approximates first-order spatial derivatives, and the second layer approximates second-order derivatives, validated against central differences.
• A depth refinement technique is introduced to progressively increase network depth during training, improving convergence speed and accuracy.
• The model is trained using a physics-informed loss function that enforces PDE satisfaction at sparse time points, with no need for full trajectory supervision.

Experimental results

Research questions

• RQ1 Can a graph-based deep learning framework effectively model spatial dynamics in irregular domains for time-dependent PDEs?
• RQ2 How does combining graph attention with neural ODEs improve solution accuracy and training efficiency compared to standard PINNs or FNNs?
• RQ3 To what extent can the proposed depth refinement strategy reduce training time while maintaining or improving solution accuracy?
• RQ4 Can the framework generalize across larger domains and longer time horizons without retraining, using only two initial snapshots?
• RQ5 How well do the learned graph network layers approximate true spatial derivatives compared to numerical schemes like central differences?

Key findings

• GrADE achieves high accuracy in solving 1D and 2D Burgers’ equations using only two time snapshots as training data, significantly reducing data requirements.
• The model generalizes well to larger domains and longer time horizons without retraining, demonstrating strong scalability.
• The graph network layers accurately approximate first- and second-order spatial derivatives, with results closely matching central difference computations.
• Using a feed-forward network for attention weights yields slightly better accuracy than a Taylor net-based alternative in the GAT layer.
• Depth refinement reduces training time and improves convergence, enabling faster and more accurate training compared to constant-depth training.
• The neural ODE component enables constant memory cost and allows trade-offs between numerical precision and speed, enhancing computational efficiency.