Enhancing training of physics-informed neural networks using
  domain-decomposition based preconditioning strategies
We propose to enhance the training of physics-informed neural networks
(PINNs). To this aim, we introduce nonlinear additive and multiplicative
preconditioning strategies for the widely used L-BFGS optimizer. The nonlinear
preconditioners are constructed by utilizing the Schwarz domain-decomposition
framework, where the parameters of the network are decomposed in a layer-wise
manner. Through a series of numerical experiments, we demonstrate that both,
additive and multiplicative preconditioners significantly improve the
convergence of the standard L-BFGS optimizer, while providing more accurate
solutions of the underlying partial differential equations. Moreover, the
additive preconditioner is inherently parallel, thus giving rise to a novel
approach to model parallelism.
Alena Kopaničáková, Hardik Kothari, George Em Karniadakis, Rolf Krause
Enhancing training of physics-informed neural networks using domain-decomposition based preconditioning strategies +

###### Abstract

We propose to enhance the training of physics-informed neural networks (PINNs). To this aim, we introduce nonlinear additive and multiplicative preconditioning strategies for the widely used L-BFGS optimizer. The nonlinear preconditioners are constructed by utilizing the Schwarz domain-decomposition framework, where the parameters of the network are decomposed in a layer-wise manner. Through a series of numerical experiments, we demonstrate that both, additive and multiplicative preconditioners significantly improve the convergence of the standard L-BFGS optimizer, while providing more accurate solutions of the underlying partial differential equations. Moreover, the additive preconditioner is inherently parallel, thus giving rise to a novel approach to model parallelism.

s cientific machine learning, nonlinear preconditioning, Schwarz methods, domain-decomposition pacs: 90C30, 90C26, 90C06, 65M55, 68T07

## 1 Introduction

In recent years, deep learning approaches for solving partial differential equations (PDEs) have gained popularity due to their simplicity, mesh-free nature, and efficiency in solving inverse and high-dimensional problems [34, 14]. The notable deep learning approaches in the context of scientific machine learning include physics-informed neural networks (PINNs) [64], deep Galerkin method [69], and deep Ritz method [15]. The idea behind these methods is to utilize deep neural networks (DNNs) to approximate solutions of the PDEs by minimizing a loss function, which incorporates the physical laws and domain knowledge. More specifically, in the context of PINNs, which we consider in this work, the loss function incorporates a PDE residual, boundary/initial conditions, and sometimes also observational data.

To approximate a solution of a PDE using PINNs, the optimal parameters of DNN have to be found during the training. This is typically carried out using Adam [35] and/or L-BFGS [46] optimizers. However, it has been shown that the performance of these optimizers deteriorates for highly ill-conditioned and stiff problems [73]. Several algorithmic enhancements have been therefore proposed to alleviate such difficulties, for instance, adaptive weighing strategies [74], empirical learning-rate annealing scheme [73], multiscale architectures [12, 23], ensemble learning [16] based on gradient boosting [6], adaptive activation functions [32, 31] etc. Furthermore, novel optimizers have also been developed, for instance, Gauss-Newton methods [65, 2], Levenberg-Marquardt optimizers [58, 76, 18], variable projection methods [60, 21, 61], block
coordinate methods [7], etc.

In this work, our goal is to enhance the convergence rate of the optimizers used for the training of PINNs by developing novel nonlinear preconditioning strategies. In the field of scientific computing, such nonlinear preconditioning strategies are often used to accelerate the convergence of the nonlinear iterative methods [3]. Specifically, left preconditioners [5, 47] enhance the conditioning of nonlinear problems by rebalancing nonlinearities or by transforming the basis of the solution space. In contrast, right preconditioners [51, 36] provide an improved initial guess for a given iterate. Both, left and right, preconditioners can be constructed using various approaches. Prominent approaches typically rely on the decomposition of the solution space into subspaces, such as field-split [39, 47], domain-decomposition (DD) [5, 42, 43], or multilevel decomposition [77, 41, 17, 22, 40]. Among these approaches, the DD approach is of particular interest due to its parallelization capabilities. As a consequence, preconditioners constructed using the DD approach not only improve the convergence rate of a given iterative method but also allow for a reduction of the overall computation time, as the workload can be distributed between larger amounts of computational resources.

In the field of machine learning, the parallelization of the training algorithms has attracted a lot of attention recently, see [8, 57, 1] for a comprehensive overview. The developed approaches can be broadly classified into two main groups: data parallel and model parallel. The data parallel approach has gained significant popularity due to its architecture-agnostic nature. In contrast, model parallelism is more intricate as it involves decomposing the network across different compute nodes. As a consequence, this approach has to be typically tailored to a specific network architecture [1]. For instance, the TorchBraid framework [26] enables layer-parallel training of ODE-based networks by utilizing a multigrid-in-time approach. In [25, 24], a model parallel approach for convolutional networks is proposed, where the network's width is partitioned into subnetworks, which can be trained in parallel. The parameters obtained from training these subnetworks are then used to initialize the global network, which is further trained sequentially. Similarly, in [37], a DD-based network architecture is proposed for parallel training of image recognition tasks.

In the context of PINNs, DD-based approaches have also been utilized to enhance training efficiency and scalability to larger computational domains and multi-scale problems. These approaches leverage the spatiotemporal decomposition of the computational domain. For instance, conservative PINNs (cPINNs) [33] utilize a non-overlapping Schwarz-based DD approach for conservation laws. The extended PINNs (XPINNs) expand the cPINN methodology to generic PDEs and arbitrary space-time domains. The idea behind XPINNs [30, 33, 68, 29] is to associate each subdomain with a sub-PINN, which can be trained independently of each other, i.e., in parallel. After the training of the subdomains is concluded, the sub-PINNs are stitched together by enforcing continuity on the interfaces and by averaging solutions along neighboring subdomain interfaces. The XPINN framework was later expanded in several ways. For instance, the adaptive PINNs (APPINs) [28], GatedPINNs [70], or extensions incorporating extreme learning machines [13] have been proposed.

An alternative approach that utilizes the variational principle and overlapping Schwarz-based DD approach has also been explored, see [45]. This method was further extended by introducing coarse space acceleration to improve convergence with respect to an increasing number of subdomains [55]. Building upon the variational principles and the deep Ritz method, the D3M [44] also employs the overlapping Schwarz method. Another approach, called finite basis PINNs (FBPINN) [59, 11],
represents the solution of a PDE as a sum of basis functions with compact support defined over small, overlapping subdomains. These basis functions are learned in parallel using neural networks. Finally, we mention PFNN-2 method [67], which enhances the training of PINNs by utilizing overlapping DD together with the penalty-free network approach (PFNN) [66].

In this work, we aim to enhance the training of the PINNs by developing novel nonlinear additive and multiplicative preconditioning strategies. More specifically, we utilize the right preconditioning approach to accelerate the convergence of the widely used L-BFGS method. The devised preconditioners are constructed by decomposing the network parameters in a layer-wise manner. Methodologically, our approach differs from the aforementioned DD approaches for PINNs that utilize the decomposition of the computational domain. As a consequence, the proposed preconditioners are agnostic to the physics-based priors and have therefore potential to be also applied to other machine-learning tasks. Using the proposed DD approach, we demonstrate that both additive and multiplicative preconditioners can significantly improve the convergence rate of the standard L-BFGS optimizer. Moreover, utilizing the additive preconditioner naturally enables model parallelism, which allows for additional acceleration of the training, if a larger amount of computational resources is available.

