Learning An Efficient-And-Rigorous Neural Multigrid Solver

Xi Han, Fei Hou, Hong Qin

Primary Area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Keywords: Partial Differential Equations, Numerical Solver, Neural Solver, Multigrid Method

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

TL;DR: We propose UGrid, which is an efficient-and-rigorous neural multigrid solver for linear partial differential equations.

Abstract:

Partial Differential Equations (PDEs) and their efficient
numerical solutions are of fundamental significance to science and
engineering involving heavy computation. To date, the historical
reliance on legacy generic numerical solvers has circumscribed
possible integration of big data knowledge and exhibits sub-optimal
efficiency for certain PDE formulations. In contrast, AI-inspired
neural methods have the potential to learn such knowledge from big data
and endow numerical solvers with compact structures and high
efficiency, but still with unconquered challenges including, a lack of
sound mathematical backbone, no guarantee of correctness or
convergence, and low accuracy, thus unable to handle complex, unseen
scenarios. This paper articulates a mathematically rigorous neural PDE
solver by integrating iterative solvers and the Multigrid Method
with Convolutional Neural Networks (CNNs).
Our novel UGrid neural 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 and strong generalization power
to complicated cases not observed during the training phase. 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. We conduct extensive experiments on
Poisson's equations, and our comprehensive evaluations have confirmed
all of the aforementioned theoretical and numerical advantages.
Finally, a mathematically-sound proof affords our new method to
generalize to other types of linear PDEs.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 4258