Research map / Learning theory and algorithms / Combinatorial and search optimization

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

ICML 2024 · Learning theory and algorithms › Combinatorial and search optimization · approach: Predict-then-optimize learning

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes UGrid, a convolutional neural multigrid solver built from U-Net and multigrid V-cycle structure with a provably convergent masked iterator, trained self-supervised with a residual loss for linear PDEs (Poisson, Helmholtz, convection-diffusion-reaction).

Problem

Numerically solving linear PDEs efficiently and reliably; existing neural PDE solvers lack convergence/correctness guarantees and generalize poorly to unseen geometries, sizes, and PDE formulations.

Approach

Maps the multigrid V-cycle onto a fully convolutional U-Net-like network with fixed masked convolutional smoothers (provably convergent Jacobi-based iterator) and learnable bias-free convolutions in the coarse-grid correction, trained with a self-supervised absolute residual loss instead of supervised solution-matching loss.

Key results

• On large-scale Poisson, UGrid reaches 1e-4 relative residual 10-20x faster than AMGCL and 5-10x faster than AmgX (e.g., Bag 18.66 ms vs 202.95 ms AMGCL and 92.14 ms AmgX).
• Residual loss gives up to 2x speedup over legacy loss in ablation; vanilla U-Net diverged on all large-scale Poisson testcases.
• Neural baseline Hsieh et al. diverged on most testcases; UGrid generalizes to unseen geometries, and to XL and XXL problem sizes without retraining.

Measured effects reported in the paper

Surprising findings

• A vanilla U-Net trained with the same residual loss diverged on all Poisson testcases, indicating the architecture (not just the loss) is essential.
• Generalization to much larger problem sizes came at a price of reduced efficiency on small-scale problems (e.g., Cat and L-shape), where UGrid was sometimes slower than AMGCL.

Limitations

• Restricted to linear PDEs; the convergence theorem does not hold for nonlinear PDEs.
• No mathematical guarantee on convergence speed; not faster on small-scale problems.
• Timing comparisons exclude UGrid training time and compare against AMG solvers that rebuild hierarchies.

Experimental setup

Datasets: synthesized boundary masks, Dirichlet boundary values and right-hand sides (16000 pairs)

Benchmarks: 2.5D freeform surface Poisson problems (small, large, XL, XXL) Helmholtz freeform surface problems steady-state convection-diffusion-reaction freeform surface problems

Metrics: wall-clock time to relative residual <= 1e-4 final relative residual error

Baselines: AMGCL AmgX Hsieh et al. vanilla U-Net UGrid trained with legacy loss

Builds on: Hsieh et al. neural linear solver U-Net multigrid V-cycle Jacobi iteration

Closest papers by meaning

• Breakeven complexity: A new perspective on neural partial differential equation solvers (NeurIPS 2026 · similarity 0.89)
• INC: An Indirect Neural Corrector for Auto-Regressive Hybrid PDE Solvers (NeurIPS 2025 · similarity 0.88)
• Towards General Neural Surrogate Solvers with Specialized Neural Accelerators (ICML 2024 · similarity 0.87)
• Toward Efficient Kernel-Based Solvers for Nonlinear PDEs (ICML 2025 · similarity 0.87)
• Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers (ICML 2025 · similarity 0.87)
• Mechanistic PDE Networks for Discovery of Governing Equations (ICML 2025 · similarity 0.87)