This paper is organized as follows. Section 2 provides a brief overview of the PINNs. In section 3, we propose nonlinear additive and multiplicative preconditioners using a layer-wise decomposition of the network. Section 4 discusses implementation details, while section 5 introduces a set of benchmark problems used for numerical experiments. Section 6 demonstrates the numerical performance of the proposed preconditioning strategies. Finally, section 7 summarizes the presented work.

## 2 Physics-informed neural networks

We consider the differential equations of a generic form and pose the problem as follows. Find \(u:\Omega\to\mathbb{R}\), such that

\[\mathcal{P}(u(\mathbf{x})) =f(\mathbf{x}), \mathbf{x}\in\Omega, \tag{1}\] \[\mathcal{B}^{k}(u(\mathbf{x})) =g^{k}(\mathbf{x}), \mathbf{x}\in\Gamma^{k}\subseteq\partial\Omega, \text{for }k=1,2,\ldots,n_{\Gamma}.\]

The symbol \(\mathcal{P}\) represents a nonlinear differential operator, while \(\{\mathcal{B}^{k}\}_{k=1}^{n_{\Gamma}}\) denotes a set of boundary condition operators. In the case of time-dependent problems, we treat time \(t\) as an additional component of the vector \(\mathbf{x}\in\mathbb{R}^{d}\), where \(d\) corresponds dimension of the problem. Therefore, the computational domain \(\Omega\) encompasses both spatial and temporal dimensions, and the initial conditions can be considered as a specific type of boundary conditions within the spatiotemporal domain.

### DNN approximation of solution

Following [64], our goal is to approximate a solution of (1) using PINNs, i.e., \(u(\mathbf{x})\approx u_{\texttt{NN}}(\mathbf{\theta},\mathbf{x})\), where the nonlinear mapping \(u_{\texttt{NN}}:\mathbb{R}^{n}\times\mathbb{R}^{d}\to\mathbb{R}\) is parametrized by the trainable parameters \(\mathbf{\theta}\in\mathbb{R}^{n}\). Typically, the exact form of \(u_{\texttt{NN}}\) is induced by the network architecture. Here, we consider a network that is constructed as a composition of an input layer, \(L-1\) hidden layers, and an output layer, thus as

\[\begin{array}{llll}\text{input layer:}&\mathbf{y}_{0}&=\mathbf{W}_{0}\mathbf{x}, \\ \text{hidden layers:}&\mathbf{y}_{l}&=\mathcal{N}(\mathbf{y}_{l-1}),\\ \text{output layer:}&\mathbf{y}_{L}&=\mathbf{W}_{L}\mathbf{y}_{L-1}+\mathbf{b}_{L}.\end{array} \tag{2}\]

The input layer maps input features \(\mathbf{x}\in\Omega\subset\mathbb{R}^{d}\) into the dimension of the hidden layers using weights \(\mathbf{W}_{0}\in\mathbb{R}^{d\times nh}\). The \(l^{th}\) hidden layer nonlinearly transforms the output
of the \((l-1)^{th}\) layer \(\mathbf{y}_{l-1}\) into \(\mathbf{y}_{l}\) by means of nonlinear operator \(\mathcal{N}:\mathbb{R}^{nh}\to\mathbb{R}^{nh}\). The operator \(\mathcal{N}\) can take on multiple forms. We utilize a single-layer perceptron with a skip connection, given as

\[\mathcal{N}(\mathbf{y}_{l-1}):=\mathbf{y}_{l-1}+\sigma_{l}(\mathbf{W}_{l}\mathbf{y}_{l-1}+ \mathbf{b}_{l}). \tag{3}\]

The parameters associated with layer \(l\) can be given as \(\mathbf{\theta}_{l}:=(\text{flat}(\mathbf{W}_{l}),\text{flat}(\mathbf{b}_{l}))\) collectively, where \(\mathbf{W}_{l}\in\mathbb{R}^{nh\times nh}\) and \(\mathbf{b}_{l}\in\mathbb{R}^{nh}\) denote weights and biases, respectively. The function \(\text{flat}(\cdot)\) represents a flattening operator and \(\sigma_{l}:\mathbb{R}^{nh}\to\mathbb{R}^{nh}\) denotes an activation function. Finally, the output of the last hidden-layer \(\mathbf{y}_{L-1}\) is linearly transformed using \(\mathbf{W}_{L}\) and \(\mathbf{b}_{L}\), giving rise to the network's output \(\mathbf{y}_{L}\).

### Training data and loss functional

In order to train the network such that it can approximate a solution of (1), we construct a dataset \(\mathcal{D}_{\text{int}}=\{\mathbf{x}_{j}\}_{j=1}^{n_{\text{int}}}\) of \(n_{\text{int}}\) collocations points, which lie in the interior of the computational domain \(\Omega\). These points are used to enforce the physics captured by the PDE. In addition, we consider datasets \(\{\mathcal{D}_{\text{bc}}^{k}\}_{k=1}^{n_{\text{r}}}\), where each \(\mathcal{D}_{\text{bc}}^{k}=\{(\mathbf{x}_{j}^{k},\mathbf{g}_{j}^{k})\}_{j=1}^{n_{ \text{bc}}}\) contains \(n_{\text{bc}}^{k}\) points, which we use to enforce the boundary conditions on \(\Gamma^{k}\).

Using \(\mathcal{D}_{\text{int}}\) and \(\{\mathcal{D}_{\text{bc}}^{k}\}_{k=1}^{n_{\text{r}}}\), the optimal parameters of the network are found by solving the following minimization problem:

\[\mathbf{\theta}^{*}:=\underset{\mathbf{\theta}\in\mathbb{R}^{n}}{\arg\min}\ \underbrace{\tilde{\mathcal{L}}_{\text{int}}(\mathbf{\theta})}_{\text{ interior loss}}+\underbrace{\tilde{\mathcal{L}}_{\text{bc}}(\mathbf{\theta})}_{\text{boundary loss}}. \tag{4}\]

Here, the interior and the boundary losses are given as

