1 Introduction↩︎

Deep neural networks (DNNs) are universal approximators, capable of approximating any continuous function defined on a compact domain to arbitrary accuracy [1]–[3]. Therefore, they have been recently widely used to obtain ansatz spaces for the solutions of partial differential equations (PDEs) for a wide range of scientific and engineering problems; see [4] for an overview. The popularity of these DNN approaches can be largely attributed to their simplicity, mesh-free nature, and ability to incorporate data [5], [6], enabling them to effectively tackle complex, high-dimensional, forward and inverse problems in single and multi-query scenarios [7].

Two prominent approaches for solving parametric PDEs using DNNs appear in the literature. The first approach approximates a solution of a given PDE for a particular choice of parameters using a DNN [5], [6]. Popular representatives of this approach are physics informed neural networks (PINNs) [5], which train the DNNs by minimizing the mean square error that incorporates the PDE residual, boundary/initial conditions, and/or observed data. The second approach considers learning the differential operator [8], [9] from the parametrized source terms, boundary/initial conditions, or material properties. The operator learning DNNs, such as DeepONet (DON) [8] or FNO [10], are typically trained by minimizing the mean square error between the network’s output and the solution generated using a high-fidelity method.

Both, physics informed and operator learning, approaches are commonly trained using gradient-based optimizers, e.g., Adam [11] or LBFGS [12]. However, the computational cost of these optimizers can be exorbitant due to the non-convexity and ill-conditioning of the underlying optimization problem. To address this problem, there has been research on improving these optimizers [13]–[15]. This is especially the case when a highly accurate solution is required. In order to reduce the training time, the parallelization of the DNN training process has gained significant attention in the literature; see, for example [16], [17]. Two of the most popular parallel paradigms are data-parallel and model-parallel approaches. Data-parallel approaches [18] split the training data among multiple processing devices. Each device then evaluates the loss/gradient for a given portion of data, and the average value is used to update the network parameters. In contrast, model-parallel approaches aim at splitting the parameters of the DNN model among multiple devices, see for example pipelining techniques [19], [20] or tensor-parallel approaches [16], [21].

In numerical analysis, domain decomposition (DD) methods are commonly used to accelerate the convergence and parallelize solution strategies for large-scale problems [22], [23]. Among these approaches, the additive Schwarz method (ASM) [24] is widely utilized for its simplicity and parallel computing capabilities. In ASM, the computational domain is first decomposed into a collection of smaller, possibly overlapping, subdomains. The solution of a PDE is then sought in parallel for each subdomain, and the acquired local solutions are subsequently combined to update the global solution approximation.

Although the ASM is parallelizable by design, it is not algorithmically scalable, as its convergence tends to deteriorate with an increased number of subdomains [22], [23]. The common remedy for achieving the algorithmic scalability, i.e., keeping the number of iterations bounded independently of the number of subdomains, is to introduce: i) overlap between the subdomains ii) a coarse-level correction step, giving rise to the two-level overlapping additive Schwarz method (TL-ASM) [25]. The role of the overlap is to provide the continuity between local solutions at the interface, while the coarse-level correction step ensures the transfer of data across the entire domain at marginal computational cost [26]. Building upon these advancements, our goal is to construct the nonlinear TL-ASM in order to parallelize the training of scientific machine learning applications while at the same time enhancing the convergence speed of the current state-of-the-art optimizers.

In the field of machine learning, DD approaches are increasingly utilized to enhance the training of DNNs. For example, the parallelization of training through the decomposition of deep convolutional networks has been proposed in [27], [28]. In [29]–[31], multilevel decompositions giving rise to layer-parallel training have been explored. Moreover, in the context of PINNs, several DD-based approaches that leverage the spatiotemporal decomposition of the computational domain have been proposed, such as XPINN [32], cPINN [33], APINN [34], and multilevel DD-based architectures [35], [36]. In this work, we build upon the recent work presented in [30], where single-level additive and multiplicative preconditioners for training of PINNs were introduced. These preconditioners utilize the layer-wise decomposition of DNNs and are, therefore, agnostic to the underlying physics and computational domain. This suggests their potential use beyond the training of PINNs to a broader range of scientific machine learning problems, including data-driven operator learning approaches.

Starting from the single-level additive layer-wise preconditioner [30], we devise a novel two-level overlapping additive Schwarz preconditioner by leveraging and expanding established techniques from numerical analysis related to nonlinear two-level additive Schwarz methods (TL-ASM) [37]–[41]. The main contributions of this work can be summarized as follows.

• We introduce overlapping layer-wise decomposition of DNNs, effectively improving the convergence of non-overlapping layer-wise preconditioners. Motivated by domain-decomposition (DD) approaches in scientific computing, we enforce overlap between the parameters of neighboring layers/subdomains.

• A novel subdomain-wise synchronization strategy is proposed to recombine corrections obtained from parallel subdomain training. This expands upon commonly used methodologies in numerical analysis, which combine subdomain corrections using the weighted averaging technique, see for example literature on the (nonlinear) restricted additive Schwarz method [37], [42]. In the context of the layer-wise DD methods considered in this work, introducing a subdomain-wise synchronization strategy improves the convergence speed of the overlapping DD method as the size of the overlap increases.

• We incorporate a coarse-level training step that enables the global communication between subdomains while maintaining the feed-forward nature of feature flow in the DNNs. This leads to significantly faster convergence compared to using single-level domain decomposition methods.

• The applicability of the single-level and eventually two-level layer-wise preconditioners is extended beyond PINNs to DeepONets, which demonstrate the applicability of the proposed method to wider range of machine-learning problems.

• Through a series of numerical experiments, we demonstrate that the proposed two-level preconditioner yields more accurate DNN models, while also improving the convergence rate of the standard LBFGS optimizer and the baseline single-level preconditioner [30].

This paper is organized as follows. 2 provides an introduction to scientific machine learning, with a particular focus on PINNs and DeepONets. In 3, we propose a novel two-level overlapping additive Schwarz preconditioner for training DNNs. 4 presents a set of benchmark problems and discusses the implementation details. Finally, in 5, we demonstrate the numerical performance of the proposed preconditioner. The conclusion and future work are discussed in 7.

2 Scientific machine learning↩︎

Scientific machine learning is an emerging research field that integrates scientific computing with machine learning techniques to solve complex scientific and engineering problems. In this work, we focus our attention on solving the parameterized PDEs using PINNs [5] and DONs [8]. To this aim, let \(\Omega \subset \mathbb{R}^{d}\), \(d=1,2\), or \(3\), be a closed bounded domain and let the parameter set \(\boldsymbol{H}\) be a compact subset of \(\mathbb{R}^{p}\), where \(p \geq 0\). The parametric PDE is given in its abstract form as: For \(\boldsymbol{\eta}\in \boldsymbol{H}\), find \(u(\boldsymbol{x}, \boldsymbol{\eta}) \colon \Omega \times \boldsymbol{H}\to \mathbb{R}\), such that \[\label{eq:pde} \begin{align} \pazocal{A}(u(\boldsymbol{x}, \boldsymbol{\eta}), \boldsymbol{\eta}) & = f(\boldsymbol{x}, \boldsymbol{\eta}), \quad \; \boldsymbol{x}\in \Omega, \\ \pazocal{B}^{k} (u(\boldsymbol{x}, \boldsymbol{\eta}), \boldsymbol{\eta}) & = g^{k}(\boldsymbol{x}, \boldsymbol{\eta}), \quad \boldsymbol{x}\in \Gamma^{k} \subseteq \partial \Omega, \quad \text{ for } k=1, 2, \ldots, n_{\Gamma}, \end{align}\tag{1}\] where \(\pazocal{A}\) and \(\{\pazocal{B}\}_{k=1}^{n_{\Gamma}}\) denote a nonlinear differential operator and a set of boundary condition operators¹, respectively. Moreover, the right hand side function \(f(\boldsymbol{x}, \boldsymbol{\eta})\) and the functions \(\{g^{k}(\boldsymbol{x}, \boldsymbol{\eta})\}_{k=1}^{n_{\Gamma}}\) are given.

2.1 Physics informed neural networks↩︎

