July 18, 2026
This work develops a neural operator preconditioned subspace method for sequences of linear systems arising from the discretization of the two-dimensional Helmholtz equation in transcranial ultrasound applications. The problem involves strongly heterogeneous, patient-dependent velocity fields that induce severe wave distortion and pose significant challenges for standard iterative solvers. Building on neural network preconditioning framework of Giraud et al. (HAL RR-9593, 2025) and the idealized skull dataset used for the learned optimizer of Stanziola et al. (JCP 441, 2021), neural operator preconditioners are trained on six mixed velocity-source datasets combining randomized source configurations and idealized skull-based velocity fields with random noise. The proposed mixed-dataset strategy aims to improve both computational efficiency and generalization across varying configurations. The neural operator is trained on a coarse grid using a physics-informed loss based on the relative residual of the discrete Helmholtz equation and is incorporated as a nonlinear preconditioner within flexible GMRES (FGMRES). Numerical experiments demonstrate that the resulting hybrid method efficiently solves practical transcranial ultrasound problems on grids 64 times larger than those used during training, whereas both classical GMRES and the learned optimizer fail to converge within comparable computational budgets. Moreover, the proposed method achieves arbitrary solution accuracies and exhibits strong out-of-distribution generalization across diverse source and velocity configurations. These results demonstrate that appropriately designed mixed training datasets substantially improve the effectiveness of learned neural operator preconditioners. More broadly, this work highlights the importance of dataset design in scientific machine learning and provides a practical framework for integrating matrix-free neural operator preconditioning with Krylov subspace methods for solving practical large-scale Helmholtz problems.
Keywords: Large-scale Helmholtz equation; Transcranial ultrasound; Neural operator preconditioning; Scientific machine learning; Flexible GMRES; Krylov subspace methods
This work focuses on developing neural operator preconditioners trained on mixed datasets to accelerate the solution of sequences of linear systems arising from the discretization of the two-dimensional parametric Helmholtz equation in transcranial ultrasound therapy [1]. In this setting, the Helmholtz equation incorporates a heterogeneous velocity field derived from human head CT data, resulting in complex wave interactions within the adult human skull. Since skull geometries and acoustic properties vary significantly across patients, the corresponding Helmholtz operators can exhibit considerable variability, making the development of robust and efficient hybrid numerical solvers, potentially enhanced by machine learning, essential for practical clinical applications. This challenging transcranial ultrasound problem was previously investigated by Stanziola et al. [2], who proposed a learned optimizer based on a recurrent neural network embedded within a Richardson iteration. For the representative transcranial ultrasound example considered in their study, the learned optimizer outperformed classical GMRES [3] in computational efficiency while exhibiting encouraging generalization capabilities. However, its attainable accuracy remained limited, and it failed to converge on the additional transcranial ultrasound examples considered in the present work (see Table ¿tbl:tab95results95test95dirac195idea95skull? and Figure 9 in Section [subsec:skull95exp]). Rather than replacing classical iterative solvers, our approach in [4] employs neural networks as nonlinear preconditioners within flexible GMRES (FGMRES) [5], combining the approximation and generalization capabilities of scientific machine learning [6], particularly neural operator learning [7], with the robustness and reliability of Krylov subspace methods. The resulting hybrid methods converge reliably to arbitrary prescribed accuracies without requiring repeated network training and exhibit strong generalization across a broad range of Helmholtz problems. Nevertheless, the previous framework was trained exclusively on randomly generated datasets, and its computational efficiency deteriorates when applied to the realistic challenging application-oriented transcranial ultrasound simulations. This observation motivates the present work. We investigate whether carefully designed mixed training datasets can improve the effectiveness of neural operator preconditioning while preserving its robustness and generalization capabilities. Specifically, we train the same neural operator architecture using six different mixed datasets for the Helmholtz equation. These datasets combine application-oriented samples that mimic realistic transcranial ultrasound configurations with randomized source and velocity fields to maintain sufficient diversity during training. Their relative performance is systematically evaluated to identify dataset compositions that provide a favorable balance between specialization for the target application and generalization to previously unseen configurations. The present work can therefore be viewed as a practical extension of our previous neural operator preconditioning framework [4]. Whereas our earlier study relied exclusively on randomly generated training data, the central contribution of the present work is the design and systematic investigation of mixed training datasets that improve efficiency for challenging transcranial ultrasound simulations while preserving strong out-of-distribution performance on other Helmholtz problems. More broadly, the present work demonstrates that careful dataset design is a key ingredient for translating scientific machine learning methods from idealized benchmarks to challenging real-world applications.
In parallel, recent advances in scientific machine learning, particularly in neural operator learning and diffusion models, have also highlighted the importance of exposing models to diverse training data distributions to improve robustness and out-of-distribution generalization across varying physical parameters, geometries, and discretizations [8]–[15]. Some related developments in neural network preconditioning for classical numerical methods are discussed in [16]–[20]. The field of scientific machine learning is evolving rapidly, and the references cited above are intended to highlight representative developments closely related to the present work rather than to provide a comprehensive survey.
The remainder of this paper is organized as follows. Section [sec:backgroud] introduces the governing Helmholtz equation arising in transcranial ultrasound therapy together with its numerical solution using nonlinear preconditioned subspace method. We present the mathematical formulation, including the boundary conditions, discretization, and iterative solution strategies, and briefly review recent developments in neural network solvers and neural operator preconditioners for the Helmholtz equation, highlighting the limitations that motivate the present work. Section 3 presents the proposed convolutional neural network (CNN)-based neural operator preconditioner and the construction of the six mixed training datasets. The resulting preconditioner is designed to accelerate FGMRES for solving sequences of heterogeneous Helmholtz equations arising in transcranial ultrasound simulations while maintaining strong generalization to other Helmholtz applications. Section [sec:experiments] evaluates the proposed approach on a variety of transcranial ultrasound examples with previously unseen physical parameters. We systematically compare the different mixed training datasets and demonstrate that the resulting preconditioners substantially accelerate FGMRES while exhibiting strong out-of-distribution generalization across a broad range of Helmholtz problems. Finally, Section [sec:conclusion] summarizes the main findings and discusses their implications for scientific machine learning, neural operator preconditioning, and large-scale PDE simulations.
Key notations used throughout this paper are summarized below. Unless otherwise specified, \(\|\cdot\|\) denotes the Euclidean norm for both vectors and matrices. The superscript \(^H\)
denotes the conjugate transpose. The symbols \(\mathbb{R}\) and \(\mathbb{C}\) denote the sets of real and complex numbers, respectively. For convenience, some Python-style
notation is adopted in the presentation of the algorithm. To improve readability, we use the following notation throughout the paper. Scalars are denoted by lowercase letters. Bold lowercase letters, e.g., \(\mathbf{x}\),
denote vectors (or columns of a matrix), while uppercase letters, e.g., \(A\), denote matrices. Functions, operators, and mappings are represented by calligraphic letters, e.g., \({\mathcal{F}}.\) Unless otherwise stated, the subscript \(_j\) denotes the quantity obtained at iteration \(j\), while the positive integer subscript \(_m\) denotes the maximum number of iterations within a Krylov cycle. The inner product between two vectors \(\mathbf{x}\) and \(\mathbf{y}\) is denoted by \(\langle \mathbf{x}, \mathbf{y}\rangle\), and the linear span of the vectors \(\mathbf{x}_1,\ldots,\mathbf{x}_j\) is denoted by \(\operatorname{span}\{\mathbf{x}_1,\ldots,\mathbf{x}_j\}.\)
We consider the two-dimensional Helmholtz equation, subject to the so-called Sommerfeld radiation condition, \[\label{eq:helmholtz} \begin{cases} \nabla^2 u + k^2 u = f \;in\mathbb{R}^2,\\[\medskipamount] \lim_{\|\mathbf{x}\|_2 \to \infty} \|\mathbf{x}\|_2^{1/2} \left( \dfrac{\partial u}{\partial \|\mathbf{x}\|_2}(\mathbf{x}) - \jmath k(\mathbf{x}) u(\mathbf{x}) \right) = 0, \end{cases}\tag{1}\] where \(\jmath \in \mathbb{C}\) denotes the imaginary unit such that \(\jmath^2=-1\), and \(u \colon \mathbb{R}^2 \to \mathbb{C}\), \(f \colon \mathbb{R}^2 \to \mathbb{C}\), \(k \colon \mathbb{R}^2 \to \mathbb{R}^+\) are scalar fields, referred to as the solution field, the source field, and the wavenumber field, respectively. The latter is defined by \(k \colon \mathbf{x}\mapsto \omega / c(\mathbf{x})\), where \(\omega \in \mathbb{R}^+\) is the angular frequency of the source, and \(c \colon \mathbb{R}^2 \to \mathbb{R}^+\) is the velocity (speed of sound) field. In what follows, we restrict ourselves to the case \(\omega = 1\). Numerically, the unbounded domain \(\mathbb{R}^2\) is truncated to a computational, square domain \(\Omega :=[-L, L]^2 \subset \mathbb{R}^2\) and the derivatives in Equation 1 are discretized using Fourier differentiation [2], [21], [22] with \(N\) points in each direction. Consequently, there is a total of \(n :=N^2\) degrees of freedom. The boundary conditions are (approximately) enforced using perfectly matched layers (PMLs) [23] in the outer region \(\Omega_{\textrm{pml}}:=\Omega \setminus \Omega_{\textrm{inner}}\), where \(\Omega_{\textrm{inner}}:=[-L_{\textrm{inner}}, L_{\textrm{inner}}]^2\), with \(L_{\textrm{inner}}:=L- \ell\) and \(\ell< L\), as illustrated in Figure 1.
Within the PML framework, the absorbing layer for truncating the domain can be characterized by a damping function \[\label{eq:pml} \sigma(\mathbf{x}) = \begin{cases} 0, \;for\mathbf{x}\in \Omega_{\textrm{inner}}\subset \Omega,\\[\medskipamount] \sigma_{\text{max}} \dfrac{(\mathbf{x}- L_{\textrm{inner}})^2}{l^2}, \;for\mathbf{x}\in \Omega_{\textrm{pml}}\subset \Omega, \end{cases}\tag{2}\] where \(\sigma_{\text{max}} \in \mathbb{R}^+\) is a prescribed parameter denoting the maximum damping coefficient. With this construction, the solution remains unchanged in the inner domain \(\Omega_{\textrm{inner}}\subset \Omega,\) while it decays exponentially in the PML region \(\Omega_{\textrm{pml}}\subset \Omega.\) As a result, outgoing waves are effectively absorbed without producing spurious reflections at the interface between the physical domain and the absorbing layer.
The discretization of the Helmholtz problem 1 with varying physical parameters (i.e., varying source filed \(f\) and velocity filed \(c\)) yields a groups of linear systems with varying left and right hand sides \[~\label{eq:discrePDEs} A(\mathbf{c})^{(\ell)} \mathbf{u}= \mathbf{b}^{(\ell)}, \text{ with } A := \nabla^2 + k^2 \in \mathbb{C}^{n\times n}, \;\ell =1,2,....L,\tag{3}\] where, associated with the \(\ell\)-th family from each discrete Helmholtz equation involving the discrete velocity field \(\mathbf{c}\in \mathbb{R}^{n}\) and the associated discretized operator (matrix) \(A(\mathbf{c}) \in \mathbb{C}^{n\times n}\), as well as the discretized source term \(\mathbf{b}\in \mathbb{C}^{n}\) and solution \(\mathbf{u}\in \mathbb{C}^{n}.\)
The solution of large linear systems of equations is commonly addressed using iterative methods. These approaches are attractive due to their modest memory requirements and their ability to stop the iterations when the quality of the solution is similar to the possible uncertainty in the matrix entries or right-hand sides, which may arise from discretization errors in PDE or from uncertainties in the input data. In this context, backward error analysis [24], [25] provides powerful techniques for designing meaningful stopping criteria. Among iterative techniques, subspace methods, such as Krylov subspace methods, have become standard tools over the past decades. For unsymmetric problems, the best known method is GMRES [3], which minimizes the residual norm over an increasing Krylov subspace. This technique relies on the so-called Arnoldi process, which incrementally builds an orthonormal basis of the nested subspaces. At iteration \(k\), the Arnoldi algorithm computes a set of orthonormal vectors \(V_k = [\mathbf{v}_1, \dots, \mathbf{v}_k ]\) that satisfy the so-called Arnoldi relation, which is written in matrix form: \[A V_k = V_{k+1}\bar{H}_k \quad \rm withV_{k+1}^H V_{k+1} = I_{k+1}\] where \(\bar{H}_k \in \mathbb{C}^{(k+1)\times k}\) is upper Hessenberg. Assume the approximated solution at iteration \(k\) is denoted as \(\mathbf{x}_k,\) the minimum residual norm iterate \(\mathbf{x}_k \in \mathop{\mathrm{span}}\{\mathbf{v}_1, \dots, \mathbf{v}_k \}\), that is \(\mathbf{x}_k = V_k \mathbf{y}_k\) with \(\mathbf{y}_k =\mathop{\mathrm{arg\,min}}_{ \mathbf{y}\in \mathbb{C}^k} (\| \bar{H}_k \mathbf{y}- \beta \mathbf{e}_1 \| )\) where \(\beta =\| \mathbf{b}\|\). In practice, a preconditioner \(M\) is used to speed up the convergence of GMRES, where \(M\) is expected to approximate \(A^{-1}\) somehow. For GMRES it is recommended to use the preconditioner on the right, i.e., GMRES solves \(AM \mathbf{t}= \mathbf{b}\), with \(M \mathbf{t}=\mathbf{x}\). Thus, at each iteration, GMRES still minimizes the residual norm of the residual associated with the original linear system. The Arnoldi relation is \[A M V_k = V_{k+1}\bar{H}_k \quad \rm withV_{k+1}^H V_{k+1} = I_{k+1}.\] In [5], Saad introduced the idea of a flexible preconditioner, which allows to have a different preconditioning matrix \(M_i\) at each iteration, so that the generalized Arnoldi relation becomes \[\label{eq:generalizedArnoldi} A Z_k = V_{k+1}\bar{H}_k \quad \rm withV_{k+1}^H V_{k+1} = I_{k+1},\tag{4}\] where \(Z_k = [\mathbf{z}_1, \dots, \mathbf{z}_k]\) with \(\mathbf{z}_i = M_i \mathbf{v}_i~(i=1,\dots,k)\). This idea can easily be extended to the situation where the preconditioner is no longer linear, that is, \(\mathbf{z}_i = {\cal M}_i(\mathbf{v}_i)\), where each \({\cal M}_i\) may be a nonlinear operator. In the latter case, Equation 4 still holds and uniquely defines each iterate as long as \(Z_k\) remains full rank. In this work, we adopt \({\cal M}_i() = \mathcal{N}_\theta()\), where \(\mathcal{N}_\theta\) denotes the trained neural network inference (see Section 3). The nonlinearity has two origins; first, the activation functions in the neurons are nonlinear, and second, the preconditioner is computed in 32-bit arithmetic while the rest of the computation is done in 64-bit arithmetic. This truncation due to casting 64 bits to 32 bits is also nonlinear. A sketch of the CNN-preconditioned FGMRES algorithm is provided in Algorithm [alg:nn-fgmres].
\(\mathbf{r}_0 = \mathbf{b}- A \mathbf{x}_0\), \(\beta = \|\mathbf{r}_0\|\), \(\mathbf{v}_1 = \mathbf{r}_0 / \beta\) \(\bar{H}_m = [h_{i,j}]_{1\leq j \leq m, 1 \leq i \leq j+1}\) \(Z_m = [\mathbf{z}_1, \dots, \mathbf{z}_m]\) \(\mathbf{y}_m = \mathop{\mathrm{arg\,min}}_{\mathbf{y}\in \mathbb{C}^m} \| \bar{H}_m \mathbf{y}- \beta \mathbf{e}_1 \|\) \(\mathbf{x}_m = \mathbf{x}_0 + Z_m \mathbf{y}_m\)
The present work is motivated by a sequence of previous studies on machine learning approaches for solving Helmholtz equations arising in transcranial ultrasound simulations. An early contribution in this direction is the learned optimizer proposed by Stanziola et al. [2]. Their approach employed a modified U-Net architecture [26] within a Richardson iteration framework, in which a recurrent neural network generated a corrective update at each iteration. The model was trained using a fixed source configuration and velocity fields derived from idealized skull geometries. While this work demonstrated the potential of neural-network-based solvers for Helmholtz problems, it also revealed several important limitations. In particular, the attainable accuracy was limited (approximately \(10^{-4}\)), and the method successfully solved only one practical transcranial ultrasound example while failing on the remaining cases considered in the present work (see Table ¿tbl:tab95results95test95dirac195idea95skull? and Figure 9 in Section [subsec:skull95exp]). Moreover, the training procedure exhibited instabilities, including producing divergent training losses and sensitivity to initialization (see [21]), which occasionally required training attempts to obtain a satisfactory model. To address these issues, an optimal step-size strategy was introduced in [21], leading to a minimum-residual Richardson iteration. By explicitly minimizing the residual at each iteration, this formulation substantially improved both the attainable accuracy and the stability of the training procedure. Beyond using neural networks as correction operators, this work also introduced the idea of employing them as nonlinear preconditioners within two types of Krylov subspace methods. Subsequent studies [4], [22] extended this neural operator preconditioning framework to more general Helmholtz problems and other classes of PDEs. Unlike purely data-driven solvers, these hybrid approaches combine the approximation capability and generalization of neural operators with the robustness of Krylov subspace methods, enabling convergence to arbitrary prescribed tolerances while retaining strong performance across a broad range of benchmark problems. Although these methods demonstrated excellent robustness and generalization across diverse Helmholtz problems, their computational efficiency remained limited for challenging application-specific transcranial ultrasound simulations. The objective of the present work is therefore to investigate whether application-oriented mixed training datasets can improve the efficiency of neural operator preconditioning while preserving its robustness and out-of-distribution generalization.
Assume the CNN-based preconditioning operator is denoted as \(\mathcal{N}_\theta\) with trainable parameters \({\theta}\) (i.e., the trainable weights and biases of the CNNs). We focus on learning an operator \(\mathcal{N}_\theta\) \[~\label{eq:operator} \mathcal{N}_\theta: {\mathcal{P}}\longrightarrow {\mathcal{U}},\tag{5}\] with coefficient input \(\left\{\mathbf{b}^{(\ell)}, \text{**kwarg}\right\} \in {\mathcal{P}}, \text{ **kwarg} = \left\{\mathbf{c}^{(\ell)}, \sigma \right\}, \;\ell =1,\ldots, L\) (here \(L\) denotes the size of the training dataset, similar to the number of the linear systems shown in Equation 3 ) and output solution \(\mathbf{u}^{(\ell)}_{\theta} \in {\mathcal{U}}\) returned from CNNs with \(\theta.\) Our target is to find the learning operator such that \[~\label{eq:operator95approximation} \mathcal{N}_\theta(\mathbf{b}^{(\ell)}, \text{**kwarg}) \sim \mathbf{u}^{(\ell)}_{\theta},with\text{**kwarg} = \left\{\mathbf{c}^{(\ell)}, \sigma \right\}, \;i =1,\ldots, L,\tag{6}\] which means the learned operator \(\mathcal{N}_\theta\) approximates to the inverse of the coefficient matrix \(A^{-1}\) since we have \(\mathbf{u}= A^{-1} (\mathbf{b}, \text{**kwarg})\) from Equation 3 . Because of this, the learned operator \(\mathcal{N}_\theta\) can be used as a preconditioner in the FGMRES method described in Algorithm [alg:nn-fgmres], where the preconditioned Krylov basis is computed as: \[~\label{eq:nns95pre} \mathbf{z}_j = \mathcal{N}_\theta(\mathbf{v}_j, \mathbf{c}, \sigma)\tag{7}\] for the application in the parametric Helmholtz equation 1 . Except for use the Krylov basis as the input of the trained inference, refer to [21] for more details about other possibilities with similar numerical performance.
To distinguish from the original FGMRES [5], we denote FGMRES with the preconditioning step shown in Equation 7 for Algorithm [alg:nn-fgmres] as the CNN-FGMRES algorithm. Compared to the original FGMRES, these steps are performed in 32-bit, consistent with the precision used during the training processes. The other remaining operations in the CNN-FGMRES algorithm are conducted in 64-bit. Thus, this CNN-FGMRES algorithm is realized under mixed-precision calculations.
As introduced in Section 2.2, our approach uses a neural network as a nonlinear preconditioner. To do so, we design a convolutional neural network based on the U-Net architecture [26], which takes the discretized source field, velocity field and the PML damping function \(\sigma\) as inputs. To make them suitable to the convolutional nature of the U-Net, these 2D fields are recast as image-like tensors. Specifically, recalling that \(n=N^2\), the complex-valued source term \(\mathbf{b}\in \mathbb{C}^n\) may be recast as a \(2 \times N \times N\) tensor composed of two \(N \times N\) channels representing the real and imaginary parts of the field. So does the complex \(\sigma \in \mathbb{C}^n\) function. Similarly, the real-valued velocity field \(\mathbf{c}\in \mathbb{R}^n\) may be recast as a \(1 \times N \times N\) tensor. These five inputs are concatenated along channel dimension to get a \(5 \times N \times N\) input tensor. The output of the network is a \(2 \times N \times N\) tensor, representing a 2D, complex-valued field. As we shall see in Section 3.3, the network will be trained to provide approximations (predictions) \(\hat{\mathbf{u}}_\theta\) of discrete solution fields \(\mathbf{u}\) of the linear system \(A(\mathbf{c}) \mathbf{u}= \mathbf{b}\) corresponding to a discretized version of the Helmholtz problem 1 with seetings described in Table ¿tbl:tab:helm95model95parameters?.
The primary hyperparameter of the U-Net architecture is its depth. For the Helmholtz problems considered in Table ¿tbl:tab:helm95model95parameters?, a depth of four is sufficient to construct an effective neural preconditioner. The resulting U-Net architecture is illustrated in Figure 2. It consists of an encoder–decoder structure, where the encoder progressively downsamples the input and the decoder reconstructs the output through successive upsampling operations. The encoder and decoder are linked through a bottleneck layer and a set of skip connections. These skip connections play a crucial role in preserving fine-scale information while combining it with the high-level features extracted by the encoder, thereby enabling the network to capture both global context and local details. The corresponding hyperparameter settings are summarized in Table [tbl:tab:model95hyperparameters]. With the Helmholtz problem parameters and network hyperparameters specified in Tables ¿tbl:tab:helm95model95parameters? and [tbl:tab:model95hyperparameters], our U-Net architecture contains approximately 832 K trainable parameters, which is independent of \(n\) owing to its convolutional nature.
The abstract, functional representation of the neural network preconditioning with trainable parameters \(\theta\) is thus a nonlinear operator, \[\label{eq:nn95operator} \begin{align} \mathcal{N}_\theta\colon \mathbb{C}^n \times \mathbb{R}^n \times \mathbb{C}^n & \to \mathbb{C}^n,\\ (\mathbf{b}, \mathbf{c}, \sigma) & \mapsto \mathcal{N}_\theta(\mathbf{b}, \mathbf{c}, \sigma) = \mathcal{N}_\theta(\mathbf{c}, \sigma)(\mathbf{b}), \end{align}\tag{8}\] where the conversions to and from image-like tensors are implicit.
In this work, mixed training datasets are employed to train the neural operator preconditioners. Each training dataset is composed of pairs of velocity fields \(\mathbf{c}\) and source fields \(\mathbf{b}\), denoted by \((\mathbf{c},\mathbf{b})\). By defining the distributions of \(\mathbf{c}\) and \(\mathbf{b}\) independently, different combinations of structured, application-oriented samples and randomized samples can be constructed. We first define a fully mixed dataset by considering the following distributions for \(\mathbf{c}\) and \(\mathbf{b}\):
For the velocity field \(\mathbf{c}\in \mathbb{R}^n:\) \(\mathbf{c}\in \left\{{\cal U}([1,2]), \;\text{\texttt{TruncIdealSkulls}}, \;\text{\texttt{randConstant}}, \;\text{\texttt{square}} \right\},\) the four configurations are sampled with probabilities \([0.4, 0.4, 0.1, 0.1],\) respectively:
\(\mathbf{c}\stackrel{\textrm{iid}}{\sim}{\cal U}([1, 2])\): each entry of \(\mathbf{c}\) is independently sampled from the uniform distribution on \([1, 2]\);
\(\mathbf{c}\stackrel{\textrm{iid}}{\sim}\text{\texttt{TruncIdealSkulls}}\): \(\mathbf{c}\) is generated from the centrally truncated regions of the training dataset of idealized skull models introduced in [2]. The truncated region preserves the main geometric features of the idealized skull structures;
\(\mathbf{c}\stackrel{\textrm{iid}}{\sim}\text{\texttt{randConstant}}\): all entries of \(\mathbf{c}\) are assigned the same value, sampled uniformly at random from \([1, 2]\);
\(\mathbf{c}\stackrel{\textrm{iid}}{\sim}\text{\texttt{square}}\): a square region centered in the computational domain is assigned the value 2, while all remaining entries are set to 1.
For the source field \(\mathbf{b}\in \mathbb{C}^n:\) \(\mathbf{b}\in \left\{{\cal N}(0,1), \;\text{\texttt{dirac}}, \;\text{\texttt{dirac1}}, \;\text{\texttt{dirac4}} \right\},\) the four configurations are sampled with probabilities \([0.6, 0.3, 0.05, 0.05],\) respectively:
\(\mathbf{b}\stackrel{\textrm{iid}}{\sim}{\cal N}(0,1)\): each entry of \(\mathbf{b}\) is independently sampled from a standard normal distribution with zero mean and unit variance;
\(\mathbf{b}\stackrel{\textrm{iid}}{\sim}\text{\texttt{dirac}}\): a single entry of \(\mathbf{b}\) is assigned the value 60 at a randomly selected grid location, while all remaining entries are set to zero;
\(\mathbf{b}\stackrel{\textrm{iid}}{\sim}\text{\texttt{dirac1}}\): a single entry of \(\mathbf{b}\) is assigned the value 60 at the center of the computational domain, while all remaining entries are set to zero;
\(\mathbf{b}\stackrel{\textrm{iid}}{\sim}\text{\texttt{dirac4}}\): four entries of \(\mathbf{b}\) are assigned the value \(30\) at four fixed grid locations within the computational domain, while all remaining entries are set to zero.
Refer to Figure 3 for a visualisation of representative samples from the different types of training data pairs \((\mathbf{c}, \mathbf{b})\) described in blocks [list95data95c]-[list95data95b]. We note that the amplitudes 60 or 30 was applied to
the dirac-type source configurations so that the resulting source field has a magnitude comparable to that of the standard normal setting \({\cal N}(0,1)\) over the computational domain. This is important in
the training process, which normalization helps maintain the effectiveness of the neural operator preconditioning across different source distributions. Such localized dirac-type sources are frequently encountered in practical wave propagation
problems, including transcranial ultrasound simulations and spherical wave propagation, where the source term is concentrated at a small number of points or within a localized region and vanishes elsewhere.
According to the velocity and source distributions defined in blocks [list95data95c]-[list95data95b], five additional mixed datasets are then obtained by selecting different subsets or combinations of these distributions. The final six training datasets considered in this work are summarized in Table 1. Each dataset is denoted by D0-D5, where D2 corresponds to the mixed setting described in blocks [list95data95c]-[list95data95b]. As shown in Table 1, datasets D0 and
D1 share the same configuration of \(\mathbf{c}\) but differ in the source term \(\mathbf{b}.\) Similarly, D2 and D3 share the same configuration of
\(\mathbf{c}\) while differing in \(\mathbf{b}.\) In contrast, D4 and D5 use the same configuration of \(\mathbf{b}\) but vary in
\(\mathbf{c}.\) In the following, we describe how neural networks are trained using these six mixed datasets to construct effective neural operator preconditioners, with a primary focus on applications to transcranial
ultrasound simulation.
| #Dataset | Source term \(\bb\) (with probability) | Velocity field (or Speed of sound) \(\bc\) (with probability) |
|---|---|---|
| D0 | \(\left\{{\cal N}(0,1), \;\text{dirac}, \;\text{dirac1}, \;\text{dirac4} \right\}\) (\([0.6, 0.3, 0.05, 0.05]\)) | \(\left\{{\cal U}([1,2]), \;\text{TruncIdealSkulls} \right\}\) (\([0.5, 0.5]\)) |
| D1 | \({\cal N}(0,1)\) | \(\left\{{\cal U}([1,2]), \;\text{TruncIdealSkulls} \right\}\) (\([0.5, 0.5]\)) |
| D2 | \(\left\{{\cal N}(0,1), \;\text{dirac}, \;\text{dirac1}, \;\text{dirac4} \right\}\) (\([0.6, 0.3, 0.05, 0.05]\)) | \(\left\{{\cal U}([1,2]), \;\text{TruncIdealSkulls}, \;\text{randConstant}, \;\text{square} \right\}\) (\([0.4, 0.4, 0.1, 0.1]\)) |
| D3 | \({\cal N}(0,1)\) | \(\left\{{\cal U}([1,2]), \;\text{TruncIdealSkulls}, \;\text{randConstant}, \;\text{square} \right\}\) (\([0.4, 0.4, 0.1, 0.1]\)) |
| D4 | \({\cal N}(0,1)\) | |
| D5 | \({\cal N}(0,1)\) | \({\cal U}([1,2])\) |
The training process aims at finding the parameters \(\theta^* \in \Theta\) of the neural network that minimize a certain cost functional \(J\colon \Theta \to \mathbb{R}\), where \(\Theta\) denotes the parameter space, that is, \[\label{eq:nn95optim} \theta^* \in \mathop{\mathrm{arg\,min}}_{\theta \in \Theta} J(\theta).\tag{9}\] For our specific problem, the cost function takes the form \[\label{eq:nn95cost95fn} J(\theta) = \mathbb{E}_{\mathbf{b}, \mathbf{c}, \sigma} [\mathcal{L}_\theta(\mathbf{b}, \mathbf{c}, \sigma)],\tag{10}\] where \(\mathbf{b},\) \(\mathbf{c}\) and \(\sigma\) are the inputs of the network, \(\mathcal{L}_\theta\) is a loss function and \(\mathbb{E}_{\mathbf{b}, \mathbf{c}, \sigma}\) denotes the expectation operator with respect to \(\mathbf{b},\) \(\mathbf{c}\) and \(\sigma.\) In this work, we consider the loss function defined by \[\label{eq:nn95loss95fn} \mathcal{L}_\theta(\mathbf{b}, \mathbf{c}, \sigma) = \dfrac{\|\mathbf{b}- A(\mathbf{c})\mathcal{N}_\theta(\mathbf{b},\mathbf{c}, \sigma)\|^2}{\|\mathbf{b}\|^2} = \dfrac{\|\mathbf{b}- A(\mathbf{c})\hat{\mathbf{u}}_\theta \|^2}{\|\mathbf{b}\|^2} = \dfrac{\|\mathbf{b}- \hat{\mathbf{b}}_\theta \|^2}{\|\mathbf{b}\|^2},\tag{11}\] where \(\hat{\mathbf{u}}_\theta :=\mathcal{N}_\theta(\mathbf{b}, \mathbf{c}, \sigma) = \mathcal{N}_\theta(\mathbf{c}, \sigma)(\mathbf{b})\) denotes the prediction of the neural network, and \(\hat{\mathbf{b}}_\theta = A(\mathbf{c})\hat{\mathbf{u}}_\theta\). The neural network is thus trained to be a good approximation of \(A^{-1}\), i.e., roughly speaking, so that \(\mathcal{N}_\theta(\mathbf{c}, \sigma) \approx A^{-1}(\mathbf{c})\) for any \(\mathbf{c}\) with fixed \(\sigma\). The training process can here be interpreted as unsupervised, in the sense that the prediction \(\hat{\mathbf{u}}_\theta\) are not (directly) compared to reference values \(\mathbf{u}^*\) in the loss function. We remark that the residual of the linear system for a given prediction \(\hat{\mathbf{u}}_\theta = \mathcal{N}_\theta(\mathbf{c}, \sigma)(\mathbf{b})\), namely \(\mathbf{b}- A(\mathbf{c})\hat{\mathbf{u}}_\theta\), which is at the core of the loss definition, corresponds to an algebraic (discretized) version of the residual of the PDE, including the boundary conditions. As such, the loss may be interpreted as being physics-informed. In terms of numerical linear algebra, we further note that for any given \(\mathbf{c}\in \mathbb{R}^n\), \(\mathcal{L}_\theta(\mathbf{b},\mathbf{c}, \sigma) = \eta_b(\hat{\mathbf{u}}_\theta)^2\), where \(\eta_b(\tilde{\mathbf{u}})\) denotes the backward error with respect to \(\mathbf{b}\) of \(\tilde{\mathbf{u}}\) as an approximate solution to the linear system \(A(\mathbf{c}) \mathbf{u}= \mathbf{b}\) [25], [27]. An alternative interpretation that can be drawn from the last representation of the loss function 11 is that the neural network is trained so that \(A(\mathbf{c})\mathcal{N}_\theta(\mathbf{c}, \sigma)\) be a good approximation of the identity operator, similar to an autoencoder [28]. Thus, the neural network \(\mathcal{N}_\theta\) can then be interpreted as an encoder, trained such that it is decoded by \(A\). The training process is summarized in Figure 4.
To train the CNN-based neural operator preconditioning described in Equation 8 or Figure 4, 10,000 random samples are generated for each training dataset listed in Table 1, including distributions for \(b\) and \(c.\) The training is conducted using the Adam optimizer, with the maximum number of epochs (denoted as
max_epoch) set to 500, that is \(\text{\texttt{max\_epoch}}=500\). The learning rate is initialized at 0.001 and gradually decreases to a minimum value of \(10^{-5}.\) Refer to
Table ¿tbl:tab:helm95model95parameters?-[tbl:tab:training95hyperparameters] for more details about the Helmholtz
operator parameters, the model, and training hyper-parameters of the U-Net model designed for the Helmholtz operators on a two-dimensional domain. This U-Net preconditioning model, configured with a \(\text{\texttt{depth}}\) of 4, is trained under the single precision (that is float32 (32-bit) and complex64) on 4 NVIDIA V100 GPUs (32.5 GB) with (Model name: Intel(R) Xeon(R) Gold 6130 CPU @ 2.10GHz)
located on the CAIUS cluster1 of the University of Strasbourg Computing Centre (CCUS). The network architecture and the Helmholtz operator training utilities are implemented
using the PyTorch and PyTorch Lightning libraries. The training time required for this U-Net preconditioning model under these different training datasets for the discrete Helmholtz equations is summarized in
Table ¿tbl:tab:train95time?.
This numerical section focus on illustrating different numerical features of these six trained U-Net neural operator preconditioners, including testing similar but unseen scenarios and
checking network generalizability from various applications. Most importantly, the dataset from the transcranial ultrasound therapy [1].
Specifically, in these testing phases, numerical experiments are carried out with the discrete parametric Helmholtz equations 3 that the neural operators have never seen before. Without special notes, we display results of
GMRES and (that is FGMRES preconditioned by the trained U-Net preconditioner with 4 depth) without restart. A classical stopping criterion for the numerical linear algebra solvers is based on backward error analysis and consists of stopping
the iteration when \[\label{eq:BEb} {\eta_b} = \frac{\|A{\mathbf{u}} - {\boldsymbol{b}}\|}{\| {\boldsymbol{b}} \|} \le \varepsilon, \text{ with } \varepsilon= 10^{-12} \text{ by
default}.\tag{12}\] The performance of the involved algorithms is evaluated in terms of the number of iterations, denoted as \(Its\), as well as the execution time, denoted as \(ET\), required to converge. The maximum dimension of the Krylov search space \(m\) shown in Algorithm [alg:nn-fgmres] is set to be 512, that is \(m=512.\) The number of restart is \(l=1 \text{ or } 10\) (here \(l=1\) refers
to the case without restart). Thus, the whole maximum iteration (denoted as \(maxIts\)) is \(l \times 512\). We stop the algorithm when satisfying formula 12 or when
\(maxIts= l \times 512.\) Notation \(Its_{max}, Its_{avg}, Its_{min}\) refers to the maximal, average, and minimum number of consumed \(Its\) when testing
multiple examples (that is \(L>1\)). Symbol \((Its_{max}, Its_{avg}, Its_{min})_{E}^{*}\) denotes the algorithm diverges with \(E\) linear systems among
the whole \(L\) testing systems (\(0 < E \leq L\)), or these \(E\) linear systems are unable to reach the targeted accuracy within \(maxIts\).
Recall that this 2D U-Net inference was trained on GPUs, whereas the classical subspace solvers have been generally implemented on CPUs. In the following subsections, the GMRES and solvers are implemented in a Python prototype, supporting
both CPU and GPU backends for computations with numpy and cupy Python libraries. Without special notes, those two involved solvers in each of the subsections are running on the same 4 V100 GPUs devices used in the training process
for a fair comparison.
We first assess the effectiveness of the six trained neural operator preconditioning inferences on problems similar to those used during the training, that is, linear systems arising from the discretization on a grid, with physical parameters \(\mathbf{c}\) and \(\mathbf{b}\) drawn from the randomly mixed dataset D0-D5 described in Table 1. The corresponding
numerical results for GMRES and CNN-FGMRES applied to the \(L\) linear systems are reported in Table ¿tbl:tab95results95test95idea95skull?. For the testing cases with \(L=500,\) we observe that GMRES (mostly) fails to reach the prescribed accuracy within \(maxIts.\) In contrast, CNN-FGMRES converges with significantly reduced computational cost, thanks to the
trained neural operator preconditioner. This effectiveness and robustness are maintained when increasing the number of test problems to \(L=1000,\) thereby covering a broader range of physical parameter variations. This
conclusion holds consistently across all trained models. These results are further illustrated in Figure 5, which presents the convergence history of GMRES and CNN-FGMRES (with the trained on dataset
D0) with \(L\) systems. Figure 6 provides visualizations of the obtained complex wavefield solutions under different configurations of the physical
parameters.
Figure 5: Results of the CNN-preconditioning trained with dataset D0 in Table 1: Convergence in terms of \(\eta_b=10^{-12}\) of CNN-FGMRES and GMRES
with \(maxIts=512\) without restrat for \(L\) examples from the test dataset with mixed data.. a — CNN-FGMRES with \(L=1000,\) dataset D0, b —
GMRES, CNN-FGMRES with \(L=500,\) dataset D0
Figure 6: Visualize the complex wavefield in real-imag part (upper two plots of each subplot (a) to (d)) and its corresponding velocity \(\mathbf{c}\) and source term \(\mathbf{b}\) (lower two plots of each subplot) of four examples solved by CNN-FGMRES (the CNN was trained with dataset D0), and \(\eta_b = 10^{-12},\) the four examples are from the
test dataset with mixed data \((\mathbf{c}, \mathbf{b}).\). a — \(\mathbf{c}\stackrel{\textrm{iid}}{\sim}{\cal U}([1,2]) + \mathbf{b}\stackrel{\textrm{iid}}{\sim}{\cal N}(0,1)\), b — \(\mathbf{c}\stackrel{\textrm{iid}}{\sim}\text{\texttt{TruncIdealSkulls}} + \mathbf{b}\stackrel{\textrm{iid}}{\sim}{\cal N}(0,1)\), c — \(\mathbf{c}\stackrel{\textrm{iid}}{\sim}{\cal U}([1,2]) +
\mathbf{b}\stackrel{\textrm{iid}}{\sim}\text{\texttt{dirac1}}\), d — \(\mathbf{c}\stackrel{\textrm{iid}}{\sim}\text{\texttt{TruncIdealSkulls}} +
\mathbf{b}\stackrel{\textrm{iid}}{\sim}\text{\texttt{dirac}}\)
On the other hand, the effectiveness of the neural operator preconditioner depends on the composition of the training dataset. For example, the model trained on dataset D5, which is constructed from a single configuration of the physical
parameters \((\mathbf{c}, \mathbf{b})\), yields the largest reduction in computational cost relative to GMRES. In contrast, the model trained on dataset D2 (or D0), which contains a more diverse
mixture of parameter configurations, exhibits relatively lower acceleration. This behavior is expected: training on a single configuration allows the network to specialize to that particular distribution, whereas mixed datasets require the model to
represent a broader range of physical configurations. Consequently, mixed datasets generally trade a small amount of efficiency on individual configurations for improved robustness and generalization across diverse problems. This trend is further
illustrated by comparing the results obtained with datasets D0 and D1, as well as with datasets D2 and D3. In each pair, the velocity distributions \(\mathbf{c}\) are
generated using the same mixed strategy, while the source fields \(\mathbf{b}\) differ between a mixed distribution and a single structured configuration (see Table 1).
The following subsections investigate the generalization capabilities of the trained CNN-based preconditioning operators in the primary application of head CT scan examples. To evaluate their generalizability, we test their performance under various heterogeneous parameter configurations in the parametric Helmholtz equation 3 , extending beyond the settings used during training. We first consider the truncated idealized skulls on a \(64 \times 64\) grid with variations in terms of the source field in Section 4.2.1. Then, in Section [subsec:skull95exp], we focus on the primary application involving a practical transcranial ultrasound dataset derived from adult human head CT scans, where both the source and velocity fields exhibit structured heterogeneity on a \(512 \times 512\) grid, which is 64 times larger than that used in training.
It is worth noting that, for challenging large-scale problems such as the practical transcranial ultrasound applications considered in Section [subsec:skull95exp], many classical algebraic preconditioners—including Jacobi, algebraic multigrid, and incomplete LU (ILU) or incomplete Cholesky factorizations—may become impractical. Their construction and application typically require explicit access to the matrix representation of the discretized operator, which can lead to prohibitive memory requirements and setup costs on fine computational grids. In contrast, the proposed neural operator preconditioner, which approximates the inverse of the discretized differential operator, is inherently matrix-free and is applied directly through network inference, thereby naturally avoiding the assembly, storage, and factorization of large sparse matrices. For further details and a quantitative comparison between neural operator preconditioners and several classical algebraic preconditioning techniques on moderate-scale academic examples, we refer the reader to [4], [29].
Before evaluating network generalizability of the models trained with the six dataset described in Table 1 on the more challenging transcranial ultrasound dataset, we first assess their performance on 1000 idealized skull samples (truncated test idealized skulls introduced in [2]) with some given source fields. In this setting, the velocity field is given by \(\mathbf{c}\stackrel{\textrm{iid}}{\sim}\text{\texttt{TruncIdealSkulls}}.\) For these 1000 realizations of \(\mathbf{c}\), which differ from those used during training, we consider the following two configurations of the source term \(\mathbf{b}\):
\(\mathbf{b}\stackrel{\textrm{iid}}{\sim}\text{\texttt{dirac1}}\): a single nonzero entry of \(\mathbf{b}\) is set to 1 at the fixed grid location \((60, 32)\), while all other entries are zero;
\(\mathbf{b}\stackrel{\textrm{iid}}{\sim}\texttt{dirac1\_smooth}\): a smoothed version of \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}\text{\texttt{dirac1}}\), defined as a localized Gaussian source \(\mathbf{b}(\mathbf{x}, \mathbf{y}) = \exp^{{- \left((x-c_x)^2 + (x-c_y)^2\right)} / {2 \sigma^2}},\) where \(\sigma = 2\) and \((c_x,c_y) = (60,32)\) denotes the grid location.
With these physical parameter configurations, the numerical results for solving the \(L\) linear systems arising from Equation 3 are reported in Table ¿tbl:tab95results95test95dirac195idea95skull? in terms of \((Its_{max}, Its_{avg}, Its_{min})\) and the elapsed time \(ET\).
Figure 7: Convergence in terms of \(\eta_b=10^{-12}\) of CNN-FGMRES and GMRES with \(maxIts=512\) without restrat for \(L\) examples with
varying velocity \(\mathbf{c}\stackrel{\textrm{iid}}{\sim}\text{\texttt{TruncIdealSkulls}}\) from the test dataset and the fixed source term \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}\text{\texttt{dirac1}}\) with fixed grid-point locates at \(60 \times 32\) within 2D domain \(64 \times 64,\) and its
polluted one \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}\texttt{dirac1\_smooth}\) with center radial decay Gaussian noise by \(\texttt{radius}=5, \texttt{noise\_level}=0.05.\) In (a)-(b), the
presented CNN-FGMRES results are for the CNN-preconditioning trained with dataset D0, and \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}\text{\texttt{dirac1}}.\) In (c), it shows corresponding results from the
CNN-preconditioning trained with the six different datasets D0-D5 described in Table 1, and \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}\texttt{dirac1\_smooth}.\). a — CNN-FGMRES with \(L=1000, \mathbf{b}\stackrel{\textrm{iid}}{\sim}\text{\texttt{dirac1}},\) dataset D0, b —
CNN-FGMRES , GMRES with \(L=20, \mathbf{b}\stackrel{\textrm{iid}}{\sim}\text{\texttt{dirac1}},\) dataset D0, c — CNN-FGMRES with \(L=20,\) \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}\texttt{dirac1\_smooth},\) and CNN was trained with the six different datasets D0-D5
As in the previous experiments, the application of the trained CNN-based operator preconditioners enables to converge with substantially fewer iterations and lower computational cost than GMRES, which fails to achieve the prescribed accuracy within
\(maxIts.\) Moreover, the robustness of the proposed approach is maintained when the number of test problems is increased to \(L=1000.\) The corresponding convergence histories for the model
trained on dataset D0 are shown in Figure 7 (a)-(b). Among the six trained models, the operator preconditioner trained on dataset D4 exhibits the best performance. This
is expected, since the tested velocity fields \(\mathbf{c}\stackrel{\textrm{iid}}{\sim}\text{\texttt{TruncIdealSkulls}}\) share structural similarities with those included in dataset D4. Nevertheless, it is
noteworthy that the operator preconditioner trained on the fully random dataset D5, which contains no particular geometric structure, still provides a significant acceleration of CNN-FGMRES , albeit with somewhat reduced effectiveness. The
differences among the six trained preconditioners can be clearly observed in Figure 7 (c), which presents the convergence histories of CNN-FGMRES for a set of 20 linear systems. These results
indicate that the choice of the training dataset does impact performance on some targed testing cases. This effect becomes even more pronounced in the practical transcranial ultrasound examples presented in the next section.
On the other hand, no significant differences are observed between the source field configurations \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}\text{\texttt{dirac1}}\) and \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}\texttt{dirac1\_smooth}\) in the numerical results reported in Table ¿tbl:tab95results95test95dirac195idea95skull?. To further investigate their effects, we visualize the corresponding wavefields in Figure 8. The smoothed-source configuration produces less localized irregularity near the source location, thereby providing a clearer visualization of both the velocity distribution and the resulting wavefield patterns.
Figure 8: Visualize the complex wavefield in real-imag part (upper two plots of each subplot (a) to (d)) and its corresponding velocity \(\mathbf{c}\) and source term \(\mathbf{b}\) (lower two plots of each subplot) of two examples solved by CNN-FGMRES (the CNN was trained with dataset D0), and \(\eta_b = 10^{-12},\) where the two varying \(\mathbf{c}\stackrel{\textrm{iid}}{\sim}\text{\texttt{TruncIdealSkulls}}\) is from the test dataset, and the fixed \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}\text{\texttt{dirac1}}\) with fixed
grid-point locates at \(60 \times 32\) within 2D domain \(64 \times 64,\) and its polluted one \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}\texttt{dirac1\_smooth}\) with center radial decay Gaussian noise by \(\texttt{radius}=5, \texttt{noise\_level}=0.05.\). a — \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}\text{\texttt{dirac1}}\), b — polluted \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}\texttt{dirac1\_smooth}\), c — \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}\text{\texttt{dirac1}}\), d — polluted \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}\texttt{dirac1\_smooth}\)
In this section, we evaluate the trained neural operator preconditioners on four practical adult human head CT scan examples obtained from qure.ai [1]. The dataset is publicly available through Kaggle.2 Each example is constructed from a transverse CT slice of an adult human skull, yielding a heterogeneous velocity \(\mathbf{c}\in \mathbb{R}^{512^2}\) distributed on a \(512 \times 512\) grid. This computational grid is 64 times larger than the \(64 \times 64\) grid used during training, providing a significantly challenging out-of-distribution testing case. Visualizations of the four velocity fields are shown in the lower-left panels of Figure 11 (a)-(d). The corresponding source field \(\mathbf{b}\in \mathbb{R}^{512^2}\) is modeled as a focused ultrasound transducer represented by a one-dimensional arc, illustrated by the red curve in the lower-right panels of Figure 11 for its visualization. The transducer aperture diameter and radius of curvature are both set to 60 mm, and the source frequency is fixed to 490 kHz.
To solve these challenging large-scale problems, restarted GMRES and CNN-FGMRES are employed with restart parameter \(l=10\), resulting in a maximum of \(maxIts=5120\) iterations
throughout this section. Furthermore, in addition to GMRES, we also benchmark our approach against the closely related state-of-the-art neural network solver: the learned optimizer proposed by Stanziola et al. in [2]. Specifically, we use their released IterativeSolver (i.e., the learned optimizer) together with the
pretrained checkpoint provided in the authors’ GitHub repository3 and run the learned optimizer solver up to 5120 nonlinear Richardson iterations. Since the downloaded
physical parameter pairs \((\mathbf{c}, \mathbf{b})\) are loaded directly into CPU memory, both the learned optimizer, GMRES and CNN-FGMRES are executed on the same CPU to avoid unnecessary data transfers between CPU and
GPU devices. Specifically, all experiments in this section are performed on the CAIUS cluster, using Intel Skylake compute nodes equipped with 2x 16-core Intel Xeon Gold 6126 processors running at 2.60 GHz and 192 GB of RAM per node.
Figure 9: Convergence in terms of \(\eta_b=10^{-3}\) for (a)-(c), and \(\eta_b=10^{-6}\) for (d) of CNN-FGMRES and GMRES with \(maxIts=5120\)
with restrat equals 10, together with the learned optimizer proposed in [2] with \(maxIts=200, 5120,\) on four (i.e., \(L=4\)) practical head CT scans problems. In (a)-(b), the presented CNN-FGMRES results are for the CNN-preconditioning trained with dataset D0.
In (c), it shows corresponding results from the CNN-preconditioning trained with the six different datasets D0-D5 described in Table 1. The (d) is the counterpart results of (c) with
datasets D0-D2 and \(\eta_b=10^{-6}.\). a — CNN-FGMRES (from dataset D0) and Learned optimizer [2] with \(L=4, \;maxIts=200\), b — CNN-FGMRES (from dataset D0), GMRES and Learned optimizer [2] with \(L=4, \;maxIts=5120\), c — CNN-FGMRES with \(L=4,\) and the CNN was trained with the six different datasets D0-D5, d — Counterpart of (c) with the three different datasets D0-D2 and \(\eta_b=10^{-6}\)
Table ¿tbl:tab95results95skull? reports the numerical results of GMRES, the learned optimizer and CNN-FGMRES for the four practical transcranial ultrasound examples, measured in terms of the iteration counts \(Its\) and CPU elapsed time \(ET\) under different prescribed accuracy requirements. Corresponding convergence histories are presented in Figure 9. The
conclusions drawn from the previous experiments in terms of comparing CNN-FGMRES to GMRES remain valid for these challenging large-scale test cases. In particular, the trained neural operator preconditioners consistently reduce both the number of
iterations and the computational time compared with unpreconditioned GMRES. For instance, CNN-FGMRES with the neural operator preconditioner trained on dataset D0 requires at most 191 \(Its\) with around
4 minutes for the most challenging example (i.e., \(\mathbf{c}\) from CT000112.dcm presented in Figure 9 (a)), and approximately 12 minutes in total to solve all four
examples with accuracy \(\eta_b=10^{-3}.\) In contrast, GMRES only successfully solves the simplest case (i.e., \(\mathbf{c}\) from CT000255.dcm presented in Figure 9 (b)), requiring 4546 \(Its\), and fails to reach the prescribed accuracy for other three examples within the maximum iteration limit \(maxIts.\)4 On the other hand, we observed that the learned optimizer proposed in [2] only successfully converges for one example shown in Figure 9 (a)-(b), which is the same test case (i.e., \(\mathbf{c}\) from CT000112.dcm) considered in [2],
reaching a stagnated attainable accuracy of approximately \(\eta_b=10^{-4}\) after 3084 Richardson iterations (noted by a green \(*\) in the green curve in Figure 9 (b)). It fails to converge for the remaining three examples (see Figure 9 (a)-(b) or Figure 13 in Appendix [subsec:nn-fgmres95three95wave95evl]). In contrast, the proposed CNN-FGMRES successfully solves all four examples with
substantially fewer iterations and achieves arbitrary prescribed accuracy.
Figure 10: Evolution of the complex wavefield predicted by the learned optimizer [2], CNN-FGMRES and GMRES at iterations 1, 2, 3, 11, 21, 61, 190 for the transcranial ultrasound example based on the head CT scan \(\mathbf{c}\) from CT000112.dcm defined 2D
domain \(512 \times 512.\). a — CNN-FGMRES with dataset D0, b — GMRES, c — Learned optimizer [2]
The evolution of the complex wavefield predicted by GMRES, the learned optimizer and CNN-FGMRES trained with dataset D0 for the most challenging example with \(\mathbf{c}\) from CT000112.dcm is
presented in Figure 10. The CNN-FGMRES iterates rapidly recover the main wavefield features: the focal region and the corresponding focal pressure are already well established after approximately 21 iterations, while
most of the fine-scale wavefield structure is captured within about 61 iterations. In contrast, GMRES does not exhibit comparable convergence within the same number of iterations. In fact, approximately 4000 iterations are required before the overall
wavefield structure becomes discernible. Furthermore, the learned optimizer requires more than 200 iterations before recovering most of the wavefield structure for the same transcranial ultrasound example (see Figure 10 (c) or [2]), and it fails to converge other three
examples.5 This qualitative comparison further highlights the effectiveness and robustness of the proposed neural operator preconditioning strategy in accelerating the
convergence of Krylov subspace methods for the Helmholtz equations arising in transcranial ultrasound simulations. Figure 11 provides visualizations of the resulting final complex wavefields together with the corresponding
physical parameter configurations.
Figure 11: Visualize four head CT scans on 2D domain \(512 \times 512\) solved by CNN-FGMRES with dataset D0, and \(\eta_b=10^{-6}.\). a — \(\mathbf{c}\) is from CT000112.dcm, b — \(\mathbf{c}\) is from CT000255.dcm, c — \(\mathbf{c}\) is from CT000200.dcm, d —
\(\mathbf{c}\) is from CT000100.dcm
On the other hand, several additional observations can be made. A key finding is that the performance differences among the six trained neural operator preconditioners become significantly more pronounced in these out-of-distribution testing scenarios.
For example, comparing the models trained on datasets D0 and D1 reveals that increasing the diversity of the source field configurations can substantially improve efficiency, especially when a tighter accuracy requirement of \(10^{-6}\) is imposed. Likewise, comparing the results obtained with datasets D0 and D2 highlights the influence of the velocity distributions used during training. The corresponding differences become
increasingly evident as the testing problems depart further from the training configurations. Furthermore, when comparing dataset D1 with datasets D4 and D5, which rely exclusively on either idealized skull
configurations or fully random velocity fields, respectively, the model trained on D1 consistently achieves better performance. This suggests that combining structured and heterogeneous training samples can be more beneficial than relying
solely on highly specialized or completely randomized datasets. Overall, these results indicate that training on specialized datasets can improve performance on closely related target problems, whereas training on more diverse datasets tends to enhance
robustness and generalization across a broader range of scenarios. Among the six models considered, the preconditioner trained on dataset D0 exhibits the best overall performance for these practical transcranial ultrasound simulations. A
possible explanation is that D0 achieves a favorable balance between specialized idealized skulls and randomized velocity fields while simultaneously incorporating the greatest diversity of source field configurations. The convergence
histories of CNN-FGMRES with different network preconditoners and convergence threshold \(\eta_b\) are presented in Figure 9 (c)-(d).
Up to this point, two important insights can be summarized from the perspective of generalization to practical simulations. First, the composition of the training dataset plays a critical role in determining the effectiveness of learned models, particularly when they are applied to challenging out-of-distribution problems. Second, for scientific machine learning approaches targeting practical applications, training datasets should carefully balance application-specific structures and randomized samples to achieve both specialization and robust generalization.
We specially note that, beyond the primary application in transcranial ultrasound, the same trained neural operator preconditioners can also be applied to accelerate simulations in other settings. We refer the reader to Appendix [subsec:other95apps] for their use in simplified models of spherical wave propagation and seismology.
In this study, we investigated the use of neural networks trained on mixed datasets as nonlinear preconditioners for solving parametric Helmholtz equations with Krylov subspace methods. Six mixed training datasets were considered to construct neural operator preconditioners, with a primary focus on large-scale transcranial ultrasound simulations. A systematic investigation of the training datasets demonstrates that dataset composition plays a critical role in determining the effectiveness and generalization capability of the learned preconditioners. In particular, a suitable balance between application-specific structured samples and randomized samples leads to improved performance on both in-distribution and out-of-distribution problems. The proposed hybrid approach combines the strong generalization capability of neural operator preconditioners with the robustness of classical iterative methods, enabling accelerated solutions of large-scale Helmholtz problems to arbitrary prescribed tolerances. Numerical experiments demonstrate that the proposed neural operator preconditioners consistently accelerate Krylov subspace methods for both idealized and practical head CT scan examples, while remaining effective for other Helmholtz-based applications.
Overall, this work highlights the importance of dataset design in scientific machine learning and provides practical guidelines for constructing training datasets for neural operator preconditioners targeting real-world large-scale PDE simulations. More broadly, it demonstrates the potential of integrating learned approximations of inverse PDE operators into established numerical linear algebra algorithms through a matrix-free framework for efficient large-scale simulations. Rather than replacing classical iterative solvers, neural operator preconditioners trained on mixed datasets serve as complementary components that provide strong generalization capability while retaining the flexibility and robustness of Krylov subspace methods.
The author would like to thank Victor Michel-Dansac and Killian Lutz from Inria MACARON team and Université de Strasbourg for the valuable discussions and exchanges regarding the physical parameters settings and numerical experiments considered in Section 4.2.1.
The author would also like to express her sincere gratitude to her Ph.D. supervisors, Luc Giraud (Inria Concace team, France) and Paul Mycek (Concace, Cerfacs, France). Research on transcranial ultrasound simulations was initiated during her doctoral studies under their supervision and continued to be developed after graduation. Their encouragement, support, and confidence in the potential of this research direction provided lasting motivation for pursuing the present study. This work is therefore, in many respects, a continuation of ideas and research interests that originated during the author’s Ph.D.
Figure 12: Counterpart of Figure 10: Evolution of the complex wavefield predicted by CNN-FGMRES trained with dataset D0 at iterations 1, 2, 3, 11, 21, 61, and 158 for (a) or 165 for (b)
161 for (c) for the other three transcranial ultrasound examples defined 2D domain \(512 \times 512.\). a — \(\mathbf{c}\) is from CT000255.dcm, b — \(\mathbf{c}\) is from CT000200.dcm, c — \(\mathbf{c}\) is from CT000100.dcm
Figure 13: Convergence histories of the learned optimizer [2] on the head CT scan examples in Section [subsec:skull95exp].. a — \(maxIts=200\) Richardson iterations, b — \(maxIts=5120\) Richardson iterations
This section explores further applications of the same trained CNN-based preconditioner for solving Helmholtz equations arising from other applications, including the simplified
models of spherical wave propagation with single or multiple source terms, as well as from idealized seismological settings involving wave propagation in heterogeneous media. To represent these scenarios, we define the source term \(\mathbf{b}\) following the \(\texttt{dirac1\_smooth}\) and \(m\) dirac configurations, where \(m \ge 1 \in
\mathbb{N}\) is an integer. Specifically, \(m\) entries of \(\mathbf{b}\) are set to \(\pm 1\) at randomly selected grid locations, while all
remaining entries are zero. To model spherical wave propagation in a homogeneous medium, we first consider the \(\gamma\)-constant case, in which the wave speed \(\mathbf{c}\)
is spatially constant. All entries of \(\mathbf{c}\) are set to a value \(\gamma\) drawn uniformly at random from the interval \([1,2].\) To mimic more
realistic seismological configurations, we construct heterogeneous wave speed fields \(\mathbf{c}\) with layered structures as shown in the publicly available OpenFWI dataset [30]. In each case, the layers are well aligned, and all elements within a given layer share a constant value sampled uniformly from \([\alpha, \beta]\) over the domain. Unless otherwise specified, we take \((\alpha, \beta)=(1, 2)\). The following configurations are considered:
4layered: \(\mathbf{c}\) consists of four aligned layers, each with a constant value;
4faulted: \(\mathbf{c}\) consists of four aligned but discontinuous (faulted) layers, with shifted segments.
We refer the reader to the lower-left panel of each subplot of Figure 15 for visualizations of these homogeneous and heterogeneous wave speed fields with layered structures. We note that these configurations are intended as simplified yet representative test cases for assessing the robustness of the proposed method under both homogeneous and heterogeneous wave propagation regimes.
The numerical results of GMRES and CNN-FGMRES for the \(L=20\) linear systems with these physical parameter settings are reported in Table ¿tbl:tab95results95test95gen95sos95dirac2?. Among the six
trained models, CNN-FGMRES with network preconditioner trained under dataset D3 exhabites the best performance for these testing cases. The corresponding convergence histories for the model trained with dataset D3 are displayed in
Figure 14 (a)-(d), while the convergence histories for all six trained models are presented in Figure 14 (e)-(h). The conclusions drawn in the previous
numerical sections remain valid in this section. For completeness, Figure 15 provides visualizations of the resulting wavefields under some representative physical parameter configurations.
Figure 14: Convergence in terms of \(\eta_b=10^{-12}\) of CNN-FGMRES and GMRES with \(maxIts=512\) without restart for \(L=20\) simplified
models of spherical waves and seismology problems. In (a)-(d), the presented CNN-FGMRES results are for the CNN-preconditioning trained with dataset D3. In (e)-(h), it shows corresponding results from the CNN-preconditioning trained with the
six different datasets D0-D5 described in Table 1.. a — \(\mathbf{c}\stackrel{\textrm{iid}}{\sim}\texttt{\gamma-constant},\) \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}~\texttt{dirac},\) dataset D3, b — \(\mathbf{c}\stackrel{\textrm{iid}}{\sim}\texttt{\gamma-constant},\) \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}2~\texttt{dirac},\) dataset D3, c — \(\mathbf{c}\stackrel{\textrm{iid}}{\sim}\texttt{4layered},\) \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}\texttt{dirac1\_smooth},\) dataset D3, d — \(\mathbf{c}\stackrel{\textrm{iid}}{\sim}\texttt{4faulted},\) \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}\texttt{dirac1\_smooth},\) dataset D3, e — Counterpart of (a) with CNN from datasets D0-D5, f — Counterpart of (b) with CNN from datasets
D0-D5, g — Counterpart of (c) with CNN from datasets D0-D5, h — Counterpart of (d) with CNN from datasets D0-D5
Figure 15: Visualize spherical waves at (a)-(b) and seismology examples (c)-(f) on 2D domain \(64 \times 64\) solved by CNN-FGMRES with \(\eta_b=10^{-12}\) and dataset
D3.. a — \(\mathbf{c}\stackrel{\textrm{iid}}{\sim}\texttt{\gamma-constant}\) with \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}~\texttt{dirac}\), b — \(\mathbf{c}\stackrel{\textrm{iid}}{\sim}\texttt{\gamma-constant}\) with \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}2~\texttt{dirac}\), c — \(\mathbf{c}\stackrel{\textrm{iid}}{\sim}\texttt{4layered}\) with \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}\texttt{dirac1\_smooth}\), d — \(\mathbf{c}\stackrel{\textrm{iid}}{\sim}\texttt{4faulted}\) with \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}\texttt{dirac1\_smooth}\), e — \(\mathbf{c}\stackrel{\textrm{iid}}{\sim}\texttt{4layered}\) with \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}2~\texttt{dirac}\), f — \(\mathbf{c}\stackrel{\textrm{iid}}{\sim}\texttt{4faulted}\) with \(\mathbf{b}\stackrel{\textrm{iid}}{\sim}2~\texttt{dirac}\)
le Cluster de cAlcul Intensif à l’Université de Strasbourg (CAIUS): https://hpc.pages.unistra.fr↩︎
To further illustrate the difficulty of these problems, we note that, even when allowing additional restarts, GMRES requires 14,849 \(Its\) and approximately 13.16 hours to solve the most challenging example shown in Figure 9 (a).↩︎
Refer to Figure 12 in Appendix [subsec:nn-fgmres95three95wave95evl] for the evolution of the complex wavefields predicted by CNN-FGMRES trained on dataset D0 for the other three head CT scan examples.↩︎