\[\tilde{\mathcal{L}}_{\text{int}}(\mathbf{\theta}) :=\frac{w_{\text{int}}}{|\mathcal{D}_{\text{int}}|}\sum_{\mathbf{x}_ {j}\in\mathcal{D}_{\text{int}}}\|\mathcal{P}(\tilde{u}_{\text{NN}}(\mathbf{\theta },\mathbf{x}_{j}))-f(\mathbf{x}_{j})\|^{2}, \tag{5}\] \[\tilde{\mathcal{L}}_{\text{bc}}(\mathbf{\theta}) :=\sum_{k=1}^{n_{\text{r}}}\Bigg{(}\frac{w_{\text{bc}}}{| \mathcal{D}_{\text{bc}}^{k}|}\sum_{(\mathbf{x}_{j}^{k},\mathbf{g}_{j}^{k})\in\mathcal{ D}_{\text{bc}}^{k}}\|(\tilde{u}_{\text{NN}}(\mathbf{\theta},\mathbf{x}_{j}^{k})-g^{k}(\mathbf{x}_{j }^{k})\|^{2}\Bigg{)}.\]

Thus, the interior loss \(\tilde{\mathcal{L}}_{\text{int}}\) is expressed as the mean PDE residual obtained by Monte Carlo integration, i.e., by the averaging over the residuals evaluated at collocation points sampled in the interior of the domain. Note that an evaluation of the PDE residual typically requires knowledge of the partial derivatives of the network output \(\tilde{u}_{\text{NN}}(\mathbf{\theta},\mathbf{x}_{j})\) with respect to the input \(\mathbf{x}_{j}\). These derivatives can be efficiently obtained using advanced automatic differentiation techniques. The second term, boundary loss \(\tilde{\mathcal{L}}_{\text{bc}}\), ensures that the boundary conditions are satisfied at a set of points sampled along boundaries \(\{\Gamma^{k}\}_{k=1}^{n_{\Gamma}}\).

The weights \(w_{\text{int}}\) and \(w_{\text{bc}}\) in (5) are used to balance the interior and boundary loss terms. Their optimal value can be determined either during the hyper-parameter tuning or directly during the training by employing adaptive strategies, such as one proposed in [54]. In this work, we overcome the difficulty of finding suitable values of \(w_{\text{int}}\) and \(w_{\text{bc}}\) by imposing the boundary conditions in a penalty-free manner, thus by eliminating the boundary loss term [50]. To this aim, the output of the network \(\tilde{u}_{\text{NN}}(\mathbf{\theta},\mathbf{x})\) is modified using the length factor function \(\ell^{k}:\Gamma^{k}\subset\mathbb{R}^{d}\to\mathbb{R}\) as follows

\[u_{\text{NN}}(\mathbf{\theta},\mathbf{x}):=\sum_{k=1}^{n_{\Gamma}}g^{k}(\mathbf{x})+\frac{ \tilde{\ell}(\mathbf{x})}{\max_{\mathbf{x}\in\Omega}\tilde{\ell}(\mathbf{x})}\tilde{u}_{ \text{NN}}(\mathbf{\theta},\mathbf{x}),\ \text{where}\ \tilde{\ell}(\mathbf{x})=\prod_{k=1}^{n_{\Gamma}}(1-(1-\ell^{k}(\mathbf{x}))). \tag{6}\]
The function \(\ell^{k}\) is constructed such that

\[\ell^{k}(\mathbf{x}) =0, \mathbf{x}\in\Gamma^{k} \tag{7}\] \[\ell^{k}(\mathbf{x}) =1, \mathbf{x}\in\Gamma^{j},\quad j\in\{1,2,\ldots,n_{\Gamma}\},\quad\text{ for }k\neq j\] \[\ell^{k}(\mathbf{x}) \in(0,1) \text{otherwise.}\]

The details regarding the construction of \(\ell^{k}\) for complex geometries can be found in [71].

Using (6), the minimization problem (4) can be now reformulated as

\[\mathbf{\theta}^{*}:=\operatorname*{arg\,min}_{\mathbf{\theta}\in\mathbb{R}^{n}}\ \mathcal{L}(\mathbf{\theta}):=\frac{1}{|\mathcal{D}_{\texttt{int}}|}\sum_{\mathbf{x}_ {j}\in\mathcal{D}_{\texttt{int}}}\|\mathcal{P}(u_{\texttt{NN}}(\mathbf{\theta}, \mathbf{x}_{j}))-f(\mathbf{x}_{j})\|^{2}. \tag{8}\]

After training, the optimal parameters \(\mathbf{\theta}^{*}\) can be used to infer an approximate solution \(u_{\texttt{NN}}\) of (1). The quality of this solution can be assessed by means of the error \(\mathcal{E}(u_{\texttt{NN}},u^{*}):=\|u_{\texttt{NN}}-u^{*}\|\), where \(u^{*}\) denotes an exact solution. In general, the error \(\mathcal{E}\) consists of three components: the discretization error, network approximation error, and optimization error [56]. More specifically, the discretization error is determined by the number and location of the collocation points. The approximation error of the network is associated with the representation capacity of the DNN, i.e., it directly depends on the network architecture (number of layers, network width). The optimization error quantifies the quality of the solution provided by the optimizer. Our goal, in this work, is to reduce the optimization error by developing enhanced training algorithms.

## 3 Nonlinearly preconditioned training using layer-wise decomposition of network

The minimization problem (8) is traditionally solved using Adam and/or L-BFGS optimizers. As it has been pointed out in [73], the convergence of these optimizers deteriorates with an increasing problem stiffness, which is caused by the multiscale and multirate nature of the considered problems. We aim to overcome this difficulty by proposing a novel nonlinear preconditioning strategy.

Let us recall that a minimizer of (8) has to satisfy the following first-order criticality condition

\[\nabla\mathcal{L}(\mathbf{\theta})=0. \tag{9}\]

As solving (9) is computationally demanding, we instead construct and solve an alternative, a nonlinearly preconditioned system of equations, given as

\[\mathcal{F}(\mathbf{\theta})=0. \tag{10}\]

The preconditioned system of equations (10) shall be constructed such that it is easier to solve than (9), but such that it has the same solution as an original nonlinear system. Thus, the preconditioned problem should have more balanced nonlinearities, and the numerical computations involved in solving (10) should be computationally cheaper.

### Network decomposition

In this work, we propose to construct the nonlinearly preconditioned system of equations (10) by decomposing the parameter space of the network. In particular, subnetworks are constructed by layer-wise decomposition of the network into \(N_{sd}\) disjoint groups, often called subdomains. Thus, each layer is assigned a unique subnetwork identifier, and the parameters of the network are
also split into \(N_{sd}\) disjoint groups, i.e., \(\mathbf{\theta}=[\mathbf{\theta}_{1},\ldots,\mathbf{\theta}_{s},\ldots,\mathbf{\theta}_{N_{sd}}]^ {\top}\), where \(\mathbf{\theta}_{s}\in\mathbb{R}^{n_{s}}\) denotes parameters associated with a \(s^{th}\) subnetwork. Please, refer to Figure 1 for the illustration of layer-wise decomposition. In passing, we note that the different types of network decomposition can be potentially utilized, e.g., intra-layer decomposition with or without overlap.

Transfer of the data between the subnetwork and the global network is performed using two transfer operators. Restriction operator \(\mathbf{R}_{s}:\mathbb{R}^{n}\rightarrow\mathbb{R}^{n_{s}}\) extracts the parameters of network associated with the \(s^{th}\) subnetwork, i.e.,

\[\mathbf{\theta}_{s}=\mathbf{R}_{s}\mathbf{\theta},\qquad\text{for }s=1,2,\ldots,N_{sd}. \tag{12}\]

In contrast, an extension operator \(\mathbf{E}_{s}:\mathbb{R}^{n_{s}}\rightarrow\mathbb{R}^{n}\) is used to extend the quantities from the subnetworks to the global network, i.e.,

\[\mathbf{\theta}=\sum_{s=1}^{N_{sd}}\mathbf{E}_{s}\mathbf{\theta}_{s}. \tag{13}\]

Using the proposed network decomposition, we now define local minimization problems associated with a given subnetwork \(s\) as

\[\mathbf{\theta}_{s}^{*}=\operatorname*{arg\,min}_{\mathbf{\theta}_{s}}\mathcal{L}( \mathbf{\theta}_{1},\ldots,\mathbf{\theta}_{s},\ldots,\mathbf{\theta}_{N_{sd}}). \tag{14}\]

Solving (14) corresponds to minimizing the loss functional \(\mathcal{L}\) with respect to the parameters \(\mathbf{\theta}_{s}\), while all other parameters of the network are kept fixed. This minimization process can be carried out by any optimizer of a choice, such as Adam [35], L-BFGS [46], or Newton [10]. For the purpose of this work, we employ the L-BFGS method, see also section 4. Note, this process loosely resembles the transfer-learning approaches [75], where only a particular subset of the network parameters is being (re-)trained.

### Right-preconditioned training optimizers

In this work, we construct a nonlinearly preconditioned system of equations as follows

\[\mathcal{F}(\mathbf{\theta}):=\nabla\mathcal{L}(G(\mathbf{\theta})), \tag{15}\]

where the operator \(G:\mathbb{R}^{n}\rightarrow\mathbb{R}^{n}\) denotes a right preconditioner. Using the network-decomposition framework introduced in subsection 3.1, we define operator \(G\) as an iterative method of the following form:

\[G(\mathbf{\theta}^{(k)})=\mathbf{\theta}^{(k)}+\alpha^{(k)}\sum_{s=1}^{N_{sd}}\mathbf{ E}_{s}(\mathbf{\theta}_{s}^{*}-\mathbf{R}_{s}\mathbf{\theta}^{(k)}). \tag{16}\]

Figure 1: An example of the layer-wise decomposition of the network.
Here, a symbol \(\mathbf{\theta}_{s}^{*}\) represents a solution of the local minimization problem (11), associated with \(s^{th}\) subnetwork, while \(\alpha^{(k)}\) denotes a step-size.

The subnetwork solutions \(\{\mathbf{\theta}_{s}^{*}\}_{s=1}^{N_{sd}}\) can be obtained independently from each other, i.e., in parallel, which gives rise to an additive preconditioner \(G\). Alternatively, one might obtain \(\{\mathbf{\theta}_{s}^{*}\}_{s=1}^{N_{sd}}\) in a multiplicative fashion. In this case, the local minimization problems are solved sequentially and the already acquired local solutions \(\{\mathbf{\theta}_{z}^{*}\}_{z=1}^{s-1}\) are taken into account while performing local training for the \(s^{th}\) subnetwork.

As we can see from (10), \(\mathcal{F}\) is a composite operator. Thus, for a given iterate \(\mathbf{\theta}^{(k)}\), the preconditioner \(G\) is used to obtain an improved iterate, denoted by \(\mathbf{\theta}^{(k+\nicefrac{{1}}{{2}})}=G(\mathbf{\theta}^{(k)})\). The improved iterate \(\mathbf{\theta}^{(k+\nicefrac{{1}}{{2}})}\) is then used as an initial guess for the next global optimization step, giving rise to the following update rule:

\[\mathbf{\theta}^{(k+1)}=\mathbf{\theta}^{(k+\nicefrac{{1}}{{2}})}+\alpha^{(k+ \nicefrac{{1}}{{2}})}\mathbf{p}^{(k+\nicefrac{{1}}{{2}})}, \tag{12}\]

where \(\mathbf{p}^{(k+\nicefrac{{1}}{{2}})}\) is obtained by a global optimizer and \(\alpha^{(k+\nicefrac{{1}}{{2}})}\) denotes appropriately selected global step-size.

#### 3.2.1 Right-preconditioned L-BFGS

There are many possible choices of a global optimizer, e.g., Adam [35], or AdaGrad [52]. Here, we utilize the BFGS method, as it is widely employed for the training of PINNs. The BFGS method belongs to the family of quasi-Newton (QN) methods, which have gained popularity due to the fact that they do not require explicit knowledge of the Hessian. Instead, a Hessian approximation \(\mathbf{B}^{(k+1)}\) is constructed such that the following secant equation is satisfied:

\[\mathbf{B}^{(k+1)}\mathbf{s}^{(k)}=\mathbf{y}^{(k)}, \tag{13}\]

where \(\mathbf{s}^{(k)}=\mathbf{\theta}^{(k+1)}-\mathbf{\theta}^{(k+\nicefrac{{1}}{{2}})}\) and \(\mathbf{y}^{(k)}=\nabla\mathcal{L}(\mathbf{\theta}^{(k+1)})-\nabla\mathcal{L}(\mathbf{ \theta}^{(k+\nicefrac{{1}}{{2}})})\). Note, that in the context of right-preconditioned nonlinear systems, the secant pairs are obtained by taking into account iterate \(\mathbf{\theta}^{(k+\nicefrac{{1}}{{2}})}\), obtained as an outcome of the nonlinear preconditioning step, see [42] for details. Since the secant equation (13) is indeterminate, additional constraints must be applied to uniquely determine the approximation \(\mathbf{B}^{(k+1)}\). A particular choice of constraints leads to different variants of the QN methods. For instance, the BFGS method is obtained by imposing the symmetry and positive definiteness of \(\mathbf{B}^{(k+1)}\).

In order to reduce the memory footprint of the proposed training method, we employ a limited-memory variant of the BFGS method, termed L-BFGS. Using compact matrix representation [46], the approximation \(\mathbf{B}^{(k+1)}\) is given as

\[\mathbf{B}^{(k+1)}=\mathbf{B}^{(0)}-\begin{bmatrix}\mathbf{B}^{(0)}\mathbf{S} ^{(k)}\\ \mathbf{Y}^{(k)}\end{bmatrix}^{\top}\begin{bmatrix}(\mathbf{S}^{(k)})^{\top} \mathbf{B}^{(0)}\mathbf{S}^{(k)}&\mathbf{L}^{(k)}\\ (\mathbf{L}^{(k)})^{\top}&-\mathbf{D}^{(k)}\end{bmatrix}^{-1}\begin{bmatrix} (\mathbf{S}^{(k)})^{\top}\mathbf{B}^{(0)}\\ (\mathbf{Y}^{(k)})^{\top}\end{bmatrix}, \tag{14}\]

where the matrices \(\mathbf{S}^{(k)}:=[\mathbf{s}^{(k-m)},\dots,\mathbf{s}^{(k-1)}]\), and \(\mathbf{Y}^{(k)}:=[\mathbf{y}^{(k-m)},\dots,\mathbf{y}^{(k-1)}]\) are constructed using \(m\) most recent secant pairs \(\{(\mathbf{s}^{(i)},\mathbf{y}^{(i)})\}_{i=k-m}^{k-1}\). The matrices \(\mathbf{L}^{(k)}\) and \(\mathbf{D}^{(k)}\) denote the strictly lower triangular, and the diagonal part of the matrix \((\mathbf{S}^{(k)})^{\top}\mathbf{Y}^{(k)}\), while \(\mathbf{B}^{(0)}\) is an initial approximation of the Hessian. Usually \(\mathbf{B}^{(0)}=\gamma\mathbf{I}\), where \(\gamma=\frac{\langle\mathbf{y}^{(k)},\mathbf{y}^{(k)}\rangle}{\langle\mathbf{y}^{(k)},\bm {s}^{(k)}\rangle}\), see [63] for alternatives.

Using (14), we are now able to state the global L-BFGS step as follows

\[\mathbf{\theta}^{(k+1)}=\mathbf{\theta}^{(k+\nicefrac{{1}}{{2}})}+\alpha^{(k+ \nicefrac{{1}}{{2}})}\underbrace{\big{(}-(\mathbf{B}^{(k+1)})^{-1}\nabla \mathcal{L}(\mathbf{\theta}^{(k+\nicefrac{{1}}{{2}})})\big{)}}_{=:\mathbf{p}^{(k+ \nicefrac{{1}}{{2}})}}. \tag{15}\]
Using the compact matrix representation, the inverse Hessian \((\mathbf{B}^{(k+1)})^{-1}\) can be obtained explicitly by applying the Sherman-Morrison-Woodbury formula. Alternatively, the two-loop recursion formula [19] can be utilized.

The convergence of the right-preconditioned L-BFGS method can be further enhanced by incorporating momentum, as often done for training DNNs. Following [53, 63], we evaluate the momentum \(\boldsymbol{v}^{(k+\nicefrac{{1}}{{2}})}\in\mathbb{R}^{n}\) in recursive manner as

\[\boldsymbol{v}^{(k+\nicefrac{{1}}{{2}})}=(1-\mu)\boldsymbol{v}^{(k-\nicefrac{{ 1}}{{2}})}+\mu\boldsymbol{p}^{(k+\nicefrac{{1}}{{2}})}, \tag{12}\]

where \(\mu\in[0,1]\). Utilizing the momentum, the global step (11) can be now reformulated as

\[\boldsymbol{\theta}^{(k+1)}=\boldsymbol{\theta}^{(k+\nicefrac{{1}}{{2}})}+ \alpha^{(k+\nicefrac{{1}}{{2}})}\boldsymbol{v}^{(k+\nicefrac{{1}}{{2}})}. \tag{13}\]

Algorithm 1 summarizes the Schwarz preconditioned quasi-Newton (SPQN) methods, namely additive and multiplicative variants termed ASPQN and MSPQN, respectively.

```
0:\(\mathcal{L}:\mathbb{R}^{n}\to\mathbb{R},\ \boldsymbol{\theta}^{(0)}\in\mathbb{R}^{n},\ \epsilon\in\mathbb{R},\ \mu\in[0,1],\ k_{\max}\in\mathbb{N}^{+}, \boldsymbol{v}^{-\nicefrac{{1}}{{2}}}=\boldsymbol{0}\)
1:\(\boldsymbol{v}^{(0)}=0\)
2:for\(k=0,\dots,k_{\max}\)do
3:for\(s=1,\dots,N_{sd}\)do
4:if additive then
5:\(\boldsymbol{\theta}_{s}^{*}=\arg\min_{\boldsymbol{\theta}_{s}}\mathcal{L}( \boldsymbol{\theta}_{1}^{(k)},\dots,\boldsymbol{\theta}_{s},\dots,\boldsymbol {\theta}_{N_{sd}}^{(k)})\)\(\triangleright\) parallel
6:else
7:\(\boldsymbol{\theta}_{s}^{*}=\arg\min_{\boldsymbol{\theta}_{s}}\mathcal{L}( \boldsymbol{\theta}_{1}^{*},\dots,\boldsymbol{\theta}_{s-1}^{*},\boldsymbol{ \theta}_{s},\boldsymbol{\theta}_{s+1}^{(k)},\dots,\boldsymbol{\theta}_{N_{sd} }^{(k)})\)\(\triangleright\) sequential
8:endif
9:endfor
10:
11:\(\boldsymbol{\theta}^{(k+\nicefrac{{1}}{{2}})}=\boldsymbol{\theta}^{(k)}+ \alpha^{(k)}\sum_{s=1}^{N_{sd}}\boldsymbol{E}_{s}(\boldsymbol{R}_{s} \boldsymbol{\theta}^{(k)}-\boldsymbol{\theta}_{s}^{*})\)\(\triangleright\) Synchronization step
12:
13:\(\boldsymbol{p}^{(k+\nicefrac{{1}}{{2}})}=-(\mathbf{B}^{(k+1)})^{-1}\nabla \mathcal{L}(\boldsymbol{\theta}^{(k+\nicefrac{{1}}{{2}})})\)\(\triangleright\) Preconditioned quasi-Newton step
14:\(\boldsymbol{v}^{(k+\nicefrac{{1}}{{2}})}=(1-\mu)\boldsymbol{v}^{(k- \nicefrac{{1}}{{2}})}+\mu\boldsymbol{p}^{(k+\nicefrac{{1}}{{2}})}\)\(\triangleright\) Momentum evaluation
15:\(\boldsymbol{\theta}^{(k+1)}=\boldsymbol{\theta}^{(k+\nicefrac{{1}}{{2}})}+ \alpha^{(k+\nicefrac{{1}}{{2}})}\boldsymbol{v}^{(k+\nicefrac{{1}}{{2}})}\)\(\triangleright\) Global parameter update
16:
17: Update \(\mathbf{S}^{(k+1)}\) with \(\boldsymbol{s}^{(k)}=\boldsymbol{\theta}^{(k+1)}-\boldsymbol{\theta}^{(k+ \nicefrac{{1}}{{2}})}\)\(\triangleright\) Update of secant pairs
18: Update \(\mathbf{Y}^{(k+1)}\) with \(\boldsymbol{y}^{(k)}=\nabla\mathcal{L}(\boldsymbol{\theta}^{(k+1)})-\nabla \mathcal{L}(\boldsymbol{\theta}^{(k+\nicefrac{{1}}{{2}})})\)
19:endfor
20:return\(\boldsymbol{\theta}^{(k+1)}\)
```

**Algorithm 1** Schwarz Preconditioned Quasi-Newton (SPQN)

## 4 Implementation details, computational cost, and memory requirements.

We implement the proposed ASPQN and MSPQN optimizers as part of Dist-TraiNN1 library [38], which utilizes PyTorch[62] for the implementation of PINNs. As mentioned earlier, the MSPQN algorithm is inherently sequential, and therefore its implementation targets computing architectures with a single GPU device.

Footnote 1: The code will be made publicly available upon acceptance of the manuscript.

In contrast, the ASPQN algorithm is inherently parallel hence its implementation is designed for distributed computing environments with multiple GPU devices. Here,
we utilize the torch.distributed package with the NCCL backend and implement the ASPQN method such that each GPU is associated with a single subnetwork, i.e., the number of GPUs determines the number of subnetworks. As an evaluation of loss associated with a subnetwork requires input to pass through an entire network, we replicate the global network to all GPUs. Each subnetwork then undergoes local optimization by training the subnetwork's parameters, i.e., by minimizing (12). This is in contrast to other model parallel approaches, where each GPU gets access to only a specific part of the network [8, 1]. As a consequence, the standard model parallel approaches have a smaller memory footprint compared to our approach. After the subnetwork's training is concluded, the synchronization of local updates is performed (line 11, Algorithm 1). Subsequently, a global L-BFGS step is executed concurrently on all nodes. Figure 2 provides a sketch of the ASPQN parallelization strategy.

_Remark 4.1_.: The implementation of the SPQN methods can be seamlessly integrated with data parallel approaches.

We solve the local minimization problem (12) using the L-BFGS method. This minimization is carried out only approximately, using a fixed number of iterations, denoted by \(k_{s}\). Both, global and subnetwork, L-BFGS Hessian approximations take advantage of three secant pairs, i.e., \(m=3\). The Hessian approximation for each subnetwork is restarted every time the local training commences. Moreover, we employ cubic backtracking line-search method with strong Wolfe conditions [9, Algorithm A6.3.1, pages 325-327] to obtain \(\alpha^{(k+\nicefrac{{1}}{{2}})}\) and \(\alpha^{(k)}\). Alternatively, one could consider using fixed step sizes, which may result in a more efficient algorithm that would require fewer evaluations of the loss function. However, determining suitable values of \(\alpha^{(k+\nicefrac{{1}}{{2}})}\) and \(\alpha^{(k)}\) would necessitate a thorough hyper-parameter search, which would eventually lead to a significant rise in overall computational cost.

All presented numerical experiments were performed at the Swiss National Supercomputing Centre (CSCS) using XC50 compute nodes of the Piz Daint supercomputer. Each XC50 compute node consists of an Intel Xeon E5-2690 V3 processor and an NVIDIA Tesla P100 GPU. The memory of a node is 64 GB, while the memory of the GPU is 16 GB.

### Computational cost and memory requirements

In this section, we discuss the computational cost and memory requirements of the proposed SPQN

Figure 2: A sketch of local-to-global updates utilized by the ASPQN.
methods, as well as standard Adam and L-BFGS optimizers. In particular, the Adam optimizer requires a single evaluation of the loss function and the gradient per iteration. The update cost involves not only updating the parameters but also computing the biased and corrected first-order and second-order moment estimates, thus approximately \(5n\) flops in total. The memory requirement for the Adam optimizer is related to storing the iterate, gradient, and first-order and second-order moment estimates, leading to the total memory of size \(4n\).

The L-BFGS method requires an evaluation of one gradient per iteration. However, the loss function might be evaluated multiple times, depending on the number of required line-search iterations, denoted by \(ls_{\text{its}}\). The update cost depends on the number of stored secant pairs (\(m\)) that are used for approximating the Hessian. Using compact matrix representation, it can be approximately estimated as \(4mn\) flops, see [4, Sections 3 and 4] for details. Moreover, we have to take into account the cost associated with an evaluation of the momentum term, leading to an overall cost of approximately \(n+4mn\) flops. In the context of memory requirements, the L-BFGS method necessitates the storage of the momentum term \(\mathbf{v}^{(k)}\) as well as matrices \(\mathbf{\mathrm{S}}^{(k)}\) and \(\mathbf{\mathrm{Y}}^{(k)}\), which accounts for the memory of size \(n+2mn\).

In the case of the SPQN methods, the computational cost consists of global and local parts. The global update cost and memory requirements consist of that of the standard L-BFGS method. Moreover, one has to take into account the cost related to the evaluation of the preconditioned iterate \(\mathbf{\theta}^{(k+\nicefrac{{1}}{{2}})}\) (line 11, Algorithm 3.1).

The local update cost is associated with the cost required by the local optimizer, denoted by \(\mathrm{UC}_{s}\). Thus, for the \(s^{th}\) subnetwork, the update cost is scaled proportionally to the number of local parameters \(n_{s}\) and the number of local iterations \(k_{s}\). Assuming that all subnetworks consist of the same number of local trainable parameters (\(n_{s}\)), the local update cost of the ASPQN method is therefore \(\binom{n_{s}/n}{k_{s}}\mathrm{UC}_{s}\). For the MSPQN method, the local update cost has to be summed up over all subnetworks, due to the sequential nature of the algorithm, thus \(\sum_{s=1}^{N_{sd}}\left(\binom{n_{s}/n}{k_{s}}\mathrm{UC}_{s}\right)=k_{s} \mathrm{UC}_{s}\). Regarding the memory cost of both SPQN methods, we have to take into account the storage of global and local Hessian approximations, the momentum, and the preconditioned iterate \(\mathbf{\theta}^{(k+\nicefrac{{1}}{{2}})}\).

In contrast to the update cost (UC), the estimate for the number of the loss and gradient evaluations required by SPQN methods is not scaled proportionally to \(n_{s}\). This is due to the fact that evaluation of the loss requires a forward pass through the entire network. In contrast, the evaluation of the local gradient, denoted as \(\nabla_{\mathbf{\theta}_{s}}\mathcal{L}\), does not necessitate traversing the entire network as several code optimization strategies can be utilized to avoid unnecessary computations. Moreover, the network's structure can be leveraged to construct local gradient approximations at a lower computation cost. For instance, in the case of ODE-based network architectures, the parallel-in-time strategies [26] could be explored. We do not investigate such code optimization techniques in the present work, as their effectiveness depends heavily on the particular implementation and the choice of network architecture. However, we plan to pursue this avenue of research in the future. Taking all aforementioned points into account, we note that the estimate \(\#g_{e}\) considered in Table 1 is fairly pessimistic.

The summary of computational cost and memory requirements for all aforementioned optimizers can be found in Table 1.

## 5 Numerical experiments

In this section, we describe numerical examples used to test the proposed SPQN methods. In particular, we consider the following benchmark problems.
* **Burgers' equation:** Burgers' equation has the following form: \[\frac{\partial u}{\partial t}+u\nabla u-\nu\nabla^{2}u =0, \forall\ (t,x)\in(0,1]\times(-1,1),\] (5.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\},\] 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:** The steady-state diffusion-advection problem is given in two-spatial dimensions as (5.2) \[-\nabla\cdot\mu\nabla u+\mathbf{b}\cdot\nabla u =f, \forall\ (x_{1},x_{2})\in(0,1)\times(0,1),\] \[u =0, \text{on }\partial\Omega,\] where \(\mathbf{b}=(1,1)^{\top}\). The data is considered to be constant, i.e., \(f=1\). Moreover, the symbol \(\mu\) denotes viscosity. In this work, we employ \(\mu=10^{-2}\).
* **Klein-Gordon:** The nonlinear second-order hyperbolic equation, called Klein-Gordon, is defined as \[\frac{\partial^{2}u}{\partial t^{2}}+\alpha\nabla^{2}u+\beta u+ \gamma u^{2} =f(t,x), \forall\ (t,x)\in(0,12]\times(-1,1),\] (5.3) \[u =x, \forall\ (t,x)\in\{0\}\times[-1,1],\] \[\frac{\partial u}{\partial t} =0, \forall\ (t,x)\in\{0\}\times[-1,1],\] \[u =-\cos(t), \forall\ (t,x)\in(0,12]\times\{-1\},\] \[u =\cos(t), \forall\ (t,x)\in(0,12]\times\{1\},\] where \(u=u(t,x)\) denotes the wave displacement at time \(t\) and location \(x\). The symbols \(\alpha,\beta,\gamma\) denote the predetermined constants. For the purpose of this work, we employ \(\alpha=-1,\beta=0,\gamma=1\) and \(f(t,x):=-x\cos(t)+x^{2}\cos^{2}(t)\). Given this particular choice of the parameters, the analytical solution is given as \(u(t,x)=x\cos(t)\), see [27].