PINNs [5] leverage empirical data and knowledge of physical priors to solve PDEs using DNNs. Here, we focus on solving one particular instance of 1 with a preselected set of parameters \(\boldsymbol{\eta}\). A solution \(u(\boldsymbol{x}, \boldsymbol{\eta})\) of a PDE for a given \(\boldsymbol{\eta}\) at any \(\boldsymbol{x}\in \Omega\) is approximated by the DNN \(u_{\boldsymbol{\theta}} \colon \mathbb{R}^{d} \to \mathbb{R}\), which is parameterized with respect to parameters \(\boldsymbol{\theta}\in \mathbb{R}^n\). An actual form of the DNN depends on the user-specified neural network architecture. For instance, \(u_{\boldsymbol{\theta}}\) can be represented by a feed-forward network or a residual network [43].

In order to ensure that the network \(u_{\boldsymbol{\theta}}\) approximates the PDE, we have to find the optimal set of DNN’s parameters \(\boldsymbol{\theta}^{\ast}\). To this aim, we minimize the loss \(\pazocal{L}\colon \mathbb{R}^n \to \mathbb{R}\) associated with a residual of a given PDE and boundary conditions, i.e., \[\label{eq:pinn} \boldsymbol{\theta}^{\ast} = \mathop{\mathrm{arg\,min}}_{\boldsymbol{\theta}\in \mathbb{R}^n} \pazocal{L}(\boldsymbol{\theta}) := w_{\mathrm{int}}\pazocal{L}_{\mathrm{int}}(\boldsymbol{\theta}) + w_{\mathrm{bc}}\pazocal{L}_{\mathrm{bc}}(\boldsymbol{\theta}),\tag{2}\] where \[\label{eq:pinn2} \begin{align} \pazocal{L}_{\mathrm{int}}(\boldsymbol{\theta}) & := \Vert \pazocal{A}(u_{\boldsymbol{\theta}}) - f \Vert^{2}_{L^{2}(\Omega)} = \int_{\Omega} (\pazocal{A}(u_{\boldsymbol{\theta}}(\boldsymbol{x}))-f(\boldsymbol{x}))^{2} d \boldsymbol{x}, \\ \pazocal{L}_{\mathrm{bc}}(\boldsymbol{\theta}) & := \Vert \pazocal{B}(u_{\boldsymbol{\theta}}) - g \Vert^{2}_{L^{2}(\partial \Omega)} = \int_{\partial \Omega} (\pazocal{B}(u_{\boldsymbol{\theta}}(\boldsymbol{x}))-g(\boldsymbol{x}))^{2} d \boldsymbol{x}. \end{align}\tag{3}\]

To discretize the interior loss \(\pazocal{L}_{\mathrm{int}}\), we consider a set of points \(\pazocal{D}_{\mathrm{int}}=\{ \boldsymbol{x}_{j}\}^{n_{\mathrm{int}}}_{j=1}\) sampled from the interior of the domain \(\Omega\). Similarly, in order to discretize the boundary loss \(\pazocal{L}_{\mathrm{bc}}\), we construct \(\{\pazocal{D}^{k}_{\mathrm{bc}}\}_{k=1}^{n_{\Gamma}}\), where each point of \(\pazocal{D}^{k}_{\mathrm{bc}}=\{ \boldsymbol{x}^{k}_{j} \}^{n_{\mathrm{bc}}}_{j=1}\) is sampled from the boundary \(\Gamma^k\). The continuous losses 3 can be then discretized, for example, by using the Monte-Carlo method, giving rise to \[\begin{align} \pazocal{L}_{\mathrm{int}}(\boldsymbol{\theta}) & = \frac{1}{\vert \pazocal{D}_{\mathrm{int}}\vert} \sum_{\boldsymbol{x}_{j} \in \pazocal{D}_{\mathrm{int}}} (\pazocal{A}(u_{\boldsymbol{\theta}}(\boldsymbol{x}_{j}))-f(\boldsymbol{x}_{j}))^{2}, \\ \pazocal{L}_{\mathrm{bc}}(\boldsymbol{\theta}) & = \sum_{k=1}^{n_{\Gamma}} \left( \frac{1}{\vert \pazocal{D}^{k}_{\mathrm{bc}}\vert} \sum_{\boldsymbol{x}^{k}_{j} \in \pazocal{D}^{k}_{\mathrm{bc}}} (\pazocal{B}^{k} (u_{\boldsymbol{\theta}}(\boldsymbol{x}^{k}_{j}))-g^{k}(\boldsymbol{x}^{k}_{j}))^{2} \right). \end{align}\]

The weights \(w_{\mathrm{int}}\) and \(w_{\mathrm{bc}}\) in 2 are used to balance the interior and boundary loss terms. In practice, the quality of the DNN approximation and the success of the training process depends on the appropriate selection of these weights [44]. Consequently, many approaches for manual or automatic selection of weights have been proposed in the literature; see, for instance, NTK-based weighting strategies [45] or residual-based attention mechanism [46]. In this work, Dirichlet boundary conditions are directly embedded into the network architecture, thereby avoiding explicit selection of \(w_{\mathrm{int}}\) and \(w_{\mathrm{bc}}\). Following [44], [47], we modify the network output as \[\label{eq:exact} \tilde{u}_{\boldsymbol{\theta}}(\boldsymbol{x}) := \sum_{k=1}^{n_{\Gamma}} g^{k}(\boldsymbol{x}) + \ell(\boldsymbol{x}) u_{\boldsymbol{\theta}}(\boldsymbol{x}), \quad \text{ where } \ell(\boldsymbol{x}) :=\prod_{k=1}^{n_{\Gamma}}\ell^{k}(\boldsymbol{x}).\tag{4}\] Each function \(\ell^{k} \colon \Omega \to [0,1]\) is defined such that \[\ell^{k}(\boldsymbol{x}) = \begin{cases} 0, & \boldsymbol{x}\in \Gamma^{k}, \\ 1, & \boldsymbol{x}\in \Gamma^{j}, \quad j \in \{ 1,2,\ldots,n_{\Gamma} \}, \quad \text{ for } k \ne j, \\ t \in (0,1), & \text{otherwise}. \end{cases}\] Note that \(t\) is determined so that \(\ell^{k}(\boldsymbol{x})\) is a smooth function. For example, when \(\Omega = (0,1)^{2}\) and \(\partial \Omega = \cup_{k=1}^{n_{\Gamma}} \Gamma^{k}\), \(\ell(\boldsymbol{x})=\ell(x_{1}, x_{2})=x_{1}(1-x_{1})x_{2}(1-x_{2})\). Using \(\tilde{u}_{\boldsymbol{\theta}}\) as defined in 4 , we can now reformulate the minimization problem 2 as \[\boldsymbol{\theta}^{\ast} = \mathop{\mathrm{arg\,min}}_{\boldsymbol{\theta}\in \mathbb{R}^n} \pazocal{L}(\boldsymbol{\theta}) := \frac{1}{\vert \pazocal{D}_{\mathrm{int}} \vert} \sum_{\boldsymbol{x}_{j} \in \pazocal{D}_{\mathrm{int}}} (\pazocal{A}(\tilde{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_{j})) - f(\boldsymbol{x}_{j}) )^{2}. \label{eq:loss95pin95exact95bc}\tag{5}\]

2.2 Operator learning↩︎

Operator learning approaches use DNNs to approximate a mapping between infinite-dimensional function spaces. Let \(\pazocal{Y}\) be an infinite-dimensional Banach space, which can represent a space of parametrized right-hand sides, boundary conditions, or material parameters. The goal is to approximate the operator \(\pazocal{G}\colon \pazocal{Y}\to \pazocal{V}\), where \(\pazocal{V}\) denotes the solution space of the parametric PDE, specified by 1 . In this work, we approximate the operator \(\pazocal{G}\) using the DON, initially proposed in [8].

The architecture of DON consists of branch and trunk networks. The branch network \(B \colon \pazocal{Y}_{m} \to \mathbb{R}^{p}\) provides a set of \(p\) coefficients for a given discretized function \({\boldsymbol{y}^{m} \in \pazocal{Y}_{m}}\). Here, \(\pazocal{Y}_{m}\) is a \(m\)-dimensional space that is assumed to satisfy the following condition: \[\forall y \in \pazocal{Y}, \; \exists \boldsymbol{y}^{m} \in \pazocal{Y}_{m}\subseteq \pazocal{Y} \text{ such that } \boldsymbol{y}^{m} \to y \text{ as } m \to \infty.\] For simplicity, we write \(\boldsymbol{y}\) in place of \(\boldsymbol{y}^{m}\), as \(m\) is a fixed and determines the dimension of the discrete space. In addition, the trunk network \(T \colon \Omega \to \mathbb{R}^{p}\) is used to encode the computational domain \(\Omega\). Therefore, it takes as an input a coordinate \(\boldsymbol{\xi}\in \Omega\) and returns a set of \(p\) basis functions.

