Multi-Level GNN Preconditioner for Solving Large Scale Problems

TL;DR

This work tackles solving large-scale Poisson problems by bridging HPC-established domain decomposition with graph neural networks. It introduces DDM-GNN, a two-level preconditioner that uses a coarse correction plus a GNN-based solver for local subproblems, enabling GPU-accelerated, scalable Krylov solves with controllable accuracy. The DSS_theta model provides a discretization-agnostic, iteratively refined local solution, trained via a physics-informed residual loss, and its integration into CG/PCG yields robust convergence across varying mesh sizes and geometries. While showing promising scalability and competitive iteration counts, the study also highlights current bottlenecks in HPC-oriented GNN inference engines, pointing to future work on efficient GNN runtimes and industrial CFD integration.

Abstract

Large-scale numerical simulations often come at the expense of daunting computations. High-Performance Computing has enhanced the process, but adapting legacy codes to leverage parallel GPU computations remains challenging. Meanwhile, Machine Learning models can harness GPU computations effectively but often struggle with generalization and accuracy. Graph Neural Networks (GNNs), in particular, are great for learning from unstructured data like meshes but are often limited to small-scale problems. Moreover, the capabilities of the trained model usually restrict the accuracy of the data-driven solution. To benefit from both worlds, this paper introduces a novel preconditioner integrating a GNN model within a multi-level Domain Decomposition framework. The proposed GNN-based preconditioner is used to enhance the efficiency of a Krylov method, resulting in a hybrid solver that can converge with any desired level of accuracy. The efficiency of the Krylov method greatly benefits from the GNN preconditioner, which is adaptable to meshes of any size and shape, is executed on GPUs, and features a multi-level approach to enforce the scalability of the entire process. Several experiments are conducted to validate the numerical behavior of the hybrid solver, and an in-depth analysis of its performance is proposed to assess its competitiveness against a C++ legacy solver.

Figures (6)

• Figure 1: Illustration of the proposed hybrid solver: A global Poisson problem (left) is solved using a Preconditioned Conjugate Gradient (PCG) method, enhanced by the DDM-GNN preconditioner. DDM-GNN combines Domain Decomposition with a GNN-based solver, harnessing parallel computations on GPUs.
• Figure 2: Illustration of DDM-GNN: It takes a global residual vector $\mathbf{r}$ defined on $\Omega_h$ as input and produces a global correction vector $\mathbf{z}$. This is done by combining the results of solving the coarse problem through LU decomposition (bottom) with those of solving all local problems using $\text{DSS}_\theta$ (top). Note, we use $\Omega_{h,i}$ instead of $R_iAR_i^T$ for reasons detailed in \ref{['subsec:archi']}.
• Figure 3: $\text{DSS}_\theta$ architecture: The model is initialized with a null latent state $H^0$ (in blue). An iterative process then propagates information for $\bar{k}$ iterations using $M^{\bar{k}}_\theta$ distinct blocks of MPNNs. At each iteration $k$, a Decoder $D_\theta^k$ translates the latent state $H^k$ into a physical state $\widehat{\mathbf{r}}^k$, with an intermediate loss computed. The green squares denote trainable functions, and their weights are learned during training. The last $\widehat{\mathbf{r}}^{\bar{k}}$ (in red) represents the model output.
• Figure 4: Example of (\ref{['subfig:dataset_graph']}): one generated global domain with $7420$ nodes, and (\ref{['subfig:dataset_partitioning']}): its partitioning into $8$ sub-meshes. Each colored sub-mesh corresponds to a graph sample of the dataset.
• Figure 5: (\ref{['subfig:partitioning']}): partitioning of the mesh into $234$ sub-meshes and (\ref{['subfig:residual']}): the evolution of the relative residual error for the different methods until it reaches an order of $10^{-9}$.