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

TL;DR

UGrid delivers an efficient-and-rigorous neural PDE solver by integrating multigrid V-cycles with a U-Net-inspired architecture, ensuring convergence and correctness for linear PDEs via a mathematically grounded masked-operator framework. It replaces traditional smoothers with learnable, boundary-aware convolutional operators and optimizes a residual-based loss in a self-supervised manner, enabling robust generalization to unseen geometries. Empirical results across Poisson, Helmholtz, and convection-diffusion-reaction problems demonstrate substantial speedups (up to 10–20x over strong baselines) and scalability to XL/XXL problems without retraining. The work offers a practical pathway for data-driven PDE solvers that retain mathematical guarantees, with potential extensions to nonlinear PDEs in future work.

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.

Figures (10)

• Figure 1: Overview of our novel method. Given PDE parameters and its current numerical estimation, the smoothing operations are applied multiple times first. Then, the current residual is fed into our UGrid submodule (together with the boundary mask). Next, the solution is corrected by the correction term and post-smoothed. Collectively, it comprises one iteration of the neural solver. The UGrid submodule (detailed in Fig. \ref{['fig:ugrid_submodule']}) aims to mimic the multigrid V-cycle, and its parameters are learnable, so as to endow our framework with the ability to learn from data.
• Figure 2: An overview of our recursive UGrid submodule. The residual is smoothed by unbiased convolution layers, downsampled to be recursively updated by a $2$x-coarser UGrid submodule (note the orange recursive invocation of UGrid submodule in the middle), then upsampled back to the fine grid, smoothed, and added with the initial residual by skip-connection. Boundary values are enforced by interior mask via element-wise multiplication. The convolution layers (shown in orange) are learnable; other layers (shown in blue) are the fixed mathematical backbone.
• Figure 3: (a-j) illustrates the $\vb{f}$ field distributions of our testcases. The boundaries are shown in bold red lines for a better view. (k-m) are self-explanatory. The complete illustration of boundary values is available in Sec. \ref{['sec:supplemental_material:more_specs_on_testcases']}.
• Figure 4: Illustration of the Dirichlet boundary values of our ten testcases. Again, all boundaries are shown in bold for a better view. Note that these boundaries are not required to be constant and could have discontinuities, which could be observed in testcases (g-j).
• Figure 5: Convergence map on large-scale Poisson problem. The $x$ coordinates are time(s), shown in $\mathrm{ms}$; the $y$ coordinates are the relative residual errors, shown in logarithm ($\log(r \times 10^{5})$) for a better view.

Theorems & Definitions (9)

• Theorem 4.1
• Theorem 4.2
• Lemma 1.1
• proof
• Lemma 1.2
• proof
• Theorem 1.3
• proof
• proof