UGrid: An Efficient-And-Rigorous Neural Multigrid Solver for Linear PDEs

TL;DR

UGrid introduces a mathematically rigorous neural multigrid solver for linear PDEs, combining U-Net and MultiGrid principles to ensure convergence, correctness, and high accuracy, with improved training stability.

Contribution

This work presents the first neural PDE solver with proven convergence and correctness, integrating multigrid methods with neural networks for enhanced reliability.

Findings

Proves convergence and correctness of the neural solver.

Achieves high numerical accuracy across various PDEs.

Demonstrates strong generalization to different geometries and parameters.

Abstract

Numerical solvers of Partial Differential Equations (PDEs) are of fundamental significance to science and engineering. To date, the historical reliance on legacy techniques has circumscribed possible integration of big data knowledge and exhibits sub-optimal efficiency for certain PDE formulations, while data-driven neural methods typically lack mathematical guarantee of convergence and correctness. This paper articulates a mathematically rigorous neural solver for linear PDEs. The proposed UGrid solver, built upon the principled integration of U-Net and MultiGrid, manifests a mathematically rigorous proof of both convergence and correctness, and showcases high numerical accuracy, as well as strong generalization power to various input geometry/values and multiple PDE formulations. In addition, we devise a new residual loss metric, which enables unsupervised training and affords more stability and a larger solution space over the legacy losses.