Hybrid Iterative Solvers with Geometry-Aware Neural Preconditioners for Parametric PDEs

AI summary

Overview

Research area: Scientific machine learning, specifically neural operators for solving parametric partial differential equations (PDEs), combined with classical numerical linear algebra (iterative solvers and preconditioning).

Technical level: Advanced. The paper assumes familiarity with the finite element method, Krylov subspace solvers, preconditioning theory, and operator learning architectures.

Scope: The paper introduces Geo-DeepONet, a geometry-aware deep operator network that encodes unstructured mesh connectivity, and uses it as a preconditioner inside hybrid iterative solvers for parametric PDEs on domains never seen during training.

Published as arXiv:2512.14596v1 [cs.LG] on 16 Dec 2025, and listed for the International Journal for Numerical Methods in Engineering. The authors are Youngkyu Lee, Francesc Levrero Florencio, Jay Pathak, and George Em Karniadakis, with affiliations at Brown University, Ansys UK Ltd, and Ansys Inc (corresponding author: george_karniadakis@brown.edu). License: CC BY-NC-ND 4.0. Keywords listed: Iterative method, Neural operator, Hybridization, Unstructured domain, Graph neural network.

What This Paper Is About

Classical iterative solvers for parametric PDEs converge slowly when the computational domain has complex geometry, because the stiffness matrix of the resulting linear system becomes badly conditioned. A previous hybrid approach by the authors combined classical solvers with a neural operator, but it only worked well on the specific geometry it was trained on and under-performed on new geometries. This paper's goal is to make neural preconditioners generalize to arbitrary unstructured meshes without retraining, by encoding the domain's geometry directly into the operator network.

Key Contributions

1. Geo-DeepONet: a geometry-aware deep operator network that takes coordinates and the node-connectivity (adjacency) matrix from a finite element discretization as input, allowing accurate operator learning across arbitrary unstructured meshes without retraining.

2. A graph-trunk network: a modified DeepONet trunk network built from a Chebyshev spectral graph convolutional network that produces a geometry-feature vector from the coordinates and adjacency matrix, followed by a fully connected network that produces basis functions.

3. A projection and scaling mechanism for the branch network: a canonical projection Π_h that interpolates nodal features from the unstructured mesh onto a uniformly discretized Cartesian grid so a CNN branch network can be used, plus a scaling network inspired by squeeze-and-excitation networks that acts like an attention mechanism on the branch output.

4. A class of geometry-aware hybrid preconditioned iterative solvers: Geo-DeepONet is coupled with relaxation schemes (Jacobi, SOR) via the HINTS/multiplicative subspace correction framework and with Krylov methods (CG, FCG, GMRES, FGMRES) via the trunk-basis (TB) prolongation-operator approach.

Main Findings

• Ablation against standard DeepONet: On the 2D Poisson equation in both a unit square and an unstructured domain, Geo-DeepONet converged substantially better than a standard DeepONet even though both were trained on the identical dataset of arbitrary unstructured domains until reaching less than 8% relative L2 error, and even though the test domain was unseen by both. The authors attribute the gain to the scaling network and the graph-trunk network encoding arbitrary geometries.

• Relaxation-method speedups: The hybrid preconditioned relaxation methods (Jacobi and SOR) achieved a speedup of approximately 7x on the benchmark problems. The authors note these methods are commonly used as multigrid smoothers, so the speedup may transfer to faster multigrid convergence.

• Krylov-method speedups (Poisson): Hybrid preconditioning achieved 10x acceleration in the CG family and 20x in the GMRES family.

• Krylov-method speedups (linear elasticity): With Neumann boundary conditions imposed on arbitrary boundary regions, the GMRES families experienced an acceleration of approximately 30x.

• Real-world 3D meshes: On heat sink, nut, and plastic embellisher meshes, HP-CG achieved about 16x acceleration compared to vanilla CG, while the ILU preconditioner in PETSc sometimes failed to converge. Speedups were notably larger on symmetric shapes (the nut and the plastic embellisher), which the authors attribute to Geo-DeepONet having been trained on randomly generated surfaces of revolution that capture such symmetric characteristics.

