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

Authors: Xi Han, Fei Hou, Hong Qin

ICML 2024 | Conference PDF | Archive PDF | Plain Text | LLM Run Details

Reproducibility Variable Result LLM Response
Research Type Experimental Extensive experiments and comprehensive evaluations have verified all of our claimed advantages, and confirmed that our proposed method outperforms the SOTA. For each of the three PDEs mentioned in Sec. 4.1: (i) Poisson problem, (ii) inhomogeneous Helmholtz problem with varying wave numbers, and (iii) inhomogeneous steady-state diffusion-convection-reaction problem, we train one UGrid model specialized for its formulation. We apply our model and the baselines to the task of 2.5D freeform surface modeling. These surfaces are modeled by the three types of PDEs as 2D height fields, with non-trivial geometry/topology. Each surface is discretized into: (1) Small-scale problem: A linear system of size 66,049 x 66,049; (2) Large-scale problem: A linear system of size 1,050,625 x 1,050,625; (3) XL-scale problem: A linear system of size 4,198,401 x 4,198,401; and (4) XXL-scale problem: A linear system of size 16,785,409 x 16,785,409. UGrid is trained on the large scale only. Other problem sizes are designed to evaluate UGrid s generalization power and scalability. In addition, we have conducted an ablation study on the residual loss metric (v.s. legacy loss) as well as the UGrid architecture itself (v.s. vanilla U-Net).
Researcher Affiliation Academia 1Department of Computer Science, Stony Brook University (SUNY), Stony Brook, NY 11794, USA. 2Key Laboratory of System Software (Chinese Academy of Sciences) and State Key Laboratory of Computer Science, Institute of Software, Chinese Academy of Sciences, Beijing, 100190, China. 3University of Chinese Academy of Sciences, Beijing, 100049, China.
Pseudocode No The paper describes the UGrid Iteration as a sequence of operations and illustrates the UGrid Submodule with a figure, but it does not provide formal pseudocode or an algorithm block labeled as such.
Open Source Code Yes Our code is available as open-source at https://github.com/AXIHIXA/UGrid.
Open Datasets No Before training, we synthesized a dataset with 16000 (M, b, f) pairs. For Helmholtz and diffusion problems, we further random-sample their unique coefficient fields (more details available in Sec. A.8 and Sec. A.9.) To examine the generalization power of UGrid, the geometries of boundary conditions in our training data are limited to Donuts-like shapes as shown in Fig. 3 (h). Moreover, all training data are restricted to zero f-fields only, i.e., f extasciitilde{} 0.
Dataset Splits No The paper mentions a 'training phase' and 'testcases' but does not specify explicit training, validation, and test dataset splits (e.g., 80/10/10 percentages or specific sample counts for each split).
Hardware Specification Yes We train our model and perform all experiments on a personal computer with 64GB RAM, AMD Ryzen 9 3950x 16-core processor, and NVIDIA Ge Force RTX 2080 Ti GPU.
Software Dependencies No The paper mentions 'Py Torch' but does not specify a version number. While it refers to other software like 'AMGCL (Demidov, 2019)' and 'NVIDIA Amg X (NVIDIA Developer, 2022)' as baselines, it does not provide version numbers for these or for any other ancillary software used in their own implementation.
Experiment Setup Yes We train our model for 300 epochs with the Adam optimizer. The learning rate is initially set to 0.001, and decays by 0.1 for every 50 epochs.