Neural Network Preconditioning: A Case Study for the Solution of the Parametric Helmholtz Equation

The gist: A convolutional neural network based on the U-Net architecture is trained in a self-supervised manner to approximate the inverse of the discretized Helmholtz operator, and this trained network is used as a nonlinear preconditioner within the flexible GMRES (FGMRES) algorithm to significantly improve convergence for solving linear systems arising from the discretization of the two-dimensional parametric Helmholtz equation.

How it works

The core approach involves hybridizing machine learning techniques with classical numerical algebra methods, specifically using a trained neural network (NN) as a nonlinear preconditioner for the flexible GMRES (FGMRES) subspace iteration method to solve linear systems arising from the discretization of the two-dimensional parametric Helmholtz equation The approach is based on training an NN in a self-supervised manner using a loss function involving the residual norm of the linear system to approximate the inverse of the discretized Helmholtz operator This trained network is then employed as a nonlinear preconditioner within FGMRES to speed up convergence.

The neural network components are designed based on the input fields of the discretized source and velocity fields The NN is designed using a U-Net architecture that takes the discretized source and velocity fields as inputs, recasting them as image-like tensors to be suitable for the convolutional nature of the U-Net. The network's functional representation is a nonlinear operator Nθ: C n × R n → C n, defined by Nθ(b, c) = Nθ(c)(b).

Training and Loss Function

The training process aims at finding the parameters θ∗ ∈ Θ of the neural network that minimize a certain cost functional J: Θ → R The cost function takes the form J(θ) = Eb,c[Lθ(b, c)]. The loss function is defined by Lθ(b, c) = ∥b − A(c)Nθ(b, c)∥2 / 2∥b∥2 / 2. This loss function corresponds to the backward error with respect to b of x˜ as an approximate solution to the linear system A(c)x = b. The network is trained so that Nθ(c) ≈ A−1(c) for any c. The training process is summarized in Figure 3.

The network was trained using training and validation data generated with the following parameters involved in the discretization of the Helmholtz equation The distribution of (b, c) is defined component-wise and independently for b and c such that the entries of b are independent and normally distributed with zero mean and unit variance (we denote b iid∼ N (0, 1)), the ones of c are independent and uniformly distributed between 1 and 2 (i.e., c iid∼ U([1, 2])).

Numerical Experiments

Numerical experiments evaluate the quality and robustness of the neural network, first as a standalone solver as a surrogate model, then as a nonlinear preconditioner The quality of an approximate solution x˜ to the linear system A(c)x = b will be evaluated through the normwise backward error ηb(˜x):= ∥b − A(c)˜x∥2 / ∥b∥2.

The trained networks are tested for generalization capabilities with respect to the parameter c and with respect to the domain size For the set of trained parameters, the numerical quality of the predicted solutions by the surrogate model is extremely poor, with a backward error ηb(ˆxθ) that varies between 0.1 and 0.8. The results presented in the last two subsections indicate that the network cannot be used as a stand-alone solver, as it lacks the robustness needed for accurate solutions.

Preconditioning Effectiveness

The effectiveness of the neural preconditioner is reported by comparing its performance with respect to an algebraic preconditioner in Section 4.2.1. The NN consistently outperforms all the algebraic preconditioners, including the best-performing one obtained with ϵ = 0.005. The NN is the same for all linear systems—an important feature that makes Nθ particularly suitable as a preconditioner when one needs to solve a sequence of (possibly massively many) systems with different left- and/or right-hand sides.

In Figure 7, the convergence history of ηb for the solution of 100 linear systems derived from the discretization of the Helmholtz problem with a random velocity field c iid∼ U([1, 2]), using FGMRES without restart is displayed. All convergence histories using Nθ (in orange) are closely aligned and convergence to 10−10 is achieved in fewer than 20 iterations. The main advantage of the nonlinear NN preconditioner is that it is trained once and for all, and can subsequently be used to solve arbitrarily many problems.

Robustness and Generalization

The NN-FGMRES solver is highly robust to the PDE parameters c and b, not only in terms of ranges of values, but also in terms of structure. The convergence history for the pair (c0 = 1.25, c1 = 1.75), which fully lies in the training range [1, 2], shows fast convergence (ηb reaches 10−10 in fewer than 30 iterations for both trainings). Furthermore, the NN-FGMRES is still able to somewhat accelerate the convergence of the subspace solver. The results suggest that the NN-FGMRES is highly robust to the PDE parameters c and b, not only in terms of ranges of values, but also in terms of structure.

The methodology presented here builds on this growing body of work by demonstrating that approximate learned inverses can be effectively integrated into established numerical schemes. The intersection of numerical linear algebra and machine learning continues to pave the way for solving complex, high-dimensional scientific problems.

Conclusion

This study illustrates the potential of using neural networks as nonlinear preconditioners to solve parametric Helmholtz equations using subspace methods. This hybrid approach combines the generalisation capabilities and computational speed of neural networks with the numerical reliability of classical iterative methods. The generic framework can be extended to other classes of linear systems beyond the Helmholtz equation, offering a promising approach to machine learning in scientific computing. Several promising avenues for future research have emerged.

--- Page 1 ---

Neural network preconditioning: a case study for the solution of the parametric Helmholtz equation Luc Giraud Carola Kruse Paul Mycek Maksym Shpakovych and Yanfei Xiang Inria Concace joint team between Airbus CR&T Cerfacs and Inria Current affiliation Opticalp France Current affiliation IRMA Université de Strasbourg MACARON Inria France Corresponding author mycek@cerfacs.fr 1 arXiv:2610.16542v1 [math.

1. No generic, efficient preconditioner currently exists for linear systems arising from the discretization of parametric PDEs like Helmholtz equations; this work introduces a hybrid numerical approach where a trained neural network is used as a nonlinear preconditioner within the flexible GMRES (FGMRES) algorithm to significantly improve the convergence of FGMRES.

2. The improved AI system can solve large linear systems arising from 2D parametric Helmholtz equation discretization by using the trained neural network as a preconditioner, demonstrating robust performance and generalization capabilities with respect to variations in the velocity field and the domain size.

3. The system can accelerate convergence of subspace methods like FGMRES by approximating the inverse of the discretized Helmholtz operator, as shown in Equation (4): In this paper we use this equality with Mi = Nθ, which is the trained neural network.

4. The system exhibits robustness when applied to problems outside its training scope, as the learned NN preconditioner exhibits remarkable robustness when applied to larger problems (in terms of domain size and number of degrees of freedom), or to problems with parameter values or structure outside the training scope.

This work presents a hybrid numerical approach for solving linear systems arising from the discretization of the two-dimensional parametric Helmholtz equation. A convolutional neural network based on the U-Net architecture is trained in an self-supervised manner (i.e., without explicitly generating reference solutions) to approximate the inverse of the discretized Helmholtz operator, using a loss function involving the residual norm of the linear system. The trained network is used as a nonlinear preconditioner within the flexible GMRES (FGMRES) algorithm. Numerical experiments show that while the neural network is not accurate enough to act as a standalone solver, it significantly improves the convergence of FGMRES when employed as a preconditioner. The neural preconditioner demonstrates robust performance and generalization capabilities with respect to variations in the velocity field and the domain size. Comparisons with algebraic preconditioners, based on LU factorizations of approximations of the original system matrices, indicate the superior efficiency of the neural approach under equivalent conditions. We believe that the proposed method is not tied to a specific neural architecture and can be extended to other parametric PDEs.