The coefficients \(B_{k}(\boldsymbol{y}) \in \mathbb{R}^p\) and the basis functions \(T_{k}(\boldsymbol{\xi}) \in \mathbb{R}^p\) obtained as an output of branch and trunk networks are then combined as \[\pazocal{G}_{\boldsymbol{\theta}} (\boldsymbol{y}) (\boldsymbol{\xi}) := \sum_{k=1}^{p} B_{k}(\boldsymbol{y}) T_{k}(\boldsymbol{\xi}), \quad \boldsymbol{y}\in \pazocal{Y}_{m}, \boldsymbol{\xi}\in \Omega,\] to approximate a solution of the parametric PDE at a point \(\boldsymbol{\xi}\) for a given function \(\boldsymbol{y}\). Figure 1 illustrates a prototypical DON architecture. Note that we are free to choose the architecture of branch and trunk networks. For example, the choice of the branch network can be guided by the structure of the input \(\boldsymbol{y}\), e.g., it can be a fully connected network in case of a scalar input, or a convolutional network in case of multidimensional discretized functions. Since the vector of coordinates \(\boldsymbol{\xi}\in \mathbb{R}^d\) is usually low dimensional, fully connected networks are commonly used in place of the trunk network. Alternatively, one can utilize the precomputed basis functions, e.g., obtained by means of the proper orthogonal decomposition (POD) performed on the training data [48] or by means of parameterizing the integral kernel in the Fourier space [49].

To train DON such that it approximates solutions of 1 , we construct a dataset \[\pazocal{D}= \{ (\boldsymbol{y}_{j}, \bar{\boldsymbol{\xi}}_{j}, \boldsymbol{u}_{j}) \}_{j=1}^{n_{s}},\] consisting of \(n_{s}\) samples. Each triple \((\boldsymbol{y}_{j}, \bar{\boldsymbol{\xi}}_{j}, \boldsymbol{u}_{j})\) consists of the discretized function \({\boldsymbol{y}_{j} \in \mathbb{R}^m}\), sampled from \(\pazocal{Y}_{m}\). Moreover, \(\bar{\boldsymbol{\xi}}_{j} = [ \boldsymbol{\xi}_{j,1}, \ldots, \boldsymbol{\xi}_{j, n_{c}} ]^T \in \mathbb{R}^{n_{c} \times d}\) represents a set of \(n_{c}\) collocation points and \(\boldsymbol{u}_{j} \in \mathbb{R}^{n_{c}}\) is a target solution evaluated at collocation points constituting \(\bar{\boldsymbol{\xi}}_{j}\). Note that the target solution is typically obtained using high-fidelity discretization techniques, such as finite element or finite difference methods.

Using dataset \(\pazocal{D}\), the optimal parameters of the DON, denoted by \(\boldsymbol{\theta}^{\ast} \in \mathbb{R}^n\), are found by solving the following minimization problem: \[\label{eq:donloss} \boldsymbol{\theta}^{\ast} = \mathop{\mathrm{arg\,min}}_{\boldsymbol{\theta}\in \mathbb{R}^n} \pazocal{L}(\boldsymbol{\theta}) := \frac{1}{n_{s} n_{c}} \sum_{j=1}^{n_{s}} \sum_{t=1}^{n_{c}} \left(\pazocal{G}_{\boldsymbol{\theta}}(\boldsymbol{y}_{j})(\boldsymbol{\xi}_{j,t}) - \boldsymbol{u}_{j}\right)^{2}.\tag{6}\] Note that instead of using \(\boldsymbol{u}_{j}\) given by the dataset, we can utilize the PINN-based loss function; see [9] for details. This might simplify the generation of the dataset but makes the training process more tedious. We do not pursue this possibility here but rather investigate the performance of the proposed two-level preconditioners in data-driven settings. Notably, the proposed preconditioner is expected to enhance the performance of standard optimizers more significantly in the case of PINN-based loss due to increased ill-conditioning of the training problem [9].

3 Two-level overlapping additive Schwarz preconditioner↩︎

In this work, we focus on training the DNN by solving the following minimization problem: \[\begin{align} \min_{\boldsymbol{\theta}} \pazocal{L}(\boldsymbol{\theta}), \label{eq:min95problem} \end{align}\tag{7}\] where the loss \(\pazocal{L}\) is as defined in 5 or 6 for PINN or DON, respectively. The minimization problem 7 is traditionally solved using gradient-based optimizers such as Adam [11] or LBFGS [12]. However, the convergence of these optimizers tends to deteriorate as problem stiffness and ill-conditioning increase [50], resulting in slow training and unsatisfactory accuracy of the resulting PINN/DON models. For instance, 1 illustrates that increasing the number of layers in a PINN model exacerbates the relative error when using a standard optimizer. As the depth of the network grows, the ill-conditioning and non-convexity of the underlying minimization problem become more significant. This, in turn, presents significant challenges for the standard LBFGS method to converge, leading to poorer accuracy of the PINN as the representation capacity of the DNN increases.

In scientific computing, Schwarz methods [24], [51] provide an efficient framework for solving large-scale problems by decomposing the solution space into smaller subspaces, within which the problem can be solved more effectively. Emulating the classical additive Schwarz methods, the single-level Schwarz preconditioning strategy, which utilizes the layer-wise decomposition of the DNN, was introduced in [30]. In this work, we propose enhancing the performance of this strategy by introducing overlap, a coarse-level correction step, and a subdomain-based recombination strategy. The culmination of these techniques results in a novel two-level overlapping Schwarz preconditioner, the performance of which does not deteriorate with increasing network depth, as shown in 1.

To present the proposed two-level overlapping Schwarz preconditioning strategy, let us consider the nonlinear systems of equations associated with the first-order optimality conditions of 7 given as \[\label{eq:fcond} \nabla \pazocal{L}(\boldsymbol{\theta}) = 0.\tag{8}\] Since solving 8 is computationally demanding, we instead construct and solve the following nonlinearly-preconditioned system of equations: \[\label{eq:precond} \pazocal{F}(\boldsymbol{\theta}) := \nabla \pazocal{L}(\pazocal{M}(\boldsymbol{\theta})) = 0,\tag{9}\] where the preconditioner \(\pazocal{M} \colon \mathbb{R}^{n} \to \mathbb{R}^{n}\) is designed such that 9 has more balanced nonlinearities than 8 . In practice, this means that the nonlinear terms are of comparable magnitude both to each other and to the linear terms, thereby avoiding dominance by any single contribution.

In this work, we solve the nonlinear system 9 iteratively. Thus, on each \(k\)-th iteration, we start from a given iterate \(\boldsymbol{\theta}^{(k)}\), and apply the preconditioner \(\pazocal{M}\). As a result, we obtain the intermediate iterate \(\boldsymbol{\theta}^{(k+1/2)} := \pazocal{M}(\boldsymbol{\theta}^{(k)})\), which is required to reduce the loss, i.e., it has to satisfy the following condition: \[\pazocal{L}(\boldsymbol{\theta}^{(k+1/2)}) \leq \pazocal{L}(\boldsymbol{\theta}^{(k)}).\] The smoothing step is followed by a global optimization step. To this aim, we apply an operator \(\pazocal{K}\colon \mathbb{R}^n \to \mathbb{R}^n\), which is associated with one or multiple iterations of a user-specified global optimizer. The operator \(\pazocal{K}\) takes as an input the gradient of the loss function and provides a direction \(\boldsymbol{p}^{(k)}\), i.e., \({\boldsymbol{p}^{(k)} := \pazocal{K}(\nabla \pazocal{L}(\boldsymbol{\theta}^{(k+1/2)}))}\). For instance, if a single step of the LBFGS optimizer is employed, then the operator \(\pazocal{K}\) is defined as \[\boldsymbol{p}^{(k)} := \pazocal{K}\big(\nabla \pazocal{L}( \boldsymbol{\theta}^{(k+1/2)}) \big) := - \bigg( \boldsymbol{\mathrm{H}}^{(k+1/2)} \bigg)^{-1} \nabla \pazocal{L}(\boldsymbol{\theta}^{(k+1/2)}),\] where \(\boldsymbol{\mathrm{H}}^{(k+1/2)}\) denotes the Hessian approximation of \(\pazocal{L}\) at \(\boldsymbol{\theta}^{(k+1/2)}\). Afterward, the updated global iterate \(\boldsymbol{\theta}^{(k+1)}\) is obtained as \[\boldsymbol{\theta}^{(k+1)} = \boldsymbol{\theta}^{(k+1/2)} + \alpha^{(k)}\boldsymbol{p}^{(k)},\] where \(\alpha^{(k)}\) denotes the step size, determined using the line-search globalization strategy.