\begin{table}
\begin{tabular}{|l|l|l|l|l|} \hline Method & \(\#\mathcal{L}_{\text{e}}\) & \(\#g_{\text{e}}\) & UC & MC \\ \hline \hline Adam & \(1\) & \(1\) & \(5n\) & \(4n\) \\ \hline L-BFGS & \(1+ls_{\text{its}}\) & \(1\) & \(n+4mn\) & \(n+2mn\) \\ \hline ASPQN & \(1+ls_{\text{its}}+\) & \(2+\) & \(2n+4mn+\) & \(2n+2mn+\) \\ \(k_{s}(\#\mathcal{L}_{\text{e}_{s}})\) & \(k_{s}(\#g_{\text{e}_{s}})\) & \((^{n_{s}/n})k_{s}\text{UC}_{s}\) & \((^{n_{s}/n})\text{MC}_{s}\) \\ \hline MSPQN & \(1+ls_{\text{its}}+\) & \(2+\) & \(2n+4mn+\) & \(2n+2mn+\) \\ \(\sum_{s=1}^{N_{sd}}k_{s}(\#\mathcal{L}_{\text{e}_{s}})\) & \(\sum_{s=1}^{N_{sd}}k_{s}(\#g_{\text{e}_{s}})\) & \(k_{s}\text{UC}_{s}\) & \((^{n_{s}/n})\text{MC}_{s}\) \\ \hline \end{tabular}
\end{table}
Table 1: Estimates for the number of the loss function (\(\#\mathcal{L}_{\text{e}}\)) and gradient evaluations (\(\#g_{\text{e}}\)), update cost (UC) and memory requirements (MC) per iteration and GPU device for various optimizers.
* **Allen-Cahn:** We consider the Allen-Cahn equation of the following form: \[\frac{\partial u}{\partial t}-D\nabla^{2}u-5(u-u^{3}) =0, \forall\ (t,x) \in(0,1]\times(-1,1),\] (10) \[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\},\]

where the diffusion coefficient \(D\) is chosen as \(D=0.001\).

Table 2 summarizes the network architectures used for all four benchmark problems. As an activation function, we employ an adaptive variant of _tanh_[32]. All networks are initialized using the Xavier initialization strategy [20]. To construct the training set, we use a uniform non-adaptive sampling strategy based on quasi-random Hammersley low-discrepancy sequences [56]. The number of collocation points is chosen to be \(10,000\) for all numerical examples.

## 6 Numerical results

In this section, we investigate the numerical performance of the proposed SPQN methods. During our investigation, we monitor the loss functional \(\mathcal{L}\) as well as the quality of PINN's solution (\(u_{\mathtt{NN}}\)), assessed by means of the relative \(L^{2}\)-error, given as

\[\mathcal{E}_{\mathrm{rel}}(u_{\mathtt{NN}},u^{*})=\frac{\|u_{\mathtt{NN}}-u^{ *}\|_{L^{2}(\Omega)}}{\|u_{\mathtt{NN}}\|_{L^{2}(\Omega)}}. \tag{11}\]

Here, \(u^{*}\) denotes an exact solution, obtained by utilizing an analytical expression (Klein-Gordon ) or by solving the underlying PDE using a high-fidelity finite element solver (Burgers', Diffusion-advection, and Allen-Cahn).

### Sensitivity with respect to a varying number of subnetworks and local iterations

We start our numerical investigation by examining the performance of the SPQN methods with respect to a varying number of local iterations (\(k_{s}\)) and subnetworks (\(N_{sd}\)). During these experiments, we consider \(k_{s}\in\{10,50,100\}\) and three different types of network decomposition. In particular, we employ a decomposition into two subnetworks (the minimal decomposition), decomposition by grouping two neighboring layers into one subnetwork, and decomposition of one layer per subnetwork (the maximal layer-wise decomposition).

Figure 3 illustrates the obtained results for different configurations of the ASPQN method. More precisely, we monitor the loss \(\mathcal{L}(\mathbf{\theta})\) and error \(\mathcal{E}_{\mathrm{rel}}\) as a function of number of global iterations. As expected, we observe that the loss decreases monotonically

\begin{table}
\begin{tabular}{|l|l|l|l|l|l|l|} \hline  & \multicolumn{6}{c|}{Parameter type} \\ \cline{2-7} Benchmark problem & Depth & Width & Adam & ASPQN & MSPQN \\  & L & nh & lr & \(N_{sd}\) & \(k_{s}\) & \(N_{sd}\) & \(k_{s}\) \\ \hline \hline Burgers’ & 8 & 20 & \(5\times 10^{-4}\) & 8 & 50 & 8 & 50 \\ \hline Diffusion-advection & 10 & 50 & \(1\times 10^{-4}\) & 10 & 10 & 10 & 10 \\ \hline Klein-Gordon & 6 & 50 & \(1\times 10^{-3}\) & 6 & 50 & 6 & 50 \\ \hline Allen-Cahn & 6 & 64 & \(2.5\times 10^{-4}\) & 6 & 50 & 6 & 50 \\ \hline \end{tabular}
\end{table}
Table 2: Summary of network architecture and optimizer’s parameters used for numerical experiments reported in subsection 6.2.
due to the line-search globalization strategy. We also observe that the configuration of ASPQN with maximal decomposition consistently outperforms other configurations in terms of the error \(\mathcal{E}_{\mathrm{rel}}\). The observed behavior is independent of the number of local iterations (\(k_{s}\)) and the benchmark problem. This can be attributed to the fact that more extensive decoupling of the parameters allows us to obtain more localized learning rates and Hessian approximations, which in turn enables capturing of the underlying nonlinearities more accurately.

Figure 4 demonstrates the performance of the ASPQN method by means of the upper bound on the number of effective gradient evaluations (\(\#g_{e}\)) as well as the estimated update cost UC, both discussed in section 4. We again observe that the most efficient variants of the ASPQN method employ the maximal layer-wise decomposition. Moreover, the lowest \(\mathcal{E}_{\mathrm{rel}}\) is typically achieved using 50 or 100 local iterations.

In the case of the MSPQN method, the obtained results are reported in Figure 5 and Figure 6. As we can see from Figure 5, using a larger amount of subnetworks and

Figure 3: The mean loss \(\mathcal{L}(\boldsymbol{\theta})\) (dotted lines) and relative error \(\mathcal{E}_{\mathrm{rel}}\) (solid lines) for ASPQN. The results are obtained for varying numbers of local iterations, i.e., \(k_{s}\in\{10,50,100\}\) (from left to right). The average is obtained over 10 independent runs. Note that axes are scaled differently for different configurations of the ASPQN method.
higher values of \(k_{s}\) gives typically rise to a faster convergence, i.e., a fewer number of iterations is required to obtain a comparably equally accurate solution. By examining results reported in Figure 6, we see that the best-performing variants in terms of UC are ones with \(k_{s}=50\) or \(k_{s}=100\), and with the maximal layer-wise decomposition. This behavior is expected since the local contributions to UC decrease with an increasing number of subnetworks. However, in the context of the \(\#g_{e}\) metric, the obtained results are fairly different. As we can notice, variants with fewer subnetworks are often more effective, as the \(\#g_{e}\) estimate does not scale proportionally with the size of the subnetwork, i.e., the number of trainable subnetworks' parameters.

### SPQN methods versus the state-of-the-art optimizers

In this section, we compare the convergence and performance of the proposed SPQN methods with state-of-the-art Adam and L-BFGS optimizers. More precisely, we compare all considered optimizers by means of the \(\#g_{e}\) and UC estimates. Moreover, we also report the execution time, although our implementation of the SPQN algorithms is not

Figure 4: The mean computational cost of the ASPQN method. The results are obtained for varying numbers of local iterations, i.e., \(k_{s}\in\{10,50,100\}\) (from left to right). Average is obtained over 10 independent runs.
yet optimized, unlike PyTorch implementation of Adam and L-BFGS2. Throughout these experiments, we configure the SPQN methods according to the specifications summarized in Table 2. The learning rate for the Adam optimizer is chosen by carrying out a thorough hyper-parameter search. In particular, we choose a learning rate (reported in Table 2), which allows us to obtain the most accurate PDE solution, thus a solution with the lowest \(\mathcal{E}_{\mathrm{rel}}\). For the L-BFGS optimizer, we employ the same memory size (\(m=3\)), line-search strategy, and momentum setup as for the L-BFGS methods employed with the SPQN framework.

Footnote 2: We have adjusted PyTorch implementation of the L-BFGS method to take into account the momentum.

As we can see from the obtained results, reported in Figure 7, the SPQN methods give rise to more accurate solutions, i.e., solutions with smaller \(\mathcal{E}_{\mathrm{rel}}\). The differences in the quality of an obtained solution are on average of

Figure 5: The mean loss \(\mathcal{L}(\mathbf{\theta})\) (dotted lines) and relative error \(\mathcal{E}_{\mathrm{rel}}\) (solid lines) for MSPQN. The results are obtained for varying numbers of local iterations, i.e., \(k_{s}\in\{10,50,100\}\) (from left to right). The average is obtained over 10 independent runs. Note that axes are scaled differently for different configurations of the MSPQN method.
The Diffusion-advection problem is particularly interesting as the L-BFGS method stagnates for this example. In contrast, the variants of the SPQN optimizers are able to reach the \(\mathcal{E}_{\mathrm{rel}}\) close to \(10^{-2}\). Numerical evidence suggests that the L-BFGS stagnation is caused by the line-search method providing very small step sizes. In contrast, when using the SPQN methods, a very small step size was observed only for one particular subnetwork at the beginning of the training process. This suggests that using the DD approach enables the usage of localized learning rates and rebalances the global nonlinearities.

In terms of the computational cost, we can also observe from Figure 7, that the SPQN methods are significantly more efficient in terms of \(\#g_{e}\) and UC. The difference is more prevalent in the case of the UC metric, as this estimate scales with the number of subnetworks. We can also see that the UC metric translates closely to the computational time (the last column of Figure 7). Thus, in order to achieve a solution of comparable quality, the SPQN methods are significantly faster than Adam and L-BFGS optimizers. To highlight the difference, we also report the minimum mean \(\mathcal{E}_{\mathrm{rel}}\) reached by the L-BFGS optimizer and the average time required to reach that solution.

Figure 6: The mean computational cost of MSPQN method. The results are obtained for varying numbers of local iterations, i.e., \(k_{s}\in\{10,50,100\}\) (from left to right). Average is obtained over 10 independent runs.
tion by all considered optimizers34. As we can see from Table 3, the computational time required by ASPQN and MSPQN is significantly lower than the computational time required by the L-BFGS method. In particular, the MSPQN optimizer enjoys an average speedup of a factor of 10 across all benchmark problems. For the ASPQN optimizer, the speedup is even larger, by an average factor of approximately 28. This speedup is however obtained with an increased amount of computational resources, as the optimizer runs in parallel, i.e., using multiple GPU devices. Thus, for a fixed amount of computational resources, MSPQN can be considered to be a more efficient optimizer. However, the ASPQN method allows for a larger overall speedup by leveraging model parallelism.

Footnote 3: Comparison with the Adam optimizer is skipped, as it never reaches \(\mathcal{E}_{\mathrm{rel}}\) as low as the L-BFGS optimizer.

Figure 7: Performance comparison of different optimizers for Burgers’, Klein-Gordon, Diffusion-advection, and Allen-Cahn problems (from top to bottom). The mean relative error \(\mathcal{E}_{\mathrm{rel}}\) obtained over 10 independent runs.
## 7 Conclusion

In this work, we proposed novel additive and multiplicative preconditioning strategies for the L-BFGS optimizer with a particular focus on PINNs applications. The preconditioners were constructed by leveraging the layer-wise decomposition of the network, thus tailoring them to the characteristics of DNNs. The multiplicative preconditioner was designed to enhance the convergence of the L-BFGS optimizer in the context of single GPU computing environments. In contrast, the additive preconditioner provided a novel approach to model parallelism with a focus on multi-GPU computing architectures. To evaluate the effectiveness of the proposed training methods, we conducted numerical experiments using benchmark problems from the field of PINNs. A comparison with Adam and L-BFGS optimizers was performed and demonstrated that the SPQN methods can provide a substantial reduction in terms of training time. In particular, using the MSPQN optimizer, we obtained an average speedup of a factor of 10, while for the ASPQN optimizer, the average speedup factors were approximately 28. Moreover, we have also demonstrated that by using the proposed SPQN methods, one can obtain a significantly more accurate solution for the underlying PDEs.

We foresee the extension of the presented work in several ways. For instance, it would be interesting to investigate different network decompositions, as well as to incorporate overlap and coarse-level acceleration into the proposed algorithmic framework. Furthermore, developing strategies for load balancing and optimization of the subnetwork's loss/gradient evaluations could further enhance the performance of the proposed algorithms. Moreover, the applicability of the proposed SPQN methods could be extended beyond the PINNs to a variety of different deep-learning tasks, e.g., operator learning using DeepOnets [48, 49], or natural language processing using transformers [72]. Lastly, in the context of PINNs, the current proposal could be combined with the methods that consider the decomposition of the computational domain (data space).