• Larger problems: For problems with more than 50,000 unknowns, the hybrid preconditioned solvers outperformed baseline solvers and reduced iteration counts relative to smoothed aggregation algebraic multigrid (SA-AMG) in PETSc and BoomerAMG in HYPRE, including when the hybrid framework was combined with AMG levels.

• No boundary functions needed for training: Because DeepONet acts on the residual inside the hybrid framework, it only needs training on homogeneous boundary conditions, so boundary functions do not have to be sampled. This reduces memory requirements and simplifies the DeepONet structure.

Methodology in Plain English

The problem setup is a parametric PDE on a bounded domain Ω ⊂ R^d with d ∈ {2,3}, discretized with P1 conforming finite elements on a quasi-uniform triangulation, which yields a linear system K u = f with stiffness matrix K and load vector f.

The core idea for preconditioning is that classical iterative solvers reduce high-frequency error modes well but struggle with low-frequency modes, while neural networks show the opposite behavior due to spectral bias. Combining them in a hybrid preconditioner should capture the strengths of both. For relaxation methods, the hybrid scheme (HINTS) alternates n_r steps of a relaxation preconditioner M1^{-1} (e.g., Jacobi or SOR) with a prediction from the pre-trained neural operator M2^{-1}. For Krylov methods, which need a linear preconditioner, the authors extract basis functions from the trunk network output to build a prolongation operator P and restriction operator R = P*, defining the preconditioner M_TB^{-1} = P (R K P)^{-1} R, where R K P is a coarse-level problem.

The new network, Geo-DeepONet, changes how domain geometry enters a DeepONet. Its graph-trunk network takes coordinates x and the adjacency matrix A extracted from the finite element mesh, runs them through a Chebyshev spectral graph convolutional network to create a geometry-feature vector, and then through a fully connected network to produce basis functions. The branch network takes a projected input: because the unstructured mesh is not the same as a uniform Cartesian grid, the nodal features are interpolated onto a uniformly discretized Cartesian grid via Π_h so a CNN can be applied. A scaling network then splits the branch output into two paths, one through two fully connected layers and softmax and one through a single fully connected layer, and multiplies them element-wise. The full operator is G_geo(f, x, A) = ⟨S ∘ B(Π_h f), T̃(x, A)⟩.

Training minimizes the average point-wise relative squared error between Geo-DeepONet predictions and reference solutions over N_s samples, where each sample contains a domain, input function, coordinates, adjacency matrix, and reference solution. Domains are generated by sampling an arbitrary smooth curve from a one-dimensional Gaussian random field (mean 0.2, σ_g = 0.2, ℓ = 0.1, values ranging from 0.1 to 0.5); in 2D a closed curve centered at (0.5, 0.5) is constructed, and in 3D a surface of revolution is generated around the axis centered at (0.5, 0.5, 1). Meshes are generated by Gmsh with maximum mesh size h = 1/128 (2D) or h = 1/32 (3D).