3.1 Overlapping layer-wise additive Schwarz preconditioner↩︎

The computational efficiency associated with solving the preconditioned system of equations 9 is directly linked to the quality and the computational cost of the nonlinear preconditioner \(\pazocal{M}\). In this work, we construct the preconditioner \(\pazocal{M}\) by utilizing the overlapping layer-wise network decomposition. Thus, we decompose the DNN, consisting of \(N\) layers, into \(N_{sd}\) subnetworks (subdomains), such that each layer belongs to a unique subnetwork. The parameters of the network are also decomposed into \(N_{sd}\) disjoint groups, i.e., \[\boldsymbol{\theta}= [\boldsymbol{\theta}_{1}, \ldots, \boldsymbol{\theta}_{s}, \ldots, \boldsymbol{\theta}_{N_{sd}}]^{T},\] where \(\boldsymbol{\theta}_{s} := [\boldsymbol{\theta}_{s, 1}, \ldots, \boldsymbol{\theta}_{s, N_s}]^T \in \mathbb{R}^{n_{s}}\), for \(s = 1, \ldots, N_{sd}\). Each \(s^{th}\) non-overlapping subnetwork consists of \(N_s\) layers, with parameters of \(h^{th}\) layer being denoted as \(\boldsymbol{\theta}_{s, h}\). In the case of the overlapping decomposition, the parameters \(\boldsymbol{\theta}_{s}\) of the \(s\)-th subnetwork are concatenated with the parameters of \(\delta(\in \mathbb{N})<N_S\) neighboring layers from the preceding \((s-1)\)-th and succeeding \((s+1)\)-th subnetworks. For \(\delta\geq 1\), this gives rise to overlapping set of subdomain parameters \[\begin{align} \tilde{\boldsymbol{\theta}}_{s} = [{\boldsymbol{\theta}}_{s-1, (N_{sd}-\delta)}, \ldots, {\boldsymbol{\theta}}_{s-1, N_{sd}}, {\boldsymbol{\theta}}_{s, 1}, \ldots, {\boldsymbol{\theta}}_{s, N_{sd}}, {\boldsymbol{\theta}}_{s+1, 1}, \ldots, {\boldsymbol{\theta}}_{s+1, \delta}]^{T} \in \mathbb{R}^{\tilde{n}_{s}}, \end{align}\] where \(\tilde{N}_{s}\) and \(\tilde{n}_{s}\) denotes the number of layers in overlapping subnetwork and the number of associated subnetwork parameters, respectivelly. Note that, for the first (\(s=1\)) and the last (\(s = N_{sd}\)-th) subnetworks, we only pick layers from the second and \((N_{sd}-1)\)-th subnetworks, respectively. 2 illustrates the parameter selection for the second subnetwork of the six-layer network decomposed in a layer-wise manner into three subdomains with an overlap of size one, i.e., \(N_{sd}=3\) and \(\delta=1\).

To transfer the data between the subnetworks and the global network, we define a restriction operator \(\boldsymbol{R}_{s} \colon \mathbb{R}^{n} \to \mathbb{R}^{\tilde{n}_{s}}\), which extracts the parameters of the overlapping subnetwork from the parameters of the global network, i.e., \[\tilde{\boldsymbol{\theta}}_{s} = \boldsymbol{R}_{s} \boldsymbol{\theta}.\] Moreover, we also consider the prolongation operator \(\boldsymbol{R}^{T}_{s} \colon \mathbb{R}^{\tilde{n}_{s}} \to \mathbb{R}^{n}\), defined as adjoint of \(\boldsymbol{R}_{s}\). The role of the prolongation operator \(\boldsymbol{R}^{T}_{s}\) is to extend the parameters associated with \(s^{th}\) subnetwork to the global network, i.e., \(\boldsymbol{\theta}= \sum_{s=1}^{N_{sd}} \boldsymbol{R}^T_{s} \tilde{\boldsymbol{\theta}}_{s}\).

3.1.1 Single-level preconditoner with layer-wise synchronization strategy↩︎

Using the proposed overlapping network decomposition, we define the local training problems for all subdomains as \[\label{eq:local} \tilde{\boldsymbol{\theta}}^{\ast}_{s} := \mathop{\mathrm{arg\,min}}_{\tilde{\boldsymbol{\theta}}_{s} \in \mathbb{R}^{n_s}} \pazocal{L}(\tilde{\boldsymbol{\theta}}_{s} + (\boldsymbol{\mathrm{I}}- \boldsymbol{R}^{T}_{s}\boldsymbol{R}_{s})\boldsymbol{\theta}),\tag{10}\] where \(s \in 1, \ldots, N_s\). Thus, for the \(s^{th}\) subdomain, we minimize the loss functional \(\pazocal{L}\) with respect to parameters \(\tilde{\boldsymbol{\theta}}_{s}\), while all other parameters are held fixed. The minimization problem 10 is solved only approximately, for example, by utilizing a few steps of the LBFGS or Adam optimizer. The resulting parameters \(\tilde{\boldsymbol{\theta}}^{\ast}_{s}\) are then used to define the (single-level) overlapping additive Schwarz preconditioner \(\pazocal{M}_{SL}\) as \[\label{eq:onelevel} \pazocal{M}_{SL}(\boldsymbol{\theta}) := \boldsymbol{\theta}+ \gamma \sum_{s=1}^{N_{sd}} \boldsymbol{R}^{T}_{s}(\tilde{\boldsymbol{\theta}}^{\ast}_{s} - \boldsymbol{R}_{s} \boldsymbol{\theta}) =: \boldsymbol{\theta}+ \gamma \sum_{s=1}^{N_{sd}} \boldsymbol{d}_{s} =: \boldsymbol{\theta}+ \gamma \boldsymbol{d}.\tag{11}\] Here, the symbol \(\gamma \in [0, 1]\) denotes an appropriately selected damping parameter, such that \[\pazocal{L}(\boldsymbol{\theta}+ \gamma \boldsymbol{d}) \leq \pazocal{L}(\boldsymbol{\theta}).\]

It is crucial to note that the convergence rate of the preconditioner \(\pazocal{M}_{SL}\) is strongly affected by the value of \(\gamma\) as well as by the quality of the search direction \(\boldsymbol{d}\). In the literature on the linear additively composed preconditioners, the multipreconditioning approach [52], [53] is often employed to determine how to optimally combine locally obtained search directions \(\{\boldsymbol{d}_{s}\}_{s=1}^{N_{sd}}\) to improve the quality of the resulting search direction \(\boldsymbol{d}\). Here, we extend this methodology into nonlinear settings by searching for a set of damping parameters \(\{\gamma^*_{s}\}_{s=1}^{N_{sd}}\), such that \[\begin{align} \pazocal{L}\left(\boldsymbol{\theta}+ \sum_{s=1}^{N_{sd}}\gamma^*_{s} \boldsymbol{d}_{s} \right) < \pazocal{L}\left(\boldsymbol{\theta}+ \sum_{s=1}^{N_{sd}} \gamma_{s} \boldsymbol{d}_{s} \right),\quad \forall \gamma^*_s \neq \gamma_s, \end{align}\] where \(\gamma_{s} \in [0, 1]\) is an initial damping parameter associated with the local, subdomain-based, search direction \(\boldsymbol{d}_{s}\).

