Learning Interface Conditions in Domain Decomposition Solvers
Abstract
Domain decomposition methods are widely used and effective in the approximation of solutions to partial differential equations.
Yet the optimal construction of these methods requires tedious analysis and is often available only in simpliﬁed, structured-grid settings, limiting their use for more complex problems.
In this work, we generalize optimized Schwarz domain decomposition methods to unstructured-grid problems, using Graph Convolutional Neural Networks (GCNNs) and unsupervised learning to learn optimal modiﬁcations at subdomain interfaces.
A key ingredient in our approach is an improved loss function, enabling effective training on relatively small problems, but robust performance on arbitrarily large problems, with computational cost linear in problem size.
The performance of the learned linear solvers is compared with both classical and optimized domain decomposition algorithms, for both structured- and unstructured-grid problems.
Introduction
Domain decomposition methods (DDMs) [1, 2, 3] are highly effective in solving the linear and nonlinear systems of equations that arise from the numerical approximation of solutions to partial differential equations (PDEs).
While most effective on elliptic boundary-value problems, DDMs can also be applied to nonlinear problems, either using their nonlinear variants or successively solving linearizations
Time-dependent problems are normally solved by using a time stepping algorithm in the time domain for implicit methods, which require the solution of a spatial problem for each time step.
Of these methods, Schwarz methods are particularly popular given their relative simplicity and ease of parallelization.
The common theme is to break the global problem into subproblems, derived either by projection or by discretizing the same PDE over a physical subdomain, and to use solutions on the subdomains as a preconditioner for the global discretization.
Classical Schwarz methods generally make use of Dirichlet or Neumann boundary conditions for these subdomain problems, while Optimized Schwarz Methods (OSMs) aim to improve the convergence of the algorithm by using more general interface conditions [4].
Notably, [5] demonstrates that optimal, but non-local, interface conditions exist for more general decompositions.
Much of the OSM literature considers only one-level additive Schwarz methods, although multilevel extensions do exist.
For one-level methods (i.e., domain decomposition approaches without a “coarse grid”), restricted additive Schwarz (RAS) approaches [6] are arguably the most common; optimized restricted additive Schwarz (ORAS) methods are considered in [7].
The OSM idea has also been extended to asynchronous Schwarz methods [8], where the computations on each subdomain are done using the newest information available in a parallel computing environment without synchronizing the solves on each subdomain.
With a recent focus on machine learning (ML) techniques for solving PDE systems [9, 10], there is also effort to apply learning-based methods to improve the performance of iterative solvers for PDEs, including DDM and algebraic multigrid (AMG) methods.
Within AMG methods, ML techniques have been applied to learning interpolation operators [11, 12] and to coarse-grid selection in reductionbased AMG [13].
Of particular note here is the loss function employed in [11, 12], where they use unsupervised learning to train a graph neural network to minimize the Frobenius norm of the error-propagation operator of their iterative method.
Within DDM, signiﬁcant effort has been invested in combining ML techniques with DDM, as in [14], where two main families of approaches are given: 1) using ML within classical DDM methods to improve convergence, and 2) using deep neural networks, such as Physics Informed Neural Networks (PINNs) [9], as a discretization module and solver for DDM problems.
In [15], a fully connected neural network is used to predict the geometric locations of constraints for coarse space enhancements in an adaptive Finite Element Tearing and
Interconnecting-Dual Primal (FETI-DP) method.
Using the continuous formulation of DDM, the so-called D3M [16] uses a variational deep learning solver, implementing local neural networks on physical subdomains in a parallel fashion.
Likewise, Deep-DDM [17] utilizes PINNs to discretize and solve DDM problems, with coarse space corrections [18] being used to improve scalability.
In this paper, we advance DDM-based solvers by developing a framework for learning optimized
Schwarz preconditioners.
A key aspect of this is reconsidering the loss function to use a more effective relaxation of the ideal objective function than Frobenius norm minimization [11, 12].
Moreover, the approach introduced here offers an opportunity to reconsider existing limitations of optimized
Schwarz methods, where optimal parameter choice is based on Fourier analysis and requires a highly regular subdomain structure, such as in the classical cases of square domains split into two equal subdomains or into regular grids of square subdomains.
Our framework learns the optimized Schwarz parameters via training on small problem sizes, in a way that generalizes effectively to large problems, and in a way that allows us to consider both irregular domains and unstructured grids, with no extraordinary limitations on subdomain shape.
Furthermore, the evaluation time of our algorithm scales linearly with problem size.
This allows signiﬁcant freedom in deﬁning optimized Schwarz methods, in comparison to classical approaches, allowing us to explore the potential beneﬁts of optimized Schwarz over classical (restricted) additive methods on unstructured grids for the ﬁrst time.
Background
Let ⌦ ⇢ R 2 be an open set, and consider the positive-deﬁnite Helmholtz problem
[Equation]
with inhomogeneous Dirichlet conditions imposed on the boundary @ ⌦ .
In (1), the parameter ⌘ >represents a shift in the Helmholtz problem.
In the numerical results below, we consider both ﬁnite-difference discretizations of (1) on regular grids, as well as piecewise linear ﬁnite-element (FE) discretizations on arbitrary triangulations.
In both cases, we denote the set of degrees of freedom as
D , and note that these are in a one-to-one correspondence with the nodes of the underlying mesh.
Consider a decomposition of D into non-overlapping subdomains D
i , i 2 { 1 , 2 , . . .
, S } such that each node is contained within exactly one subdomain D
i , yielding [ D 0 i = D .
In this subdomain notation, the superscript denotes the amount of overlap in the subdomains, which is zero for the non-overlapping subdomains that we ﬁrst consider.
Let R
i be the restriction operator onto the set of degrees of freedom (DoFs) in D
i , and let
R 0 i
T be the corresponding extension operator from D {+
 i D 11+} D
[Figure] into set D . Then, an FE discretization of the Helmholtz problem leads to a linear system of the form
Ax = b , where A is the global stiffness matrix and
A
i =
R
i A
R 0 i
T is the subdomain stiffness matrix for D
i .
We note that alternate deﬁnitions to the Galerkin projection for A {+
+} i are possible, and are commonly considered in optimized Schwarz settings (as noted below).
In the case of restricted additive Schwarz (RAS) [6], the subdomains are extended to allow for overlap: nodes near the “boundary” of their subdomain are potentially included in two or more subdomains.
We denote the amount of overlap by  2 N , deﬁning subdomains D  i recursively, by
D  i =
D   1 i
[Equation]
for  > 0 .
The conventional RAS preconditioner is deﬁned as
[Equation]
where R  i is the standard restriction operator to subdomain ⌦  i ,
˜ R  i is a modiﬁed restriction operator from D to the DoFs in D  i that takes nonzero values only for DoFs in D
i , and A  i =
( R  i ) T AR  i .
Figure 1 shows an example unstructured grid with two subdomains and overlap  = 1 .
In an optimized Schwarz setting, we modify the subdomain systems, A  i .
Rather than using a
Galerkin projection onto D  i , we rediscretize (1) over the subdomain of ⌦ corresponding to these
DoFs, imposing a Robin-type boundary condition on the boundary of the subdomain.
We deﬁne this matrix to be ˜
A  i =
A i + L i , where A i is the term resulting from discretization of (1) with Neumann boundary conditions, and L i is additional the term resulting from the Robin-type condition, as in:
Dirichlet: u = g D [Equation] x ) ,
[Equation]
where ~ n is the outward normal to the edge on the boundary and g denotes inhomogeneous data.
The matrix L i has the same dimensions as A i , so that ˜
A  i is well-deﬁned.
However, it has a signiﬁcantly different sparsity pattern, with nonzero entries only in rows/columns corresponding to nodes on the boundary of subdomain D  i .
In practice, we identify a cycle or path in the graph corresponding to A i with the property that every node in the cycle is on the boundary of D  i but not the boundary of D (the discretized domain), and then restrict the nonzeros in L i to the entries corresponding to the edges in this cycle/path (including self-edges, corresponding to entries on the diagonal of L i ).
Using this notation, the ORAS preconditioner can be written as
[Equation]
Our aim is to learn the values in the matrices L i , aiming to outperform the classical choices of these values predicted by Fourier analysis, in the cases where those values are known, and to learn suitable values for cases, such as ﬁnite-element discretizations on unstructured grids, where no known optimized Schwarz parameters exist.
We optimize the values for the case of stationary (Richardson) iteration, but evaluate the performance of the resulting methods both as stationary iterations and as preconditioners for FGMRES.
Graph Neural Networks (GNNs): Applying learning techniques to graph structured data necessitates stepping beyond multilayer perceptron (MLP) and conventional convolutional neural networks (CNN) to a type of network that leverages the underlying graph nature of the problem, namely graph convolutional neural networks (GCNNs).
GCNNs are typically divided into two categories: spectral and spatial [19].
Spectral GCNNs, ﬁrst introduced by Bruna et al.
[20], consider a graph Laplacian eigenbasis and deﬁne convolution as diagonal operators.
As such, spectral GCNN methods suffer from time complexity problems due to the necessity for the eigendecomposition of the Laplacian matrix.
Nevertheless, in follow-up work [21, 22], remedies have been proposed to mitigate this.
Unlike spectral methods, spatial GCNNs consider local propagation of information in graph domains as a convolution graph.
One popular framework is the message passing neural network (MPNN) [23], which is based on sharing information among neighbor nodes in each round of a convolution pass.
This has been generalized [24] by introducing a conﬁgurable graph network block structure consisting of node and edge convolution modules and a global attribute.
In an effort to alleviate computational complexity of GCNNs [25], topology adaptive graph convolution networks (TAGCN) can be constructed by deﬁning learnable ﬁlters.
This is not only computationally simpler, but also allows for adapting to the topology of the graphs when scanning them for convolution.
Method 3.1 Optimization problem and the loss function
Throughout this paper, we use k · k to denote the ` 2 norm of a matrix or vector, k A k F for the Frobenius norm of A , and ⇢ ( A ) as the spectral radius of A .
The optimization problem that we seek to solve is to ﬁnd optimal values for the entries in the matrices L i , constrained by given sparsity patterns, to minimize ⇢ ( T ) ,
where T = I
M ORAS A is the error-propagation operator for the stationary iteration corresponding to M ORAS deﬁned in (4).
The spectral radius ⇢ ( T ) corresponds to the asymptotic convergence factor of the stationary iteration, giving a bound on the asymptotic convergence of the method.
Formally deﬁned in terms of the extremal eigenvalue of T T T (since T is not symmetric), direct minimization of ⇢ ( T ) is difﬁcult since backpropagation of an eigendecomposition is numerically unstable [26].
To overcome this, Luz et al.
[12] propose to relax the minimization of ⇢ ( T ) (for a similar AMG preconditioner) to minimizing the Frobenius norm of T , k T k F .
In our experiments, however, we ﬁnd that this is insufﬁcient, leading to preconditioners that do not scale.
One reason is that while the Frobenius norm is an upper bound on the spectral radius, it does not appear to be a suitably “tight” bound for use in this context (see Section 4.2 and Figure 6).
Instead, we use a relaxation inspired by Gelfand’s formula, that ⇢ ( T ) =
lim K !
1 k
[Equation]
k 1 K , and the common bound that
[Equation]
for some ﬁnite K 2 N .
This results in the optimization problem
[Equation]
[Equation]
3.2 Numerical evaluation of the loss function
We denote the action of evaluating the GNN by f ( ✓ ) (where ✓ represents the network parameters), and consider a discretized problem with DoF set D , of size n .
The set D can be decomposed into subdomains either by using ﬁxed geometric choices of the subdomain (e.g., for ﬁnite-difference discretizations), using the METIS graph partitioner [27], or a k -means-based clustering algorithm (best known as Lloyd’s algorithm which has O ( n ) time complexity) [28, 29].
For unstructured problems, we use a k -means-based algorithm, decomposing D to subdomains D i for i = 1 , 2 , . . .
, S , with overlap  ; see Supplementary Materials for details.
The GNN then takes D and its decomposition as inputs, as well as sparsity constraints on the matrices L i for i = 1 , 2 , . .
, S , and outputs values for these matrices:
[Equation]
Using the learned subdomain interface matrices, we then obtain the modiﬁed MLORAS (Machine
Learning Optimized Restricted Additive Schwarz) operator, M ( ✓ ) ORAS , simply using ˜ A
i =
A i + L ( ✓ ) i in (4).
We denote the associated error propagation operator by T ( ✓ ) =
I  M
( ✓ ) ORAS A .
While Gelfand’s formula and the associated upper bound in (5) are valid in any norm, it is natural to consider them with respect to the ` 2 norm in this setting.
However, this raises the same issue as encountered in [12], that it generally requires an eigendecomposition to compute the norm.
To avoid this, we use a stochastic sampling of
T ( ✓ )
K  , generated by the sample set X 2 R n ⇥ m for some m 2 N , given as X
[Equation]
Here, we randomly select m points uniformly on a unit sphere in R n , which can be done using the method in [30].
We then deﬁne
[Equation]
taking each column of X as the initial guess to the solution of the homogeneous problem Ax = 0 and taking K steps of the stationary algorithm to generate
T ( ✓ )
K x j .
Since we normalize each column of X to have k x j k = 1 , the value of
T ( ✓ )
K x j  serves as a lower bound for
T ( ✓ )
K  .
Thus, taking the maximum of the values in Y provides a practical loss function that we use below, deﬁning
[Equation]
We note similarities between this loss function and that used in [31], but that we are able to use the maximum of Y ( ✓ ) (giving a better approximation to the norm) in our context instead of averaging.
Ultimately, the cost of our algorithm depends strongly on the chosen values of m and K .
For sufﬁciently large values of m , we now show that the maximum value in Y ( ✓ ) is an arbitrarily good approximation to
[Equation]
Theorem 1.
For any nonzero matrix T , ✏ > 0 , and  < 1 , there exist M, K 2 N such that for any m > M , if one chooses m points, x j , uniformly at random from { x 2 R n , k x k = 1 } , then
[Equation]
Proof.
According to Gelfand’s theorem, there exists L
N such that 8 ` > L ,    ⇢ ( T )  sup x :
[Equation]
< ✏ 2 .
Take any K  L and let ˜ ✏ = ✏ 2 k T K k 1 K .
Since R n is ﬁnitedimensional, there exists an x ⇤ 2 R n , k x ⇤ k
= 1 such that sup x : [Equation]
T K x k =
k T K x ⇤ k .
Denote the volume of the surface of the n -dimensional sphere of unit radius around the origin in R n by C tot , and the volume of the region on this sphere within radius ˜ ✏ K of x ⇤ by C ˜ ✏ ,K .
Given  < 1 , let M 2 N satisfy M
[Equation]
Then,
[Equation]
Thus, with probability of at least 1   , we expect at least one point from a selection of m > M points uniformly distributed on the sphere of radius unit radius to be in C ˜ ✏ ,K .
Let that point be x r , giving
T K
[Equation]
[Equation]
T K
[Equation]
[Equation]

T K x ⇤
T K
[Equation]
1 K (13) 
[Equation]
[Equation] (15) using Lemma 3 (see Supplementary Materials) and the reverse triangle inequality. Finally, by the triangle inequality, with probability of at least 1   , we have:
[Equation]
[Equation]
Remark.
According to [32], since the optimal interface values are Turing-computable, if the depth of the GNN is at least the diameter of the graph, and a TAGConv layer followed by a feature encoder is Turing-complete, the optimal interface values can be learned.
In our setting, for a problem on a structured grid of size N ⇥ N with two identical rectangular subdomains, this implies that the GNN will be able to learn the optimal interface values given, if and only the GNN has depth at least 2 N , has deep enough feature encoders, and the width of the layers is unbounded. Remark.
Theorem 1 guarantees convergence of the loss function to the spectral radius, in the limits of many samples and many stationary iterations.
To the best of our knowledge, such a guarantee is not known for the previous loss functions used in the area [12, 11].
Moreover, there are substantial improvements in the numerical results using the new loss function in comparison to that of [12], as shown in Figure 6.
Theorem 2.
Assuming bounded subdomain size, the time complexity to evaluate the optimal Schwarz parameters using our method is O ( n ) , where n is the number of nodes in the grid.
Proof.
Given bounded Lloyd subdomain size and ﬁxed number of Lloyd aggregation cycles, subdomain generation has O ( n ) time complexity [28] (see the Supplementary Material).
To evaluate each TAGConv layer, one computes y =
P
L ` =1 G ` x ` + b 1 n , where L is the number of node features,
G ` 2 R n ⇥ n is the graph ﬁlter, b is a learnable bias, and x ` 2 R n are the node features.
Moreover, the graph ﬁlter is a polynomial in the adjacency matrix M of the graph: G ` =
P
Jj =
g ` ,j M j where J is a constant and g ` ,j are the coefﬁcients of ﬁlter polynomial.
Since the graph has bounded node degrees, it implies that M is sparse and the action of M j has O ( n ) cost, and therefore, the full TAGConv evaluation also has O ( n ) cost.
Moreover, the cost of edge feature and the feature networks are O ( n ) , resulting in overall O ( n ) cost.
3.3 Training Details
We use a graph neural network based on four TAGConv layers and ResNet node feature encoders consisting of eight blocks after each TAGConv layer; see Supplementary Materials for more details on the structure of the GNN, and Section 4.3 for an ablation study.
The training set in our study consists of 1000 unstructured grids with piecewise linear ﬁnite elements and with grids ranging from 90–850 nodes (and an average of 310 nodes).
The grids are generated by choosing either a regular grid (randomly selected 60% of the time) or a randomly generated convex polygon; pygmsh [33] is used to generate the mesh on the polygon interior.
A sample of the training grids is shown in Figure 2.
We train the GNN for four epochs with a mini batch size of 25 using the ADAM optimizer [34] with a ﬁxed learning rate of 10  4 .
For the numerical evaluation of the loss function (10) we use
K = 4 iterations and m = 500 samples.
The code 1 was implemented using PyTorch Geometric [35],
PyAMG [36], and NetworkX [37].
All training was performed on an 8-core i9 Macbook Pro using
CPU only.
Moreover, we train two special networks for use in Section 4.1.
The ﬁrst is a “Brute Force” network, which is a single-layer neural network without an activation function, trained only on a structured [Figure] 10 ⇥ 10 grid with two identical subdomains using the ADAM optimizer.
The purpose of training this is to obtain the optimal interface values as a benchmark for comparison, including against our method.
The second is the same network used for our method, MLORAS, but overtrained on only the problem in Section 4.1, to understand the learning capabilities of the method and choice of GNN architecture.
Results 4.1 Two-subdomain structured grids
We ﬁrst consider rectangular structured grids with two subdomains.
Although restrictive, these problems allow us to directly compare to existing OSM parameters, which are only available for structured grids and exactly two subdomains.
We follow [7] and consider the Helmholtz problem (1) on the unit square with ⌘ = 1 and homogeneous Dirichlet boundary conditions.
We discretize the problem on an N ⇥ N rectangular grid with h = 1 / ( N +1) grid spacing, using the standard ﬁve-point ﬁnite difference stencil.
Figure 3 shows the results of different methods on a 10 ⇥ 10 structured grid with two identical subdomains.
We see that the overtrained MLORAS network learns interface parameters that outperform the Restrictive Additive Schwarz (RAS) [6] method; more signiﬁcantly, MLORAS also outperforms the existing Optimized-RAS (ORAS) methods that were analytically derived for this speciﬁc problem (zeroth-order OO0 and second-order OO2 from [7]).
Moreover, in these results, the MLORAS network almost reaches the performance of the “Brute Force” network (obtained by directly optimizing all interface values for this single structured grid), indicating that using a GNN to encode the optimal interface values does not signiﬁcantly restrict the search space.
[Figure]
Brute Force [Figure] MLORAS (Ours) [Figure]
[Figure]
ORAS OO
[Figure]
To understand the performance of the methods in more detail, Figure 4 plots the interface values output by the brute force optimization, the MLORAS network, and the OO2 ORAS algorithm.
This shows that the MLORAS network is choosing interface values very close to those directly optimized by the brute force method, unlike those selected by the OO2 method.
4.2 Unstructured grids
To evaluate the performance of the MLORAS network on unstructured grids, we considerunstructured triangular grids in 2D with sizes ranging from about 90 to 40 000 nodes.
These grids are deﬁned on convex subsets of (0 , 1) ⇥ (0 , 1) ; we solve the Helmholtz problem (1) on these domains with ⌘ = 1 and homogeneous Dirichlet boundary conditions, discretized with piecewise-linear ﬁnite elements.
Example convergence plots are shown in Figure 5, where our method (MLORAS) is compared to RAS (Restricted Additive Schwarz [6]), AS (Additive Schwarz), and Jacobi methods as a preconditioner for FGMRES.
Here we see that the MLORAS network is able to learn optimized interface parameters for unstructured grids that outperform RAS, and that MLORAS can scale to problems that are much larger than those in the training set, which are all below n = 1000 nodes.
Importantly, MLORAS retains an advantage over RAS even as the grid size increases.
Figure 6 shows the performance of three different methods across all unstructured test grids.
The three methods are: (1) “F-norm optimized”: the same network as MLORAS, but trained to directly optimize the Frobenius norm of the error propagation matrix, k A k F , (2) the RAS [6] method, and (3) the MLORAS method.
We do not show ORAS results here, as it cannot be applied to unstructured grids.
Figure 6 reveals two important facts.
First, MLORAS consistently outperforms RAS over the entire testing set, both as an FGMRES preconditioner and as a stationary algorithm, and this [Figure] R E S s t e p s [Figure]
[-Second, we see that the Frobenius norm is indeed a 8 Residual block length 31 .31 .32 .
-][Figure]
Grid
[Figure]
[Figure]
[Figure] [Figure] S t a t i o n a r y e rr o r F-norm optimized
RAS
{+MLORAS (Ours) 
+}
{+Grid size 
+}
F r o b e n i u s n o r m [Figure] advantage is maintained even for large grids. Second, we see that the Frobenius norm is indeed a worse choice for the loss function than our new loss in (6).
We see this from the fact that MLORAS has a worse Frobenius norm than either of the other two methods, but it has the best convergence rate.
In addition, when we explicitly optimize the Frobenius norm (the “F-norm optimized” method), we see that we do obtain the lowest Frobenius norm, but this translates to the worst convergence rate.
4.3 Ablation study
To understand the impact of the network structure on performance, we conduct an ablation study by varying the ResNet block length and the number of TAGConv layers.
In each case, we train a network on ﬁve different training sets, each consisting of 1000 grids generated as described in Section 3.3.
The trained networks are tested on 50 unstructured grids, each with 2400 to 2600 nodes.
For each architecture, the mean performance is computed, together with error bars from the ﬁvefold repetition, with results shown in Figure 7.
We see that a higher residual block length is always better, but that the optimal number of TAGConv layers is 4, and using more layers actually decreases performance.
This phenomenon has been observed in other contexts (e.g., [38] and [39]) and can be attributed to over-smoothing in GNNs.
The use of residual blocks in our architecture is, thus, important to allow us to increase network depth without needing to increase the number of GNN layers.
Conclusion
We propose an unsupervised learning method based on Graph Convolutional Neural Networks (GCNNs) for extending optimized restricted additive Schwarz (ORAS) methods to multiple subdomains and unstructured grid cases.
Our method is trained with a novel loss function, stochastically minimizing the spectral radius of the error propagation operator obtained using the learned interface values.
The time complexity of evaluating the loss function, as well as obtaining the interface values using our neural network are all linear in problem size.
Moreover, the proposed method is able to outperform
ORAS, both as a stationary algorithm and preconditioner for FGMRES on structured grids with two subdomains, as considered in the conventional ORAS literature.
On more general cases, such as unstructured grids with arbitrary subdomains of bounded size, our method outperforms RAS consistently, both as a stationary algorithm and preconditioner for FGMRES.
The main limitations of the current work are that the method was studied for two PDE cases, namely the Helmholtz problem in the main paper and the non-uniform Poisson problem in Appendix D. We also defer the study of nonlinear or time-dependent PDEs to future work. Acknowledgement
The work of S.P.M. was partially supported by an NSERC Discovery Grant. The authors thank the referees for their insightful comments. The authors have no competing interests to declare.