Benchmarks are a Poisson equation with a variable coefficient and a linear elasticity equation, both in 2D and 3D unstructured domains. For Poisson, the diffusion coefficient κ ≥ 0.3 is sampled with Gaussian random fields with mean E[κ(x)] = 1.0 and covariance σ_g² exp(-||x - x'||²/(2ℓ²)) with σ_g = 0.3 and ℓ = 0.1, and the forcing term uses the same covariance with σ_g = 1.0 and ℓ = 0.1. For elasticity, Poisson's ratio ν = 0.3 and Young's modulus E(x) is sampled the same way as κ, with the body force sampled as the forcing term.

Implementation: 50,000 training samples and 10,000 test samples for all benchmark problems; reference solutions computed with the sparse direct solver MUMPS; training samples generated with FEniCSx; iterative solvers implemented through PETSc and petsc4py; Geo-DeepONets implemented in PyTorch and trained with AdamW, batch size 100, learning rate 10^-4; training stopped when validation average relative L2 error fell below 8% or the training loss did not improve for 100 consecutive epochs. Solvers terminate when ||r^(i)||_2 ≤ 10^-12 or ||r^(i)||_2 / ||r^(0)||_2 ≤ 10^-8. Experiments ran on nodes with an AMD EPYC 9554 64-Core Processor (256GB) and an NVIDIA L40S GPU (48GB) at Brown University's Center for Computation and Visualization. Table 1 gives the specific branch, scaling, and graph-trunk architectures for Poisson 2D, Poisson 3D, Elasticity 2D, and Elasticity 3D.

For the real-world mesh comparison, the authors reused the parameters of the Geo-DeepONet trained on surfaces of revolution, used the Chebyshev semi-iterative method for the relaxation steps, and used CG as the backbone solver, comparing against PETSc's ILU and SOR preconditioners and the preconditioner in ANSYS Mechanical.

Why This Matters

Impact on research. The work directly addresses the main limitation of prior hybrid preconditioning strategies, which required the neural operator to be trained on the same domain as the target problem. By encoding mesh connectivity through a graph-trunk network and projecting unstructured features onto a structured grid, it offers a route to operator learning that generalizes across arbitrary unstructured geometries without retraining, and it positions neural preconditioners as components usable alongside established multigrid machinery such as SA-AMG and BoomerAMG.

Real-world applications.

• 3D linear elasticity simulations on real engineering meshes, specifically the heat sink, nut, and plastic embellisher meshes studied in the paper.
• Structural analysis with mixed boundary conditions, since the elasticity benchmark imposes Neumann conditions on arbitrary boundary regions while Dirichlet conditions cover the rest.
• Parametric design studies where the same PDE must be solved repeatedly over many different geometries, which is where reusing a trained operator without retraining pays off.
• Large-scale simulations with more than 50,000 unknowns, where algebraic multigrid is typically required for scalability.

Industry relevance. Two of the four authors are affiliated with Ansys, and the preconditioner in ANSYS Mechanical is used as a comparison baseline, indicating direct industrial interest in replacing or augmenting solver preconditioners in commercial simulation software.

Future Directions

• Training Geo-DeepONet on a more diverse dataset consisting of various CAD models to enhance robustness and generalizability, which the authors state would lead to more consistent performance across a wider range of geometries such as the heat sink.
• Further exploiting HP-relaxation methods as smoothers inside multigrid, which the authors suggest could lead to faster multigrid convergence, and continuing to combine the hybrid framework with AMG restriction and prolongation operators at each multigrid level.
• Reducing the remaining geometry-specific gaps: the smaller speedup on the heat sink relative to symmetric shapes suggests open work on geometries that are not well captured by surfaces of revolution.
• Determining how well the approach extends beyond the two benchmark PDE families and the specific 2D/3D domain-generation process (closed curves in 2D, surfaces of revolution in 3D) used here, since the paper does not report results for other domain-generation schemes.

Target Audience

Researchers and practitioners in scientific machine learning and computational mechanics who already understand finite element discretization, iterative solvers, and preconditioning. It is most useful for those working on neural operators for unstructured meshes, hybrid physics-based and learned solvers, or solver acceleration in commercial simulation settings. Readers without a background in numerical linear algebra or operator learning will find the methodology sections demanding.

Authors’ abstract

The convergence behavior of classical iterative solvers for parametric partial differential equations (PDEs) is often highly sensitive to the domain and specific discretization of PDEs. Previously, we introduced hybrid solvers by combining the classical solvers with neural operators for a specific geometry 1, but they tend to under-perform in geometries not encountered during training. To address this challenge, we introduce Geo-DeepONet, a geometry-aware deep operator network that incorporates domain information extracted from finite element discretizations. Geo-DeepONet enables accurate operator learning across arbitrary unstructured meshes without requiring retraining. Building on this, we develop a class of geometry-aware hybrid preconditioned iterative solvers by coupling Geo-DeepONet with traditional methods such as relaxation schemes and Krylov subspace algorithms. Through numerical experiments on parametric PDEs posed over diverse unstructured domains, we demonstrate the enhanced robustness and efficiency of the proposed hybrid solvers for multiple real-world applications.