The optimal set of parameters \(\{\gamma^*_{s}\}_{s=1}^{N_{sd}}\) can be determined by solving the following \(N_{sd}\)-dimensional minimization problem: \[\begin{align} \min_{\gamma_{1}, \ldots, \gamma_{N_{sd}}} \pazocal{L}\left(\boldsymbol{\theta}+ \sum_{s=1}^{N_{sd}}\gamma_{s} \boldsymbol{d}_{s} \right). \label{eq:gammas95min95problem} \end{align}\tag{12}\] However, finding a solution of 12 is computationally demanding due to the nonlinear nature of the problem. This is in contrast to multi-preconditioning approaches used in the context of linear problems, where \(\{\gamma^*_{s}\}_{s=1}^{N_{sd}}\) can typically be obtained in closed form, see for instance [52], [54].

To elevate the computational burden of solving 12 exactly, one can employ multidimensional line-search heuristics, such as multidimensional backtracking proposed in [55]. Here, we opt for a simplified approach and leverage the fact that the majority of DNNs construct their bases by propagating the input features through the network from the first to the last layer in a sequential manner, which is frequently regarded as a time-dependent process [29], [31]. Therefore, we also follow the feature flow and determine \(\{\gamma_{s}^*\}_{s=1}^{N_{sd}}\) in a sequential manner, such that \[\pazocal{L}\left( \boldsymbol{\theta}+ \sum_{j=1}^{s-1} \gamma_{j}^* \boldsymbol{d}_{j} + \gamma_{s} \boldsymbol{d}_{s} \right) \leq \pazocal{L}\left( \boldsymbol{\theta}+ \sum_{j=1}^{s-1} \gamma_{j}^* \boldsymbol{d}_{j} \right),\] for \(s=1, \ldots, N_{sd}\). Thus, the values \(\{\gamma^*_{s}\}_{s=1}^{N_{sd}}\) are determined from the first subdomain to the last one, utilizing a standard, unidimensional, line-search strategy. This guarantees that each search direction \(\boldsymbol{d}_{s}\) decreases the value of \(\pazocal{L}\) associated with a global network. At the same time, we ensure that the DNN’s bases are constructed by appropriately adjusting the impact of each search direction \(\boldsymbol{d}_{s}\) on the global DNN’s approximation. Here, we remark that the search directions \(\{\boldsymbol{d}_{s}\}_{s=1}^{N_{sd}}\) are obtained in parallel, while only scalar damping parameters \(\{\gamma^*_{s}\}_{s=1}^{N_{sd}}\) are obtained sequentially during the subdomain synchronization step.

3.2 Two-level overlapping additive Schwarz preconditioner↩︎

In the field of numerical methods, the convergence of single-level preconditioners is often enhanced by incorporating a coarse level, which ensures a global communication between subdomains. Following this methodology, we propose to improve the convergence of a single-level layer-wise preconditioner by incorporating a coarse-level training step. Our design of the coarse-level network is inspired by the observation that the forward pass through the network can be interpreted as a time-stepping process [29], [31], [56]–[58]. This observation suggests that coarse-level subnetworks can be created by reducing the number of layers within a given subdomain, analogous to time-step coarsening. More precisely, we create the coarse-level subnetwork by grouping together the first layers of each subdomain. Thus, the trainable coarse-level parameters are given as \[\begin{align} \boldsymbol{\theta}_{0} = [\boldsymbol{\theta}_{1, 1}, \ldots, \boldsymbol{\theta}_{s, 1}, \ldots, \boldsymbol{\theta}_{N_{sd}, 1} ]^T, \end{align}\] where \(\boldsymbol{\theta}_{s, 1}\) is extracted from \(\boldsymbol{\theta}_{s} = [\boldsymbol{\theta}_{s, 1}, \ldots, \boldsymbol{\theta}_{s, \tilde{N}_s}]^T\), for all \(s=1, \ldots, N_{sd}\). An illustration of the coarse-level parameter selection for a six-layer network with \(N_{sd}=3\) and \(\delta=0\) is depicted in 3. The computational cost of the coarse-level problem is proportional to the number of subdomains \(N_{sd}\).

The transfer of the data between the coarse-level subnetwork and the global network, is ensured by two transfer operators. In particular, we consider the restriction operator \(\boldsymbol{R}_{0} \colon \mathbb{R}^{n} \to \mathbb{R}^{n_{0}}\), given such that \[\boldsymbol{\theta}_{0} = \boldsymbol{R}_{0} \boldsymbol{\theta}.\] The prolongation operator \(\boldsymbol{R}^{T}_{0} \colon \mathbb{R}^{n_{0}} \to \mathbb{R}^{n}\) is defined as the adjoint operator of \(\boldsymbol{R}_{0}\). Note that the number of trainable coarse-level parameters \(n_{0}\) scales proportionally to the number of subdomains \(N_{sd}\).

Using the coarse-level subnetwork, we can now define the coarse-level training problem as \[\boldsymbol{\theta}^{\ast}_{0} := \mathop{\mathrm{arg\,min}}_{\boldsymbol{\theta}_{0} \in \mathbb{R}^{n_0}} \pazocal{L}(\boldsymbol{\theta}_{0} + (\boldsymbol{\mathrm{I}}- \boldsymbol{R}^{T}_{0}\boldsymbol{R}_{0})\boldsymbol{\theta}).\] Hence, the loss functional \(\pazocal{L}\) is minimized with respect to the parameters associated with the coarse-level subnetwork while all other parameters are held fixed. The outcome of the coarse-level minimization process, the parameters \(\boldsymbol{\theta}^{\ast}_{0} \in \mathbb{R}^{n_0}\), is then used to update the parameters of the global network as \[\label{eq:coarse} \hat{\boldsymbol{\theta}} = \boldsymbol{\theta}+ \boldsymbol{R}^{T}_{0} (\boldsymbol{\theta}^{\ast}_{0} - \boldsymbol{R}_{0}\boldsymbol{\theta}).\tag{13}\] Note that since all coarse-level parameters \(\boldsymbol{\theta}_{0}\) are optimized simultaneously, the coarse-level training process is serial. Moreover, the condition \(\pazocal{L}(\hat{\boldsymbol{\theta}}) \leq \pazocal{L}(\boldsymbol{\theta})\) is automatically fulfilled.

Using 13 , we are now ready to define a two-level overlapping additive Schwarz preconditioner \(\pazocal{M}_{TL}: \mathbb{R}^{n} \to \mathbb{R}^n\), by multiplicatively composing the coarse-level training step with an application of \(\pazocal{M}_{SL}\). In particular, \(\pazocal{M}_{TL}\) is defined by applying \(\pazocal{M}_{SL}\) after the coarse-level update 13 has been performed, i.e., \[\pazocal{M}_{TL} (\boldsymbol{\theta}) := \pazocal{M}_{SL} ( \underbrace{\boldsymbol{\theta}+ \boldsymbol{R}^{T}_{0} (\boldsymbol{\theta}^{\ast}_{0} - \boldsymbol{R}_{0}\boldsymbol{\theta})}_{:=\hat{\boldsymbol{\theta}}}).\] Note that since the coarse-level training indirectly enforces the feed-forward flow of the features through the network, the parameters \(\hat{\boldsymbol{\theta}}\) serve as an improved initial guess for the parallel subdomain training phase. 4 provides a graphical illustration of the one-level and two-level Schwarz methods when applied to train the DNNs. In addition, 5 summarizes the proposed two-level preconditioning process for a user-selected optimizer \(\pazocal{K}\). Our practical implementation of the line-search step (Step 9 in 5) employs cubic backtracking with the Strong-Wolfe conditions [59].

4 Benchmark problems and implementation details↩︎

This section describes the benchmark problems used to evaluate the performance of the proposed two-level preconditioner. In particular, we consider three benchmark problems associated with training PINNs and three benchmark problems with training DONs. Additionally, we discuss the implementation aspects and configuration of all considered optimizers.

4.1 Benchmark problems for training of PINNs↩︎

For training PINNs, we consider the following benchmark problems.

• Burgers’ equation (Burg): Burgers’ equation is given in the following form: \[\begin{align} \frac{\partial u}{\partial t} + u u' - \nu u'' & = 0, \quad \quad \quad & & \forall \;(t,x) \in (0,1] \times (-1,1), \\ u & = - \sin(\pi x), \; & & \forall \;(t,x) \in \{0\} \times [-1,1], \\ u & = 0, \; & & \forall \;(t,x) \in (0,1] \times \{1\}, \\ u & = 0, \; & & \forall \;(t,x) \in (0,1] \times \{-1\}, \end{align}\] where \(u=u(t,x)\) denotes the flow velocity and \(\nu\) stands for the kinematic viscosity. Here, we choose \(\nu=0.01/\pi\).

• Diffusion-advection (DA): The diffusion-advection equation is given as \[\begin{align} -\mu \Delta u + \boldsymbol{b}\cdot \nabla u & = f, & & \forall \;(x_1,x_2) \in (0, 1) \times (0, 1), \\ u & = 0, & & \text{on } \partial \Omega, \end{align}\] where \(\boldsymbol{b}=(1,1)^\top\). The right-hand side is considered to be constant, i.e., \(f = 1\). Moreover, the symbol \(\mu\) denotes viscosity, which we choose to set as \(\mu=10^{-2}\).

• Allen-Cahn (AC): We consider the Allen-Cahn equation of the following form: \[\begin{align} \frac{\partial u}{\partial t} - D u'' - 5 (u - u^3) & = 0, \quad \quad \quad & & \forall \;(t,x) \in (0,1] \times (-1,1), \\ u & = x^2 \cos(\pi x), \; & & \forall \;(t,x) \in \{0\} \times [-1,1], \\ u & = - 1, \; & & \forall \;(t,x) \in (0,1] \times \{-1\}, \\ u & = -1, \; & & \forall \;(t,x) \in (0,1] \times \{1\}, \end{align}\] where the diffusion coefficient \(D\) is chosen as \(D=0.001\).

The amount and location of collocation points significantly influence the discretization error of PINNs. Here, we perform sampling by employing the Quasi-Monte Carlo method [60], with Hammersley low-discrepancy sequences [7]. As a consequence, it is well known that the convergence rate of Quasi-Monte Carlo is close to \(O(N^{-1})\), while the convergence rate of the standard Monte Carlo method, with pseudo-random sequences, is \(O(N^{-1/2})\) [61]. The number of collocation points is \(10{,}000\) for all three examples.

4.2 Benchmark problems for training of DONs↩︎

For the training of DONs, we consider the three benchmark problems, with datasets created using the finite element method within the Firedrake library [62].

• Anisotropic Poisson equation (Aniso): Let \(\Omega := (-1,1)^2\) be the computational domain, with boundary \(\partial \Omega\), the anisotropic Poisson equation is given as \[\begin{align} - \nabla \cdot (K(x, \boldsymbol{\eta}) \;\nabla u(x) ) & = f(x), \quad \quad & & \forall \;x \in \Omega, \\ u & = 0, \; & & \forall \;x \in \partial \Omega, \\ \end{align} \label{eq:aniso95laplace}\tag{14}\] where \(u\) denotes the solution and \(f(x):=1\) is the forcing term. The diffusion coefficient \(K(x, \boldsymbol{\eta})\) is considered to be an anisotropic tensor of the following form: \[\begin{align} K(x, \boldsymbol{\eta}) = \begin{pmatrix} \cos(\alpha) & -\sin(\alpha) \\ \sin(\alpha) & \cos(\alpha) \end{pmatrix} \begin{pmatrix} 1 & 0 \\ 0 & \beta \end{pmatrix} \begin{pmatrix} \cos(\alpha) & \sin(\alpha) \\ -\sin(\alpha) & \cos(\alpha) \end{pmatrix}, \end{align}\] where \(\boldsymbol{\eta}:=[\alpha, \beta]\). The parameter \({\beta \in [10^{-6},1]}\) specifies the anisotropic strength, and it is sampled from the following distribution: \({\log_{10} (\frac{1}{\beta}) \sim \pazocal{U}[0, 6]}\). The parameter \({\alpha \in (0,\pi)}\), sampled as \(\alpha \sim \pazocal{U}[0, \pi]\), is used to specify the angle of anisotropic direction.

  To construct the dataset, for each \(\boldsymbol{\alpha}\), we discretize 14 using a uniform mesh with \(34 \times 34\) elements. The number of training samples and test samples are \(4{,}250\) and \(750\), respectively. We have uploaded this new dataset to the public ZENODO repository [63].

  

  

  

  

  

• Helmholtz equation (Helm): Let \(\Omega\) be the computational domain, defined as an annular cylinder depicted on 6. The parametrized Helmholtz equation is given as \[\begin{align} - \Delta u(x) - k_{\text{H}}^2 u(x) & = f(x, \boldsymbol{\eta}), \quad \quad \quad & & \forall \;x \in \Omega, \\ u & = 0, \; & & \forall \;x \;\text{on} \;\Gamma, \\ \frac{\partial{u}}{\partial n} & = 0, \; & & \forall \;x \;\text{on} \;\partial \Omega \setminus\Gamma, \\ \end{align} \label{eq:helmholtz}\tag{15}\] where \(u\) is the solution and \(\Gamma\) denotes the top boundary of the cylinder. Moreover, the symbol \(k_{\text{H}}\) denotes a constant wave number, which we set to \(k_{\text{H}}=3\). We sample the forcing term from Gaussian random fields (GRFs) with mean \(\mathbb{E}[f(x, \boldsymbol{\eta})]=0\), and the covariance is given as \[\begin{align} \text{Cov}(f(x,\boldsymbol{\eta}), f(y, \boldsymbol{\eta})) = \sigma^2 \exp \bigg(- \frac{\| x - y \|^2}{2 \ell^2} \bigg), \end{align}\] with \(\sigma = 3.0\) and \(\ell = 0.1\). To construct the dataset, we discretize 15 using an unstructured mesh with \(1,503\) nodes. The sampled right-hand sides are then interpolated onto a \(32 \times 32 \times 32\) grid composed of elements within the bounding box \([-1, 1] \times [-1, 1] \times [0, 1]\), generating an input that can be processed by the 3D convolutional layers of the branch network. The number of training samples and test samples are \(4{,}250\) and \(750\), respectively. We have made this new dataset publicly accessible on the ZENODO repository [64].

• Advection equation (Adv): Following [48], we consider a wave advection equation with periodic boundary conditions, defined at the domain \(\Omega := [0,1]\) in space and \([0,1]\) in time. The problem is parameterized with respect to different initial conditions, and it is given as \[\begin{align} \frac{\partial u}{\partial t} + \frac{\partial u}{\partial x} & = 0, \quad \quad \quad & & \forall \;(x, t) \in \Omega \times [0, 1], \\ u & = u_0(x, \boldsymbol{\eta}), \; & & \forall \;(x, t) \in \Omega \times \{0\}. \\ \end{align}\] The DON is trained to learn a map from the initial condition to the solution of the PDE. To this aim, \(u_0\) is parametrized as \[u_0(x, \boldsymbol{\eta}) := h_1 1_{\{c_1 - \frac{w}{2}, c_1 + \frac{w}{2}\}} \\ + \sqrt{\max(h_2^2 - a^2(x-c_2)^2, 0)},\] where the parameters \(\boldsymbol{\eta} := [a, h_1, h_2, c_1, c_2, w]\) are randomly sampled from the uniform distributions \(a \sim \pazocal{U}(0.1, 0.2)\), \(h_1 \sim \pazocal{U}(0.5, 2.0)\), \(h_2 \sim \pazocal{U}(0.5, 2.0)\), \(c_1 \sim \pazocal{U}(0.2, 0.3)\), \(c_2 \sim \pazocal{U}(0.7, 0.8)\), and \(w \sim \pazocal{U}(0.1, 0.2)\), respectively. The dataset for training the DON is constructed using a finite difference method with uniform mesh composed of \(40 \times 40\) nodes in space-time. The number of training and test samples is \(1{,}000\). Note that for this benchmark problem, we utilize the POD-DeepONet architecture [48], which replaces the trunk network with the POD bases. This allows us to effectivelly demonstrate the capabilities of the proposed layer-wise decomposition for different types of DON architectures.

4.3 Implementation details↩︎

All benchmark problems and considered optimizers are implemented on top of the library PyTorch [65]. The code used for generating the numerical results is part of the open-source library DistTraiNN [66]. The description of DNN architectures is summarized in 1 and 2 for PINNs and DONs, respectively. All networks are initialized using Xavier initialization strategy [67]. For PINNs, the boundary conditions are imposed by employing the length factor function \(\ell\) given as \(\ell(x_1, x_2) = x_1(1-x_1)x_2(1-x_2)\) and \(\ell(t,x) = t(x+1)(x-1)\) for DA problem and Burg/AC problems, respectively. Moreover, we employ an adaptive activation function [68] and a skip-connection [43] to prevent gradient vanishing during training and improve the inference performance of the neural network.

As a baseline, we employ the LBFGS optimizer and single-level layer-wise preconditioned LBFGS (SL-LBFGS). The proposed two-level preconditioner is also tested by preconditioning the LBFGS method, giving rise to the TL-LBFGS optimizer. All LBFGS variants utilize the limited-memory implementation with the memory size \(m=3\). Moreover, we utilize momentum with parameter \(0.9\) and the cubic backtracking line-search with strong Wolfe conditions [69] to obtain all step-sizes. For the SL/TL-LBFGS optimizers, the global and local optimization problems outlined in 5, are also solved using the LBFGS, which is restarted at every epoch.

Compared to the standard LBFGS algorithm, a single iteration/epoch of the TL-LBFGS method is computationally more expensive. To analyze the iteration cost, 3 summarizes the gradient evaluations cost \((g_e)\), the update cost (UC) and memory cost (MC) of the LBFGS and TL-LBFGS optimizers. As we can see, the LBFGS optimizer requires one gradient evaluation per iteration, i.e., \(g_{e} = 1\). However, UC depends on the size of stored secant pairs (\(m\)) used for Hessian approximation. It can be approximately estimated as \(4mn\) flops when the compact matrix representation is employed; see [70] for details. By incorporating the cost associated with evaluating the momentum, the UC equals to \(n+4mn\) flops. In the perspective of MC, the LBFGS optimizer requires storing two matrices of size \(m \times n\) as well as a momentum of length \(n\), which takes up memory of size \(n+2mn\).

The computational cost of the TL-LBFGS optimizer is associated with global, local and coarse-level training phases. The global computational cost is exactly same as for the standard LBFGS optimizer. The gradient evaluation cost for local and coarse-level parts are \(k_{s} g_{e_{s}}\) and \(k_{0} g_{e_{0}}\), where \((g_{e_{s}})\) and \((g_{e_{0}})\) denote the cost of gradient evaluation of local and coarse-level optimizers, respectively. Since the number of coarse trainable parameters is the number of subdomain \(N_{sd}\), the cost of coarse-level evaluation \(g_{e_{0}}\) is generally less than that of local evaluation \(g_{e_{s}}\). In terms of UC, let \(\text{UC}_{s}\) and \(\text{UC}_{0}\) be the update cost required by the local and coarse optimizers. Since the UC for \(s^{th}\) subnetwork is scaled proportionally to the number of local trainable parameters \(n_{s}\) and the number of local iterations \(k_{s}\), the local update cost is \((\frac{\tilde{n}_s}{n} )k_s \text{UC}_{s}\). Similarly, the coarse-level update cost is \((\frac{n_{0}}{n} ) k_{0} \text{UC}_{0}\). Similarly, the coarse-level update cost is \((\frac{n_{0}}{n}) k_{0} \,\text{UC}_{0}\). Note that the MC of TL-LBFGS can be derived in the same way as the UC, i.e., the local memory cost and the coarse-level memory cost are \((\frac{\tilde{n}_s}{n}) \,\text{MC}_{s}\) and \((\frac{n_{0}}{n}) \,\text{MC}_{0}\), respectively. Moreover, note that in the case of the single-level layer-wise preconditioner for L-BFGS, termed SL-LBFGS, the computational cost can be estimated by simply setting \(k_{0}=0\).

5 Numerical results↩︎

In this section, we investigate the performance of the proposed layer-wise preconditioner. Since it has been demonstrated in [30], [71] that, for training SciML models, L-BFGS is clearly superior to Adam and other first-order methods, we focus our numerical results on the performance of the two-level layer-wise preconditioner for L-BFGS, giving rise to TL-LBFGS training algorithm. To this aim, we monitor the value of the loss function \(\pazocal{L}\) and the relative \(L^{2}\) error, given as \[\begin{align} \pazocal{E}_{\text{rel}}(u_{\boldsymbol{\theta}}, u^*) = \frac{ \| u_{\boldsymbol{\theta}} - u^* \|_{L^2(\Omega)} }{ \| u_{\boldsymbol{\theta}} \|_{L^2(\Omega)}}, \end{align}\] where \(u_{\boldsymbol{\theta}}\) is the neural network solution and \(u^{\ast}\) denotes the exact or the high-fidelity finite element solution.

5.1 Impact of the overlap on the convergence of SL-LBFGS↩︎

First, we investigate the impact of overlap on the convergence of the SL-LBFGS optimizer. During this experiment, we fix the number of local iterations at \(k_{s}=50\) and consider a varying number of subnetworks \(N_{sd} \in \{2, 4, 8\}\) and the size of the overlapping regions \(\delta \in \{0, 1, 2\}\). 7 illustrates the obtained results for Burg (PINN) and the Adv (DON) equations. Note, the layer-wise synchronization strategy proposed in 3.1.1 is not utilized. Instead, we average the parameters of the layers corresponding to the overlap, an approach traditionally used in numerical analysis. As we can see from the obtained results, in this particular case, increasing the size of the overlap \(\delta\) does not improve convergence, and in some cases, it even causes its deterioration.

For comparison, 8 demonstrates the obtained results by incorporating the novel layer-wise synchronization strategy. In contrast to the results reported in 7, we can observe that using the layer-wise synchronization strategy allows for an enhanced convergence as the size of the overlapping region \(\delta\) increases. Here, we, however, point out that as \(\delta\) increases, the computational cost of solving the local minimization problems also increases. Consequently, selecting \(\delta=1\) provides the best trade-off between enhanced convergence and the low computational cost.

5.2 Impact of coarse-level training (SL-LBFGS vs. TL-LBFGS)↩︎

As a next step, we investigate the impact of coarse-level training on the convergence of the proposed preconditioner. In this experiment, we fix the size of the overlap to \(\delta = 1\) and the number of local iterations to \(k_{s}=50\). At the same time, we vary the number of coarse-level iterations \(k_{0}\) and the number of subdomains \(N_{sd}\), such that \(k_{0} \in \{0, 25, 50\}\) and \(N_{sd} \in \{2, 4, 8\}\). 9 illustrates the obtained results. As we can see, the loss function decays more rapidly for a sufficient number of coarse iterations \(k_{0}\), suggesting that incorporating coarse-level training significantly improves the convergence rate of the proposed TL-LBFGS method. However, because the coarse-level training process is sequential and cannot be parallelized, an increased number of coarse-level iterations increases the computational cost. Moreover, we would like to remark that the coarse-level model is guaranteed to approximate the fine-level model only in the local neighborhood of the current fine-level iterate and is therefore inherently local. Consequently, performing too many \(k_{0}\) nonlinear steps on the coarse level may cause the coarse-level model to no longer be consistent with the fine-level model, which can in turn degrade overall convergence. To balance computational cost with improvements in convergence speed, we keep \(k_{0}\) fairly low for all subsequent numerical experiments, i.e., between \(25\) and \(50\).

5.3 Convergence, accuracy and computational cost of LBFGS and SL/TL-LBFGS optimizers↩︎

In this section, we compare the performance of the proposed TL-LBFGS optimizer with the standard LBFGS optimizer. The comparison is performed with respect to relative error \(\pazocal{E}_{\text{rel}}\), the number of epochs, \(g_{e}\) and UC, as discussed in 4.3. Our first experiment illustrates the impact of preconditioning on the overall performance of training. We compare the performance of the two-level additive Schwarz method as a standalone solver (TL-SA) and as a preconditioner for LBFGS (TL-LBFGS). The obtained results clearly indicate that the preconditioned variant is significantly more efficient. These observations are consistent with findings reported in the literature. For example, V. Dolean et al. [38] and X. Cai et al. [37] demonstrated that the nonlinear two-level additive Schwarz method alone is insufficient to achieve fast convergence. To overcome this limitation, they used the two-level additive Schwarz to precondition the Newton method. In view of the superiority of the preconditioning strategy, all remaining experiments employ the proposed two-level additive Schwarz method as a preconditioner for L-BFGS.

11 and 12 compare the performance of the LBFGS and TL-LBFGS for PINN’s and DON’s benchmark problems. As we can see, TL-LBFGS always provides the models with lower \(\pazocal{E}_{\text{rel}}\) than the LBFGS optimizer. For example, for the DA example with \(N_{sd}=5\), the relative error \(\pazocal{E}_{\text{rel}}\) almost immediately decreases to \(10^{-2}\), while LBFGS optimizer fails to train the model to that accuracy. Moreover, we can also observe that for all benchmark problems, the performance of TL-LBFGS improves as the number of subdomains \(N_{sd}\) increases.

As a next step, we compare the performance of TL-LBFGS with respect to its single-level counterpart, SL-LBFGS. We configure SL-LBFGS as proposed in [30], i.e., without overlap, without layer-wise synchronization strategy and without the coarse-level training. As we can see from 4, except in the case of Aniso problem, the TL-LBFGS method consistently outperforms the SL-LBFGS method by providing more accurate DNN models. While the observed improvements compared to the LBFGS method are around one order of magnitude, they are not as big as those compared to the SL-LBFGS method. However, in some cases, such as the Burg benchmark problem, where SL-LBFGS did not perform very well ², TL-LBFGS yields about four times more accurate results.

In the end, we compare \(g_{e}\) and UC required to achieve the same accuracy as the LBFGS optimizer. Tables 5 and 6 report the obtained results in terms of speedup factor \(\pazocal{S}\), defined as \(\pazocal{S}^A_B = \frac{C_A}{C_B}\), where \(C_A\) and \(C_B\) denote the \(g_{e}\) or UC required by the optimizer A and by the optimizer B, respectively. As we can see, for PINN’s problems, which are significantly more ill-conditioned than DON’s problems, utilizing the TL-LBFGS optimizer reduces \(g_{e}\) and UC by factors of \(1.5-1,099\) and \(3.5-12,121\) compared to the LBFGS method, respectively. Moreover, compared to SL-LBFGS, we achieve a speedup of \(2.9-7.3\) and \(1.6-5.1\) for \(g_{e}\) and UC, respectively. This effectively highlights the benefits of incorporating the proposed coarse-level training step, overlap, and layer-wise synchronization into the training process.

For DON’s benchmark problems, we observe lower speedup factors than for the PINNs. In particular, compared to LBFGS, we obtain speedup factors of \(12.8-18.3\) for \(g_{e}\) and \(33.2-45.5\) for UC, respectively. Compared to SL-LBFGS, we notice speedup factors of \(0.64-1.7\) for \(g_{e}\) and \(0.3-1.3\) for UC, respectively. The most unsatisfactory result is obtained for the Aniso example, where we even observe an increase in the computational cost. We attribute this behavior to the fact that this problem is the least ill-conditioned of all our benchmark problems, see 8. As a consequence, constructing a two-level preconditioner significantly increases the computational cost, while the convergence speed improves only slightly.

6 Discussion↩︎

Machine learning (ML)-based solvers for PDEs are rapidly gaining attention; however, their role remains distinct from that of classical numerical methods. Established approaches such as finite element and finite volume methods, when preconditioned by multigrid or domain decomposition techniques, continue to be the most efficient and reliable tools for solving large-scale, three-dimensional problems on unstructured geometries. At present, machine learning-based solvers should be regarded primarily as complementary methods, with particular strengths in high-dimensional problems, inverse problems, model discovery, and surrogate modeling for parametric problems.

A major drawback of current machine learning-based approaches is the absence of rigorous error analysis and convergence guarantees, which makes it challenging to assess their reliability compared to classical methods. Their performance is also highly sensitive to heuristic design choices, including the architecture of the network, the selection of hyperparameters, and the training strategy. From a computational standpoint, classical methods can often solve benchmark problems in fractions of a second, whereas training machine learning-based solvers typically require significant time and computational resources. As demonstrated in the manuscript, the development of novel, large-scale, training strategies can significantly reduce the associated computational burden. Moreover, once trained, neural networks enable extremely fast evaluations across multiple parameter configurations, which is particularly appealing for applications such as shape optimization, design under uncertainty, or real-time simulation.

In summary, although machine learning-based solvers are not yet ready to fully replace state-of-the-art numerical methods, they represent a promising complementary paradigm. A key open challenge is the development of novel training strategies that combine improved efficiency with enhanced model accuracy, ultimately enabling machine learning-based solvers to evolve into practical and reliable tools for large-scale scientific computing.

7 Conclusion↩︎

In this paper, we proposed a novel two-level overlapping additive Schwarz preconditioner for enhanced training of scientific machine learning applications. Motivated by traditional DD methods, the network parameters were decomposed into overlapping groups. The parameters of each group were then trained simultaneously, i.e., in parallel. To achieve fast convergence, we introduced a subdomain-wise synchronization strategy and incorporated a coarse-level training step. The performance of the devised two-level preconditioner was investigated using six benchmark problems, encompassing PINNs and DONs. The obtained results demonstrated that the proposed preconditioned optimizer yields significantly more accurate models while also reducing the training cost compared to the state-of-the-art LBFGS method and the existent single-level layer-wise preconditioner.

In the future, we plan to combine the proposed two-level layer-wise preconditioning framework with approaches that consider the decomposition of the computational domain/data space, such as those proposed in [32], [72]. Moreover, it would be interesting to extend the proposed training algorithm into stochastic settings, which would enable its applicability beyond scientific machine learning applications. Furthermore, developing strategies for enhanced parallel implementation, such as load balancing and optimization of the subnetwork’s loss/gradient evaluations, could further improve the performance.

Acknowledgements↩︎

Y.L. was supported in part by Basic Science Research Program through NRF funded by the Ministry of Education (No. RS2023-00247199). A.K. was supported by the Swiss National Science Foundation (SNSF) under the project “Multilevel training of DeepONets – multiscale and multiphysics applications”, as well as by the Platform for Advanced Scientific Computing (PASC) under the project ExaTrain. Moreover, work of A.K. benefited from the AI Interdisciplinary Institute ANITI, funded by the France 2030 program under Grant Agreement No. ANR-23-IACL-0002. G.E.K. is supported by the ONR Vannevar Bush Faculty Fellowship. We also acknowledge support from the DOE-MMICS SEA-CROGS DE-SC0023191 award and Ansys Inc.

8 Condition number of (layer-wise) neural tangent kernel↩︎

7 reports the logarithmic value of condition number of the neural tangent kernel (NTK) [73], [74], denoted by log(cond\((\mathbb{A})\)) for all benchmark problems considered in 4. As we can see, the problems associated with training PINNs are more ill-conditioned than those associated with training DONs. Moreover, we also see that for PINN problems, the layer-wise condition number, denoted by \(\text{cond}(\mathbb{A}_l)\), varies over a wider range. This suggests that the local training associated with solving 10 requires different adjustments of step-size. Thus, by decoupling the layers, we are able to use larger step-sizes for layers that are better conditioned, which in turn enhances the overall convergence.

9 Getting started with DistTrainNN↩︎

In this section, we briefly describe how to run an TL-LBFGS training algorithms within our open-source  DistTraiNN [66]. Firstly, the user is required to setup a Python environment with the following depedencies: PyTorch, Pandas, pyDOE, scikit-optimize, scikit-learn.

The scripts used to obtain the numerical results reported in Section 5 are available in the folder “tests/SPQN_tests”. An example code for training a PINN to approximate a solution for the Allen–Cahn equation is provided below. Please, note that the optimizer is fully configurable via command line arguments.

git clone git@bitbucket.org:alena_kopanicakova/disttrainn.git
cd disttrainn

export PYTHONPATH="${PYTHONPATH}:/path/to/disttrainn"

cd tests/TL_SPQN_tests/AllenCahn
python3 test_SPQN_twolevel_additive_serial.py --seed 1 --num_points_x 25 --num_points_y 50 --num_ref 4 --max_epochs 200000 --lr_global 1.0 --lr_local 1.0 --history 3 --width 64 --hiddenlayers_coarse 4 --num_subdomains 3 --num_local_steps 50 --num_coarse_steps 50 --use_layerwise_lr True --overlap_width 1

References↩︎