NeuralPCG: Learning Preconditioner for Solving Partial Differential Equations with Graph Neural Network
Authors: Yichen Li
Tao Du
Peter Yichen Chen
Wojciech Matusik
Venue: Submitted to ICLR 2023 (2023)
Keywords: Physics Simulation
Graph Neural Network
Applied Mathematics

Abstract: Fast and accurate partial differential equation (PDE) solvers empower scientific and engineering research. Classic numerical solvers provide unparalleled accuracy but often require extensive computation time. Machine learning solvers are significantly faster but lack convergence and accuracy guarantees. We present Neural-Network-Preconditioned Conjugate Gradient, or NeuralPCG, a novel linear second-order PDE solver that combines the benefits of classic iterative solvers and machine learning approaches. Our key observation is that both neural-network PDE solvers and classic preconditioners excel at obtaining fast but inexact solutions. NeuralPCG proposes to use neural network models to \emph{precondition} PDE systems in classic iterative solvers. Compared with neural-network PDE solvers, NeuralPCG achieves converging and accurate solutions (e.g.,1e-12 precision) by construction. Compared with classic solvers, NeuralPCG is faster via data-driven preconditioners. We demonstrate the efficacy and generalizability of NeuralPCG by conducting extensive experiments on various 2D and 3D linear second-order PDEs.

Under review as a conference paper at ICLR 2023
NEURALPCG: LEARNING PRECONDITIONERS FOR SOLVING PARTIAL DIFFERENTIAL EQUATIONS WITH GRAPH NEURAL NETWORKS
Anonymous authors Paper under double-blind review
ABSTRACT
Fast and accurate partial differential equation (PDE) solvers empower scientific and engineering research. Classic numerical solvers provide unparalleled accu- racy but often require extensive computation time. Machine learning solvers are significantly faster but lack convergence and accuracy guarantees. We present Neural-Network-Preeone
Conjugate Gradient, or NeuralPCG, novel lin- car second-order PDE solver that combines the benefits of classic iterative solvers and machine learning approaches. Our key observation is that both neural-network PDE solvers and classic preconditioners excel at obtaining fast but inexact solu- tions. NeuralPCG proposes to use neural network models to precondition PDE systems in classic iterative solvers. Compared with neural-network PDE solvers, NeuralPCG achieves converging and accurate solutions (e.g.. le- precision) by construction. Compared with classic solvers, NeuralPCG faster via data-driven preconditioners. We demonstrate the efficacy and generalizability of NeuralPCG by conducting extensive experiments on various 2D and 3D linear second-order PDESs"
INTRODUCTION
Partial differential equations (PDEs) are fundamental mathematical models with broad applications in science and engineering, for example, the Navier-Stokes equation in fluid dynamics, Poisson's equation in computational geometry, and the Black SacklSSholes equation in mathematical finance. De- spite their powerful modeling ability and wide applications, is notoriously difficult to find analyti- cal solutions general PDE. Therefore, numerical solvers have long been the mainstay of solving PDEs. Classic PDE solvers provide accurate solutions well-understood PDEs but typically the cost of long computation time. Speeding up these classic solvers is non-trive non-trivial and often requires complex numerical techniques, c.g., multi-grid methods (Briggs et al. 2000), domain decomposition (Smith 1997). and model reduction (Holmes et al.) 2012). Recently, several pioneering works (Li et al. Sanchez Saaaaannn-Gnnzalleee Gonzalez et al. 2020) introduce machine learning techniques to solving PDEs, particularly in the field of physics simulation. While this line of methods typically outperforms clas- sic solvers in speed by a large margin, it struggles with converging into a highly precise solution (e.g., le 12). A lack of theoretical analysis convergence and accuracy inhibits neural-network PDE solvers' applications in mechanical engincering, structure analysis, and aerodynamics, where precise PDE solutions have a higher priority than fast yet inexact results.
This work proposes NeuralPCG, a novel and hybrid method that combines the benefits of classic and machine-learning PDE solvers. Our key observation is that neural-network solvers are fast at estimating PDE solutions with a low to moderate accuracy. This property aligns with the pre- conditioning technique in numerical methods, which uses an easy-to-solve approximation of original PDE to speed up numerical solvers. Based on this intuition, NeuralPCG proposes to learn a neural network that preconditions a classic iterative solver. This is in sharp contrast to existing netral-network PDE solvers that replace the numerical solver entirely with neural network. The backbone of NeuralPCG remains to be a classic solver, empowering it to inherit convergence and
'Supplementary videos are available on the project webpage: https: t/ttps://s/ tes googl com/vies neuralper pog

---

Under review as a conference paper at ICLR 2023
accuracy guarantees. Moreover, the learned preconditioner even outperforms classic preconditioners because it adapts to data distributions from target PDE applications.
To demonstrate the efficacy of NeuralPCG, we evaluate it set of 2D and 3D representative lin- car second-order PDEs. We compare its performance with two baselines of very different natures:
(1) MeshGraphNet (MGN), state-of-the-art neural-network PDE solver (Pfaff al., 2020); (2)
the preconditioned conjugate gradient (PCG) algorithm with classic preconditioners; Our experi- ments show that NeuralPCG has its unique advantage over both baselines: Compared with MGN and similar machine-leariing solvers, NeuralPCG generates convergent and accurate solutions by construction and avoids eтrоr accumulation. On the other hand. NeuralPCG also outperforms clas- sic PCG solvers in speed because our method learns a preconditioner tailored to the training data distribution, which classic preconditioners typically do not exploit. Finally, NeuralPCG generalizes to physics parameters or meshes with moderate differences (5m) from the training data and is robust to outliers by design.
In summary, our work makes the following contributions:
1. We propose NeuralPCG, novel and efficient solver for linear second-order PDEs that combines the benefits of classic and machine-learning PDE solvers.
2. We present framework for studying the performance of preconditioners in classic iterative solvers, including benchmark problem set with carefully designed performance metrics.
3. We conduct extensive experiments to evaluate the efficacy and generalizability of Neu- ralPCG and demonstrate its ad advantages over existing methods.
2 RELATED WORK
Numerical PDE solvers Classic PDE solvers discretize the continuous PDEs into a numerical system. Popular discretization schemes include finite differences (Strikwerda 2004). finite elements (Hughes 2012), and finite volumes (Le Veque et al. 2002). Once the discretized numerical system formulated, one often needs to solve a linear system. Solving these linear systems accurately and efficiently the key to a robust PDE solver. Previous work has studied numerical linear algebra extensively (Trefethen & Bau III 1997, Golub & Van Loan| 2013). These solvers have formed the backbone of many successful, mature numerical PDE software, e.g., ANSYS (Madenci & Guven 2013). Abaqus (Helwany 2007). COMSOL (Pryor 2009), etc. These solvers are crucial for science and engineering, but extensive wall-clock computation time limits their wider deployment. This time bottleneck typically comes from (repeatedly) solving linear system, and the problem becomes even more severe when the size of the problem scales up. To accelerate linear solvers, prior work has constructed efficient preconditioners for iterative linear solvers via matrix decomposition (Khare)
& Rajaratnam 2012), sparsity control (Liu) 1990 Davis & Hager 1999; Schäfer et al. 202 and multiscale techniques (Chen et al. 2021).
Machine learning (ML.) methods for solving PDEs Researchers have attempted to incorporate machine learning techniques into solving PDEs. A particular line of work focuses on numerical simulation problems, e.g., solving the underlying PDEs for simulating fluid (Kim et al.) 2019). cloth (Pfaff al., 2020). and elastoplastic solids (Sanchez-Gonzalez al. 2020). Typically, these approaches aim to train surrogate neural network model that replaces the PDE solving process entirely (Brandstetter et al. 2022b Such methods enjoy a major speedup over classic numerical methods thanks to the fast inference time of a trained network model. However, while the network models are capable producing visually plausible results, they yield no guarantee on the accuracy of the solution. As such, these methods often suffer from error accumulations, inhibiting their wider usage in applications where accuracy is a priority.
Combining numerical methods with networks Given the pros and cons of classic and machine learning methods, researchers have been combining the classical numerical solver with neural net- works (Belbute-Peres et al.) 2020 Um et al. 2020, Li al. 2020a b). with physics-infomme informed neural networks being a particularly notable example (Raissi et al. 2019 Karniadakis et al., (022). Since neural networks serve the backbones of these methods, they often have hard time generating quantitatively accurate solutions (c.g.. le - 12 precision) and enforcing hard constraints (Márquez- Neila et al. 2017), both of which are relatively simple with classic numerical solvers. Additionally,
2

---

Under review as a conference paper at ICLR 2023
generalization can be issue when faced with data out of the distribution in the training set Kim et al. 2021). Our method avoids these issues by construction because we use a classic solver as the backbone and fine-tune its preconditioner with neural networks. Therefore, convergence and accuracy of the solutions are guaranteed even for unseen data the performance may degrade but the solving process will not fail. AAzuay & Treister (2022) also model preconditioners with neural networks, but their convolutional neural network (CNN) architecture only supports voxel grid dis- cretization and has not been shown to obtain performance gain. Our task and method is perhaps most similar to previous works that uses CNN to learn preconditioners (Sappl et al. 2019 Ackmann al. 2020). However, we are different from these works in three folds. Becuase CNN architec- ture only supports grid-based data domains, these methods cannot be directly applied to mesh-based data, whereas we focus on mesh-based problem domains that has several advantages over grid-based data, such as sharp interface handling. Additionally, we propose novel loss function that exploits field values, but previous methods use condition number as loss energy that
the data distribution only focus on the system matrix. Finally, proposed method is more generic and works and tested on all second order linear PDEs, whereas previous methods uses hard-coded threshold to satisfy the constraint for specific problem and is not applicable other problem settings.
METHOD
3.1 PROBLEM SETUP
In this work, we focus on solving linear second-order PDE problems, which cover some of the most common PDEs in scientific and engineering applications, e.g., the heat equation in thermodynamics, the wave equation in acoustics, and Poisson's 's equation in electrostatics.
Linear second-ordes PDEs ing format:
Formally speaking, we write linear second-order PDEs in the follow-
AV)(x) + b. v/(x) c V Q.
Here, Ω C Rd is the problem domain, f Q - R is the function to be solved, A E Rdddd and E IRd are constants, and c : Ω + R is a user-specified function. Without loss of generality, we can safely assume A to be a symmetric matrix. The definition of A classifies linear second-order PDEs into elliptic (e.g., the Poisson or Laplace equation), hyperbolic (e.g.. the wave equation), and parabolic (e.g.. the heat equation) ones. This work studies the model PDEs from these categories.
Boundary conditions
We consider a mixture of Neumann and Dirichlet boundary conditions:
0f(x) On =x(x), Vx E 02N,
(2)
f(x) -D(x), Vx E 80
(3)
Here, M2y and B2D are a partition of the domain's boundary 0., and N : 0ΩN and D : 0LD + IR are two user-specified functions. The Neumann boundary condition defines the behavior of the directional gradient of f along the normal direction n at the boundary, and the Dirichlet boundary condition sets the target value of f at the boundary explicitly.
Numerical problems after discretization The PDE defined above is a continuous problem that needs to discretize before applying numerical solver. Discretization approximates gradient op- crators with numerical computation and converts the continuous PDEs into numerical systems to be solved. Popular discretization schemes include finite difference, finite element, and finite volumes. In this work, we adopt the standard Galerkin method from the finite element theory (Johnson 2012). resulting in linear system of equations:
Kf
(4)
where K E B with being the number of degrees of freedom after discretization, is the stiff- ness matrix of the PDE system after discretization, which is large, sparse, and symmetric-positee definite (SPD). The vector E IR" discretizes the function the domain and also typically fuses
3

---

Under review as a conference paper at ICLR 2023
information from (discretized) boundary conditions. The goal is to solve f E R", the support of after discretization at each degree of freedom, from which we can reconstruct the final solution f through numerical interpolation. This linear system (K, c) becomes the input to our method and other baseline algorithms discussed in this paper.
3.2 MOTIVATION
Our strategy is to build hybrid PDE solver that combines the advantages of both machine learning approaches and classic numerical solvers. Traditional numerical methods for solving Kf = fall into two main categories: direct solvers and iterative solvers. Direct solvers factorize K into matrices that are easier to solve (e.g., triangular matrices) and are most useful only if the left-hand side matrix K remains fixed. We are more interested in iterative solvers because they are more suitable for varying K and f. Iterative solvers, e.g., the srijugatt-gradienn (CG) method, repeatedly apply matrix-vectorr products to refine an estimation of the solution until convergence. The runtime of an iterative solver largely depends on the condition number of the stiffness matrix K, motivating the preconditioning techniques in numerical methods. At a high level, preconditioning a numerical system means working on a modified system similar to the original problem but faster to solve. It uses the solution of the modified system to bias the iterative solvers towards solving a system with better conditioning, c.g., by clustering the eigenvalues of the stiffness matrix (Solomon 2013).
Our method is motivated by the observation that the preconditioners described above share remark- able similarities with neural-network PDE solvers (Sanchez-Gonzalez et al. 2020): They both gen- erate inexact solutions with moderate errors fast computation time. This observation inspires us to propose the Neural-Network Prreoor-Preconditioneed Conjugate Gradient method, or NeuralPCG, which uses CG as the backbone but with a trained neural network as its preconditioner. We see this combi- nation as mutually beneficial for both neural-network methods and classic solvers. On the network side, instead of replacing the whole PDE solver with a neural network as in many previous works (Li et al. 2018 Sanzhez-Gonzalez et al. 2020 Pfaff et al., 2020), using a network only to replace the preconditioner ensures convergence, accuracy, and robustness of the solution. It also enables us to generalize the network method to problems that favor precision over speed and to unseen data.
NeuralPCG also benefits CG solvers by proposing a new perspective for designing precondition- ers. Designing a high-performance preconditioner is challenging because its two desired properties (fast derivation speed and high similarity to the original system) often conflict. Otherwise, it would imply that the stiffness matrix K itself was too easy solve in the first place. The design of clas- sic preconditioners, e.g., Incomplete Cholesky or Symmetric Successive Over-Relaxation (SSOR) (Golu Van Loan 2013). defined on the soft-hand-side side matrix K only. We argue that this design decision unnecessarily limits the full power of preconditioners because they overlook the right- side vector e and its distribution among actual PDE problem instances. contrast to these classic preconditioners, we propose to learn a netwol- preconditioner from both left-hand-side ma- trices and right-hand-side vectors in the training data. We expect a higher performance due to such exploitation of data distributions.
3.3 THE NETRAL-NETWORK PRECONDITIONER
Network design Given linear system Kf c, our network preconditioner takes as input (K, c and outputs a preconditioner P that CG, a standard iterative solver, can use seamlessly to solve the system. The representative of K depends on the discretization scheme of the PDEs. For example, a grid discretization leads to K matrix conveniently stored as a grayscale image on the grid, and a triangulated PDE results in a K stored as the node and edge values on a graph. Our work demonstrates our preconditioner on triangle and tetrahedron finite elements and uses GNNs as the network model. However, we believe our core idea is agnostic to discretization schemes.
More concretely, consider PDE problem on Ω discretized as triangle mesh in 2D or tetrahedron mesh in 3D. We construct a GNN whose nodes and edges are mesh vertices and edges, respectively.
We store the input K as a one-dimensional edge feature: If Kij is a nonzero entry in K, we add it as an edge feature in the edge from node to j. Similarly, we store vector f as a -iedennnnii node feature on the graph. We then apply the same encoder- mestrre--eegeee-ee message-passing--eeee -ecoder network chitecture as in previous work |faff et al. (2020) to predict a -inddmnnsioonal feature on each edge, which serves as the output of our network. We leave more details about our network in Appendix.
4

---

Under review as a conference paper at ICLR 2023
The last step in network preconditioner is constructing valid preconditioner from the network output. Unfortunately, directly assembling the predicted edge feature into a matrix often fails to serve a valid preconditioner because there is no guarantee of its symmetry positive definiteness. Therefore, we first construct a lower-triangular matrix L(K, c) from the network output, We use LL preconditioner, and this construction ensures its symmetry and positive definiteness.
Loss function Designing many classic preconditioners can be cast as a problem of minimizing their discrepancy the given linear system over a set of casy-to-compun matrices:
min L(P, K) PEP
(5)
where is the stiffness matrix defined above and the system that we want to precondition upon, P is the feasible set of preconditioners, and L(.,-) -) is a loss function defined on the difference between the two input matrices. A careful choice of P ensures the preconditioners remain fast to compute. For example, one compute the Jacobi preconditioner by choosing P the set of all diagonal matrices and L as the difference between the diagonals of the two matrix inputs.
Following the idea described above, it is now tempting to consider the following loss function defi- nition for our neural-network preconditioner:
min (K.f..c))
Σ
[L, CL)g (K, c2) K,
where 0 is the network parameters to be optimized, L(0) is the lrier-triangular matrix from the net- work's output, and I1: represents the squared Frobenius norm. The index loops over training data tuples (K, f,, c,). This definition closely resembles the goal of the famous Incomplete Cholesky preconditioner, especially since L shares the same sparsity pattern as the lower triangular part of K. However, closer look at the loss function reveals some potential inefficiencies in its design:
L
LL - 1
(7)
Σ I(LL" - ?
(8)
LL e, - K,e,,
where stands for the one-hot vector with one at the j-th entry and zero elsewhere. This derivation shows that this loss encourages well-rounded preconditioner with uniformly small errors in all e, directions, regardless of the actual data distribution in the training data (K, f,, c,). Therefore, we
consider new loss function instead:
L:
Σ
LL f; K.f.
(10)
Σ LL f; 11
(11)
Comparing these two losses, we can see that the new loss replaces e; with f, from the training data. Therefore, the new loss encourages the preconditioner ensemble K not uniformly in all directions but towards frequently seen directions in the training set. Essentially, this new loss trades generalization of the preconditioner with better performance for more frequent data.
Remark It is worth revisiting the competing factors in designing high petthrrerormmcce precondition- ers and clarifying how our approach handles these factors. Traditionally, designing a preconditioner strikes balance between fast computation time and similarity with the underlying system, which often conflict with each other. For example, the Jacobi preconditioner is extremely fast to com- pute. However, it only approximates the diagonal of the system. On the other hand, the Incomplete Cholesky preconditioner well approximates nonzero entries in the system but requires a much longer computation time. Our neural-netwol preconditioner resolves the conflicts between these factors by inheriting the speed from cural networks and achieves high approximation accuracy by targeting certain data distributions learned from the training data set.
5

---

Under review as a conference paper at ICLR 2023
heat-2d
wave-2d
poisson 2d
poisson-2d
cight shaped mesh
circle mush
cylinder flow mosh
connector mesh
armadillo mesh
Figure I:
Environment Overview. Left to right: heat-2d, wave-2d, poisson-2d, and poisson-3d.
3.4 THE FULL NEURALPCG SOLVER
We ready describe the full NeuralPCG solver during test time. For any (K, c) input from a PDE instance after discretization and proper boundary handling, we first send the neural network preconditioner (K, c) obtain its LL" factorization. Next, we plug the LL preconditioner into the standard CG solver and run it until the solution reaches a target precision, concluding the full NeuralPCG algorithm. Note that the neural network model only affects the beginning the CG algorithm by generating the LL factorization, after which the CG solver no longer needs to access network. This framework enables direct and fair comparison between classic CG solvers and NeuralPCG when evaluating their performances: Any performance difference can only come from the difference between preconditioners.
EXPERIMENTS
Our experiments aim to answer the following questions:
1. How does the proposed method compare with end-to-end network models in speed and accuracy?
How does the proposed method compare with classical preconditioners speed and accuracy?
3. Is the data-dependent loss introduced in the paper effective?
Does our approach generalize well unseen inputs?
We introduce the experiment setup in Sec. 4.1 followed by answering the four questions from
Sec. 4.2 to Sec. 4.5
4.1 EXPERIMENT SETUP
Environments We provide three environments that study the representative linear second-order The heat environment studies the heat equation, a parabolic PDE, the wave environment studies the wave equation, hyperbolic PDE, and the poisson environment studies the Poisson equation, an elliptic PDE. Each environment is associated with 2D triangle mesh and/or 3D tetrahedron mesh for discretization purposes (Fig. Th. More details can be found the Appendix.
PDEs:
Baselines We consider two sets of baselines: neural network and classic numerical methods. For neural network method baseline, we consider the state-of-the art MeshGraphNet (Pfaff et al.) 2020) (MGN), which takes the input (K, c) and outputs a prediction f. For classic solvers, our approach with PCG solvers using two standard preconditioners (Jacobi and IC), which we call peg-jacobi and peg-ic, respectively. We use all all-zero vector as initial guess for the PCG solver. More details about baselines can be found the Appendix.
compare
Our method The neural network architecture for preconditioner is based on the MeshGraph- Net. We make modifications of the last layer in order to generate the LL" matrix (see Section 3.3 We then pass in the learned preconditioner into the PCG solver. The only difference between our method and peg-jacobi and pcg-ic is the preconditioner. Therefore, any performance difference can be attributed the preconditioner quality. More details can be found in the Appendix.
Evaluation Metrics We consider all baselines and our method linear solvers: given a pair of K and c, all methods needs to output with the goal of satisfying Kf = c. The performance of neural network baseline MGN can be quantified by checking the residual emor Kf elj. For

---

Under review as a conference paper at ICLR 2023
Ours
MGN
Error
OLISE
9,8
21A
Figure 2: MGN Error accumulation. Field values time step (wave-2d): the ground-truth field solved using PCG with le-10 convergence threshold (top), the field predicted by MGN (middle), and their difference (bottom), all evaluated time step 5, 30, 70, 100 (left to right).
PCG baselines and our method, we quantify the performance by comparing the total wall-clock time spent for each preconditioner to reach desired accuracy levels.
To ensure a fair comparison between all methods, we summarize the performance of PCG solvers not in single number but the following values: (a) the time spent precomputing the precondi- tioner for the given K and c; (b) the iterations and (c) the total time (including the precomputing time) reach different precision thresholds (c.g.. le 12). We then compare the performance be- tween iterative solvers and end-to-end networks (e.g.. MGN) by checking the wall-clock time for the iterative solver achieve the same residual error MGN.
4.2 COMPARISONS WITH END-TO-ENT END NETWORK MODELS
In the first experiment, we demonstrate one distinct advantage of our approach over end-to-end network methods: We ensure accurate solutions while end-to-end networks accumulate errors over time. To show this quantitatively, we solve the wave-2d equation for 100 consecutive time steps using approach and MGN, Fig. shows that while our approach agrees exactly with the ground truth (le - precision), MGN deviates from the ground truth over time. At time step 100, MGN has an error of 517.4%. We can expect MGN to be faster in wall-clock time, as MGN only needs to run network inference once, while we need to run network inference to compute the preconditioner followed by running PCG solvers. Please refer Sec. A.4.I in Appendix for detail.
To summarize, MGN good at estimating solution rapidly while our approach has the flexibility of achieving arbitrary solution precision, just like standard CG solver. Therefore, network meth- ods are suitable for applications where speed dominates accuracy, while our approach is better for application that requires high precision, e.g., in scientific computing and engineering design.
4.3 COMPARISON WITH CLASSIC SOLVERS
Next, we compare our approach with numerical preconditioners, peg-jacobi and peg-ic. Table summarizes the time cost and iteration numbers of PCG solvers using different preconditioners up to convergence threshold 1e-2, le-4, 1e-6, 1e-8, le-10, and le-12.
peg-jacobi conducts the simple diagonal preconditioning so it has little precomputation overhead. However, the quality of the preconditioner is mediocre, and PCG takes many iterations to converge. peg-ic takes fewer iterations to converge. However, its precomputation process is sequential and therefore expensive to compute. By contrast, our approach features an easily parallelizable pre- computation stage (like peg-jacobi) and produces preconditioner with quality close to peg-ic. Therefore, in terms total computing time, our approach outperforms both peg-jacobi and peg-ic across wide range of precision thresholds. We also show additional experimental results in ap-

---

Under review as a conference paper at ICLR 2023
Tisk
Method
Precompute time 4
lime Giler) 4
lime (iter) 4
time (ter.) 4
Lime (iter. 4
time (iter) 1
time (ilce))
until le.2
until Ie-4
unsil 1e-6
until 1e-8
uniil Je-10
unsil 1e:12
Deg-jacohi
0.0001
0.147 (0)
1,308 (78)
2,674(169) (169)
1.577 (250)
4,650 (302)
5,605 (373)
beat-2d
pog-ic
1.331
(32)
(78)
(99)
3.143
Ours
0.0257
0.226 (0)
0.791 (38)
1,532 (88)
2,02 (122)
2.469 (152)
3,073 (193)
peg-jacohi
0:141 (0)
0:141 (0)
D:141 (0)
0.176 (6)
(26)
0.417 (46)
wave-2d
pog-ic
0.7679
0,885 (0)
0,885 (0)
0.885 (0)
0,904 (3)
0.953 (11)
1,007 (20)
0.0147
0,081
0.081
0,081 (0)
0.100
(12)
0.211 (21)
Deg-jacohi
0.0001
0,980 (275)
1,231 (348)
1,572 (448)
1,822 (522)
2,119 (611)
2,405 (097)
possion-2d
pog-ic
1.188 (1335
1.309
1.668 (219)
246)
1.774 (311)
(349)
ourss
0.0145
0.6339 (175)
0,818 (227)
1.017 (286)
1.118 (3166
1,312 (374)
1,510 (452)
peg-jacohi peg-ic
0.0002
1,526 (0)
2. 2,693 (7)
5,496 (25)
9.552 (50)
13:636 (76)
17.080 (97)
possien 3d
3,7371
6.441 (0)
6.891 (3)
7.905 (9)
9.436 (18)
11.005 (27)
12.346 (35)
0.4137
3,010
3.320
4.815 (13)
6,908 (28)
8,749 (41)
10.406 (53)
Table 1: Comparison between preconditioners for PCG. We report precompute time, total time (peg- icl. precompute time) for each precision level, and the corresponding PCG iterations (in parenthesis). The best value in each category is in bold. 4: lower the better.
pendix condition number comparison between our method and previous methods in Appendix
4.4 ABLATION ON THE LOSS FUNCTION
To highlight the value of our loss function targeting data distributions on the training set, we train the network with the losses in Eqn 66 and in Eqn (10). respectively. Compared to Eqn Eqn (TO)
exposes the data distribution of to the training process.
Mothod
Procompatc time J.
time (itee)
time (iar)) 4
time (iter.
time (iter.) 4
time (iter.)
time (itor.)
1e-2
until
1e-6
until
until 1c-10
until Te-12
Naive loss (Egn
0.0218
0.177 (15)
0.732 (181)
1.187 (236)
1.518 (296)
1.885(296))
2.284 (361)
Our loss (Egg
0.0271
0.172 (14)
0.420 (57)
0.597 (87)
0.733 (110)
0.909 (140)
1.056 (165)
Table 2: Wall-clock time and iterations: our method with two different loss functions on heat-2d.
Table 2 reports the performances of our loss function Eqn T0p and the naive loss function Eqn (6) on the heat-2d environment. We observe that the preconditioner trained with Eqn (10) converges in fewer iterations than Eqn (6). As such, we conclude that enforcing data distribution dependence during training allows us to achieve better an-distribution inference during test time.
4.5 GENERALIZATION
Physics parameters. First, we consider generalizing the PDEs on their physics parameters, which govern system K. We use poisson-2d as an example. Our method trained on fixed density distribution, and we test the performance of our method on distributions that gradually deviates from the training distribution.
the performance of our approach on these parameters.
Since changing physics parameters does not affect the matrix sparsity, the precomputation time remains largely unchanged across different data distributions (see Table 3 Column 1). Even on the challenging out-of-distribution datasets, our approach still maintains reasonable performance, achieving high precision while using less total time than peg-ic and peg-jacobi.
Geometry. We also conduct generalization tests mesh models, i.e., training network models on one mesh and testing its performance on unseen meshes. We demonstrate this on the poisson-3d environment, where we train the network preconditioner test the trained network model on a new mesh "armadillo" (Fig.B)
"connector mesh" (Fig. II. We then
By comparing the bottom and middle rows in Fig. (MGN) struggles to generate accurate solutions when deployed on the unseen mesh (0.5315 error), whereas our approach achieves arbitrary accuracy by construction (le 9 error threshold here).
first see that the end-to-end network method
We also report the time and iteration cost of our approach and classic preconditioners on poisson-3d with armadillo mesh. (Table 4h. We notice that our method can generalize unseen mesh. Even
8

---

Under review as a conference paper at ICLR 2023
Tisk
Method
Procompute time 4
lime (iler.) 4
time (iir.)
lime (ier.)
time (ier.) 4
time (iter) 4
time (iter. 4
te-4
until Te-6
until
until
until 1o-12
peg peg-jacobi
0.0001
0.811 (230)
0.983 (281)
1.341 (388)
1,699 (174)
1832 (555)
2,176 (640)
lest
peg-ic
0.6955
1,079 (100)
1.157 (123)
1.275 (196)
1,437 (203)
1,525 (229)
1.648 (264)
0,0138
0,513 (1337)
0.615 (167)
0.814 (226)
0.987 (277)
1,109 (313)
1,297 (369)
peg-jacobi
0.0001
0.904 (253)
1.1 (310)
1,505 (130)
1,756 (505)
2,011 (589)
2,414 (702)
peg-ic
0.7093
1.104 (110)
1.178 (132)
1.338 (180)
1,461 (217)
1.551 (244)
1,715 (293)
ourss
0.0139
0,603 (165)
0.72 (200)
0.986 (250)
1.14 (326)
1,292 (372)
1.557 (152)
peg jacobi
(262)
(317)
(519)
2.123 (606)
tost 3a
ic
0.6992
1.114 (111)
1.198 (137)
1,365 (157)
1.484 (222)
1,578 (251)
1,743 (301)
0.0141
0.56 (153)
0.672
0.905 (256)
1,051
1,201 (345)
1.433 (415)
pag-jacobi
0.0001
1.082 (310)
1,455 (421)
1.761 (512)
2,06 (603)
2.531 (745)
2.763 (816)
peg-ic
1.188 (1355
1,311 (171)
(219)
1,564 (346)
1,783 (311)
(349)
Ours
0.0136
0.728 (203)
0.97 (275)
1,165 (333)
1.334 (384)
1,683 (488)
1.817 (528)
Table 3: Generalization: physpeg-ics parameters. We test the PCG approaches on testing datasets increasingly deviating from the training data distribution. stands for standard deviation of training set, test-o, test-3a, and test-5a means 1, 3, and std. dev away from training distribution, respectively.
Training Example
MGN (error = 0,5315)
Ours (error te
poisson-3d on unseen mesh: MGN (middle), our method (bottom), training mesh (top). Left to right: solution fields at different cross-section heights.
Figure 3:
in the challenging case of solving PDE unseen test mesh, still able to converge in less time than peg-ic and peg-jacobi across different precision requirements.
5 CONCLUSIONS, LIMITATIONS, AND FUTURE WORK
This work presented NeuralPCG, a hybrid ML-numeric PDE solver. While prior ML PDE solvers often do not have precision guarantees, approach attains arbitrary precision (up to floating-point error) by construction. Our key observation that the preconditioner for classic iterative solvers does require exact precision and is. therefore. ideal candidate for neural network approxi- mation. NeuralPCG approximates the preconditioner with graph neural network and embeds this preconditioner into a classic iterative conjugate gradient solver. Compared to end-to-end ML ap- proaches, our approach takes a longer wall-clock time to compute but reduces the error level from
Method
Precompute
Time until 1c-2
Time (iter.
Time J
(iter.)
until Ic-4
until 1c-6
until 1c-8
until Ic-10
until lc-12
peg-jacobi
0.0002
2.314 (6)
4.634 (22)
8.395 (48)
12,002 (72)
15.381 (95)
18.332 (115)
peg-ic
(13)
(23)
(31)
ours
0.4206
3,591 (4)
4.97 (14)
7.005 (29)
8.978 (43)
10.721 (55)
12.412 (68)
Table 4: Our approach generalizes to unseen armadillo mesh in the poisson-3d environment, and outperforms both peg-ic and peg-jacobi in total wall-clock time.

---

Under review as a conference paper at ICLR 2023
1c - to le- 12. Compared to classic preconditioners, approach is faster while achieving the same accuracy.
Currently, approach is limited to linear PDEs. Future work may consider extending to more complex PDEs, such as the elastodynamics equations and the Navier-Stokes equations shown in prior end-to-end ML approaches SSanalez-Gonzalez et al. 2020). Additionally,
enforce the pre- conditioner's sparsity to be the lower triangular sparsity of A. We consider relaxing this constraint and exploring more effective sparsity control (Schäfer et 2021 as an exciting future research direction.
REFERENCES
Jan Ackmann, Peter D. Düben, Tim N. Palmer, and Piotr K. Smolarkiewicz. Machine-learned preconditioners for linear solvers in geophysical fluid flows, 2020. URL https://arxxiv
org/abs/200 2010. 02866
Owe Axelsson. Iterative solution methods. Cambridge university press, 1996.
Yacl Azulay and Eran Treister. Multigrid-augmennee deep learning preconditioners for the helmholtz equation. arXiv preprint arXiv:2203.11023. 2022.
Filipe De Avila Pabbute-Peres, Thomas Economon, and Zico Kolter. Combining differentiable pde solvers and graph neural networks for fluid flow prediction. In international conference on ma- chine learning, pp. 2402-24222 2411. PMLR, 2020.
Johannes Brandstetter, Max Welling, and Daniel E Worrall. Lie point symmetry data augmentation for neural pde solvers. arXiv preprint arXiv:2202.07643, 2022a.
Johannes Brandstetter, Daniel Worrall, and Max Welling. Message passing neural pdc solvers. arXiv preprint arXiv:2202.03376 2022b.
William L Briggs, Van Emden Henson, and Steve F McCormick. multigrid tutorial. SIAM, 2000.
Jiong Chen, Florian Schäfer, Jin Huang, and Mathieu Desbrun. Multiscale cholesky preconditioning for ill-conditioned problems. ACM Transactions on Graphics (TOG), 40(4):1-13, 2021.
Timothy Davis and William W Hager. Modifying sparse cholesky factorization. SIAM Journal on Matrix Analysis and Applications, 20(3):606-621) 1999.
Gene H Golub and Charles F Van Loan. Matrix computations. JHU press, 2013.
Sam Helwany. Applied soil mechanics with ABAQUS applications. John Wiley & Sons, 2007.
Philip Holmes, John L Lumley, Gahl Berkooz, and Clarence W Rowley. Turbulence, coherent structures, dynamical systems and symmetry. Cambridge university press, 2012.
Thomas JR Hughes. The finite element method: linear static and dynamic finite element analysis. Courier Corporation, 2012.
Claes Johnson.
Numerical solution of partial differential equations by the finite element method. Courier Corporation, 2012.
George Em Karniadakis, Ioannis G Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang. Physics-informed machine learning. Nature Reviews Physics, ((6):422-440, 2021.
Kshitij Khare and Bala Rajaratnam. Sparse matrix decompositions and graph characterizations. Linear algebra and its applications, 437(3)9932 9479992229999909 2012.
Byungsoo Kim. Vinicius C Azevedo, Nils Thuerey, Theodore Kim, Markus Gross, and Barbara Solenthaler. Deep fluids: A generative network for parameterized fluid simulations. In Computer graphics forum, volume 38, pp. 59-70. Wiley Online Library, 2019.
Jungeun Kim, Kookjin Lee, Dongeun Lee, Sheo Yon Jhin, and Noseong Park. Dpm: novel train- ing method for physics-infiee informed neural networks in extrapolation. In Proceedings of the AAAI Conference Artificial Intelligence, volume 35, pp. 8156-8154, 2021.
10

---

Under review as a conference paper at ICLR 2023
Randall J Le Veque al. Finite volume methods for hyperbolic problems, volume 31. Cambridge university press, 2002.
Yunzhu Li, Jiajun Wu, Russ Tedrake, Joshua B Tenenbaum, and Antonio Torralba. Learning par- ticle dynamics for manipulating rigid bodies, deformable objects, and fluids. arXiv preprint orXivv1810.014 2018.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint arXiv:2010.08895 2020a.
Zongyi Li, Nikola Kovachki, Kamyar Arizzadenesheli, Burigede Liu, Andrew Stuart, Kaushik Bhat- tacharya, and Anima Anandkumar. Multipole graph neural operator for parametric partial differ- ential equations. Advances in Neural Information Processing Systems, 33:6755-6766 6766, 2020b.
Joseph WH Liu. The role of elimination trees in sparse factorization. SIAM journal on matrix analysis and applications, 11(1)1344-12, 1990.
Erdogan Madenci and Ibrahim Guven. The finite element method and applications in engineering using ANSYS@. Springer, 2015.
Pablo Márquez-Neine Neila, Mathieu Salzmann, and Pascal Fua. Imposing hard constraints on deep net- works: Promises and limitations. arXiv preprint urXiv:176..088 2017.
Tobias Pfaff, Meire Fortunato, Alvaro Banchez-Gonzalee and Peter Battaglia. Learning mesh-based simulation with graph networks. In International Conference on Learning Representations, 2020.
Roger W Pryor.
Multiphysics modeling using COMSOL@: first principles approach. Jones & Bartlett Publishers, 2009.
Maziar Raissi, Paris Perdikaris, and George Karniadakis. Physics-informe neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.
Alvaro Sanchez-Gonzales Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter Battaglia. Learning simulate complex physics with graph networks. In International Confer- ence on Machine Learning, pp. 8459-8468. PMLR, 2020.
Johannes Sappl, Laurent Seiler, Matthias Harders, and Wolfgang Rauch. Deep learning of precon- ditioners for conjugate gradient solvers urban water related problems, 2019. URL https:
7/ //axxiv.rrg/aa 1906999661100 06925
Florian Schäfer, Matthias Katzfuss, and Houman Owhadi. Sparse cholesky factorization by kullback-leibler minimization. SIAM Journal on scientific computing, 43019:420,,, 2021.
Barry Smith. Domain decomposition methods for partial differential equations. In Parallel Nu- merical Algorithms, pp. 225 243-243.2F Springer, 1997.
Justin Solomon. Numerical algorithms: methods for computer vision, machine learning, and graph- ics. CRC press, 2015.
John Strikwerda. Finite difference schemes and partial differential equations. SIAM, 2004.
Lloyd N Trefethen and David Bau III. Numerical linear algebra, volume 50. Siam, 1997.
Kiwon Um, Robert Brand, Yun Raymond Fei, Philipp Holl, and Nils Thuerey. Solver-in-the--ue Learning from differentiable physics to interact with iterative pde-solvers. Advances in Neural Information Processing Systems, 33:611 1166111612222 2020.
11

---

Under review as a conference paper at ICLR 2023
A APPENDIX
A.1
ENVIRONMENT DETAILS
Table lists the environment details for experiment section 4.2 section
section
section
PDE
Mesh Shape
Number of Nodes
Number of Elements
Physies Paramoter
Boundary Coadition
heat-2d
eight shaped mesh
7454
14351
diffusivity: 10.0
varying varying varying fixed fixed
wave-2d
circle shaped mesh
4852
8839
speed: 0.1
poisson 2d
cynlinder flow
3167
6117
density: 0.01
polssson-33
connector mesh
23300
129981
polsson-3d
armadillo mesh
1818]
97476
Table 5: Environment setup for experiment section 4.2 section 4.3 section 4.4
section 4.5
Figure shows examples demonstrating varying boundary conditions. For heat-2d and wave-2d, varying length and position mesh geometric boundary nodes are selected as Dirichlet boundaries. the inviscid-euler fluid equation are only responsible for solving the pressure that makes the velocity field incompressible, which is Poisson equation. The advection and external force steps then applied generate the data visualization. For poisson-2d, two sets of varying length and position of mesh geometric outer border boundary nodes selected influx and Dirichlet boundary. The remaining mesh geometric border nodes, including the remaining outer border and all nodes in the inner border, are obstacle boundary.
For poisson-2d equation,
demonstration. All solvers
Hest
Wave 2D
Poisson 2D
Figure 4: Varying boundary conditions: For heat-2d and wave-2d, the black vertices represent the Dirichlet boundary. The poisson-2d example shows the red vertices are obstacle boundary, the blue vertices are influx boundary, and green vertices are the Dirichlet boundary.
For experiment 4.5 across different physics parameters consider the same mesh domain used in other poisson-2d experiments with mesh size of 3167 and element size 6117. We train on a set of training set with density distribution from [0.001, 0.005], and our environment test-la of density 0.006. Test environment test-3o is of density 0.008. Test environment test-5a is of density 0.01.
A.2 IMPLEMENTATION DETAILS AND HYPERPARAMETER
In this Section, we describe the implementation details of each method.
A.2.1 BASELINES
MeshGraphNet To ensure a fair comparison, we re-implement MeshGraphNet in PyTorch, the framework we used for our method, based on their paper and released code. We have also manually tuned its hyperparameters on our dataset to ensure we see reasonably good performance from Mesh- GraphNet. We follow the are all Multi-Layer Perceptrons (MLP) with ReLU activation function. Specifically, 2 layer MLP encoder and a layer MLP decoder are used for the heat-2d, wave-2d, and poisson-2d equation. The Processor Module is composed of 5 iterations of message passing layer. For the poisson-3d equation, we use a layer encoder, 2 layer decoder, and iterations of message passing for the processor module. All MI.Ps used are of hidden dimension size of 16.
Decoder. Encoder, Processor and Decoder modules
12

---

Under review as a conference paper at ICLR 2023
Classical Methods
PCG-IC is implemented using pytorch and is fully vectorized for optimized computation. This implementation follows Golub & Van Loan (2013).
PCG-Jacobi is implemented using pytorch and fully vectorized. The Jacobi or diagonal precondi- tioner is described in Axelsson 1996).
A.2.2 OUR METHOD
Technical Details and Justification We follow the diagonal decomposition LDL as way of lower triangular decomposition for the original system K. It easy to see that this diagonal decom- position is equivalent to lower triangular decomposition.
K = LoLg
(12)
= Lg DVDDL.
(13)
DLog
(14)
The diagonal decomposition LDL has several advantages, similar lower triangular decompo- sition LL T it is easy to invert, and guarantees symmetry. Additionally, we enforce the diagonal element D be the diagonal elements of original system K. This way, we enforce the value and gradient range for the lower triangular matrix Lg to ensure positive definiteness of the learned de- composition.
Implementation Details We follow the same Encoder- Pwocessor--eeooeeee architecture implemen- tation for MeshGraphNet. A layer MLP encoder and a layer MLP decoder are used for the heat-2d, poisson-2d, and poisson-3d equation. The Processor Module is composed of 5 iterations of message passing layer. For the Poisson 3D equation, we use layer encoder, 2 layer De- coder, and iterations of message passing for the processor module. All MLPs used are of hidden dimension size of 16.
A.3 TRAINING SPECIFICATION
We use workstation equipped with 64-core AMD 3995W CPU and NVIDIA RTX A6000 GPU for all our experiments.
We adopt end-to-end training scheme with the loss functions described section We train our method for the heat-2d equation environment for 6 hours. The wave-2d equation is trained for 6 hours. The poisson-2d equation is trained for 6 hours. The poisson-3d Equation is trained for hours.
We train the NN baselinebaseline for days, 72 hours, for experiments wave-2d, poisson-2d, and ad.
A.4 ADDITIONAL EXPERIMENTAL RESULTS
A.4.1 TIME COMPARISON WITH END-TO-ENS TO-END NEURAL NETWORK
We first report Table the time cost of network methods and that of our method achieving the same residual error. Specifically, for each environment, we first run MGN on the test set. We then record the wall-clock time and the residual error from each pair of (K, c). Next, we use the residual error as the convergence threshold in NeuralPCG and record the wall-clock time cost. It is unsurprising from the table to see that MGN is substantially faster than our approach: For a given pair of and c, MGN only need to network inference once, while we need to run network inference to compute the preconditioner followed by running PCG solvers. We notice that our method usually converge the MGN residual error in iteration.
A.4.2 CONDITION NUMER
We show the condition number comparison in Table
13

---

Under review as a conference paper at ICLR 2023
епv.
heat-2d
wavc-2d
poisson-24
poisson-3d
time (MGN) (s)
0.0213
0.0171 0.0925
0.0135 0.0527
0.1154
time (ours) (s)
0.0950
1.7762
Table 6: The wall-clock time of MGN and our method stopped at the same error produced by MGN.
Method
Wave-2d
Poisson-2d
Poisson-3d Heat-2d
P (original system)
540272.25
43658.16
181.56
1008.79
peg-jacobi
96.35
18712.16
165.71
225.86
peg-ic
23.22
5662.48
42.83
43.23
Ours 23.89
8384.31
64.23
136.86
Table Condition number comparison between various methods
A.4.3 ADDITIONAL COMPARISON WITH CLASSICAL PRECONDITIONERS
Here we show additional experimental results on the four equations, heat-2d, wave-2d, poisson-2d, and poisson-3d, under other physical parameters or other meshes.
Task
Method
Procompute time
time (iter.) uncil
time Gler.)
time (iler.)
time (iter.)
lime (iker.)
time (iter.)
parameter
(5)
until
until
until Ic-10
le-12
heal-21)
Lacobi
0.0002
0.240
0.766 (109)
1,131 (167)
1,4066
1,784 (273)
2,080 (321)
circle mnchh
IC
1.0105
1.156 (12)
1.346 (44)
1,495 (69)
1.611 (88)
1,751 (112)
1,873 (132)
diffusiiit==10.00
Ours
0.0271
0.172 (14)
0.420 (57)
0.597 (87)
0.733 (110)
0.909 (140)
1,056 (165)
heat-21)
Jacob
0.0011
0.669 (32)
1,755 (132)
2,098 (202)
3,471 (257)
4.317 (333)
5.174 (.174(398)
eight shaped mush
IC
1,5453
1,931 (13)
2.369 (44)
2.757 (71)
3,063 (93)
3,346 (112)
3.696 (137)
diffusivit==1.5
Ours
0.49
(71)
(110)
(140)
(177)
(214)
wave-21)
Jacab
0.0188
0.072 (0)
0,072 (0)
0.072
0.105
0.214 (24)
0.343 (46)
circle mush
IC
0.7088
0.781 (0)
0.781 (0)
0.781 (0)
0.794 (2)
0.843 (10)
0,899 (20)
specd=0.3
Ours
0.0171
0,091 (0)
0.091 (0)
0.091 (0)
0.104 (2)
0,156 (11)
0.214 (21)
poisson-21)
Jacob
0.0011
1.339 (219)
1,643 (271)
2.211 (369)
2.670 (449)
3.071 (518)
3,653 (620)
circle mesh
IC
0.9052
1,505 (88)
1.632 (110)
1881 (141)
2,041 (178)
2.172 (200)
2,403 (238)
density-0.01
Ours
0.0143
(125)
1.01 (152)
1,316 (202)
(243)
(277)
(325)
poisson-DD
Jacob
0.0011
(3348)
2,119 (611)
2,405
syalinder flow
IC
0.6905
1.140 (115)
1.213 (135)
1.385 (183)
1.504 (215)
1.614 (246)
1.746 (282)
density-0.0008
Ours
0.0145
0.639 (175)
0,818 (227)
1,017 (286)
1,118 (316)
1,312 (374)
1,51 (452)
possion-3D urmadillo mesh
Jacobi IC
0.0001
1.749 (8)
1,365 (38)
6.642 (63)
8.955 (89)
11,363 (115)
13,832 (143)
3.3779
5.175
6.038 (13)
6.869 (23)
7,603 (31)
8.336 (40)
9.237 (50)
Ours
0.1142
2.273
3.726 (25)
4.919 (37)
(52)
7.325 (65)
8.677 (81)
Table 8: Additional experimental results comparing with classical preconditioners.
A.4.4 ADDITIONAL COMPARISON WITH LEARNING-BANE BASED PRECONDITIONER
We show proposed method compare with learning based preconditioner Sappl al. (2019) in Table Comparing to previous Bearning-base preconditioner, NN baseline, our method shows stable convergence across various settings using different second order linear PDEs. NN base- line (Sappl et al. 2019) is not orignally designed to be generic method to work in different problem domains. For example, the hard-coded diagonal clipping threshold of le 3 is carefully picked by the authors for the urban water problem and there is no general principles to picking such thresh- old for other problem settings. We can see that NN baseline Sappl et al. (2019) works on par with our method on the wave-2d setting, but it does not work in many other settings, heat-2d, poisson- 2d. For Poisson-2d setting the resulting convergence iteration even worse than the peg-jacobi method. Additionally, the method Using condition number as the loss energy is not computation- ally efficient. computing condition number for a system requires full eigen decomposition, which is On")). Computation scales cubicly with problem size. Our poisson-3d setting uses mesh of size 23,300 nodes. We train the NN baselinemethod for three days, 72 hours, and it only make through less than percent of the training set, using the same hardware setup. Empirically, we found that using condition number as loss energy is computationally infeasible with problems more than 10,000 nodes.
14

---

Under review as a conference paper at ICLR 2023
Task
Method
Precompute time +
time (iter.)
time (iter.) +
time (iter.) 4
until le-2
until le-6
until le-10
NN bascline
0.0236
0.226 (0)
1.973 (116)
3,223 (200)
heat-2d
Ours
0.0237
0.226 (0)
1.532 (88)
2.469 (152)
NN bascline
0.0165
0.087 (0)
0.087 (0)
0.165 (12)
wave-2d
0.0147
0.081 (0)
0.081 (0)
0.156 (12)
NN baseline
0.0129
1.014 (291)
1.602 (471)
2.221 (662)
possion-2d
0.0145
0.639 (175)
1.017 (286)
1.312 (374)
Table 9: Comparison between preconditioners for PCG. We report precompute time, total time (peg- icl. precompute time) for each precision level, and the corresponding PCG iterations (in parenthesis). The best value in each category is in bold. J: lower the better.
15


Reviews:
[{"review_id": "kVfw3RcrYXp", "reviewer": "Reviewer_G2sH", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes NeuralPCG, which learns a preconditioner using neural networks to improve the solution procedure of \nthe linear solver for the resulting system of linear equations from discretizing partial differential equations (PDEs). The authors\ntested the learned preconditioner for two- and three-dimensional simple PDEs.", "summary_of_the_review": "The contribution of this paper is very limited. The paper ignores the entire existing research on learning preconditioners for solving systems of linear equations. Indeed, many existing works on deep learning of preconditioners are more solid in algorithm and validation than NeuralPCG. The experimental results are not convincing either.", "strength_and_weaknesses": "Strength: \n\nThe paper is easy to read and it is straightforward to catch the basic idea of the paper.\n\nWeaknesses:\n\n1. Learning preconditioners using neural networks for improving solving systems of linear equations has been widely studied. However,\n*the entire existing works are completely ignored by the authors*. I actually doubt if this paper has any novelty over existing works. By typing \"Learning preconditioner\" on the google search engine, we immediately get many related papers, e.g., the first two papers are: https://arxiv.org/abs/1906.06925 and https://arxiv.org/abs/2010.02866. Indeed, the study in the first paper is much more solid than the paper under review. Look at figure 5 of the the paper https://arxiv.org/abs/1906.06925, you will see the existing algebraic multi-grid (AMG) preconditioner is much better than the learned ones, Jacobi. and etc. In the paper under review, the authors never contrast the learned preconditioner with the existing state-of-the-art.\n\n\n2. The authors claimed that \"Classic numerical solvers provide unparalleled accuracy but often require extensive computation time.  Machine learning solvers are significantly faster but lack convergence and accuracy guarantees. We present Neural-Network- Preconditioned Conjugate Gradient, or NeuralPCG, a novel linear second-order PDE solver that combines the benefits of classic iterative solvers and machine learning approaches. Our key observation is that both neural-network PDE solvers and classic preconditioners excel at obtaining fast but inexact solutions. NeuralPCG proposes to use neural network models to precondition PDE systems in classic iterative solvers.\" --- I wonder if there is any convergence and accuracy guarantee of PDE solvers using the learned preconditioner? If not, then the bottleneck of existing learning-based PDE solvers is also a bottleneck of the proposed approach.'\n\n\n3. The authors stated that \"While this line of methods typically outperforms classical solvers in speed by a large margin, it struggles with converging into a highly precise solution (e.g., 1e − 12). A lack of theoretical analysis on  convergence and accuracy inhibits neural-network PDE solvers’ applications in mechanical engineering, structure  analysis, and aerodynamics, where precise PDE solutions have a higher priority than fast yet inexact results.\" --- My concern is whether we really need 1e-12 accuracy in real applications. Also, if we need such an accurate solution, do learning-based PDE solvers really faster than classical PDE solvers?\n\n\n4. The authors stated that \"NeuralPCG generalizes to physics parameters or meshes with moderate differences from the training data and is robust to outliers by design.\" --- Does the learned preconditioner generalizes to PDEs of different types and problems at different scales? AMG can be applied to problems of different sizes and even different types. I would like the authors to make some comparisons of their NeuralPCG with AMG and other state-of-the-art preconditioners.\n\n\n5. Does NeuralPCG works well when to the linear PDE when A is x-dependent?\n\n\n6. I wonder if NeuralPCG works for solving high-frequency wave equations?\n\n\n7. The authors may compare the method with some fast methods for solving PDEs, e.g., fast Fourier transformer and fast multipole method. Rather than consider the method that is applicable to any type of PDEs, which is a highly unfair comparison as the learned \npreconditioner is problem-dependent.\n\n\n8. Does the learned preconditioner reduce the condition number? I want to see some numerical evidence of reducing \ncondition number of the underlying problems.\n\n\n9. Does over-smoothing of graph neural network affect the performance of the learned preconditioner? The computational cost of \ngraph neural networks can be non-negligible when the system of linear equations is massive.\n\n\n10. In equation (6), the learned matrix L can be dense, significantly reducing the scalability of the PDE solver. This can be\na significant bottleneck of NeuralPCG. \n\n\n11. I think the authors should report the computational time, including the time for learning the preconditioner. In particular, for large-scale problems. It seems to me that NeuralPCG is much more expensive than existing SOTA algorithms. E.g. algebraic multi-grid.\n\n\n12. The baseline preconditioning methods are very weak. In practice, we use much more efficient ones.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "Bkp2KN1FN_", "reviewer": "Reviewer_DNx8", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors presented NeuralPCG, a hybrid ML-numeric PDE solver. While prior ML PDE solvers\noften do not have precision guarantees, their approach attains arbitrary precision (up to floating-point\nerror) by construction.", "summary_of_the_review": "The significance and novelty of this paper are limited.", "strength_and_weaknesses": "Strength:  The authors propose a novel and efficient NeuralPCG  solver for linear second-order PDEs and present a framework for studying the performance of preconditioners in classic iterative solvers.\n\nWeaknesses: The paper is well written and easy to follow with clear comparison to the literature.  The significance and novelty of this paper are limited.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "3eeLUo7xgV", "reviewer": "Reviewer_94zp", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a so-called Neural-Network-Preconditioned Conjugate Gradient (NeuralPCG) method to solve partial differential equations. This method uses graph neural network to model the preconditioner of iterative solvers for PDE. It is noted that the NeuralPCG can solve PDE with a satisfying convergence condition and works faster than classic PDE solver. Several classic examples, including passion equation and heat conducted equation, are solved via the NeuralPCG, the comparing results with classic numerical PCG method show the generality and efficiency of the proposed method.", "summary_of_the_review": "This is a well-written and easy to read paper. However, it did not review enough on the investigations of deep learning based precondition method. Some comments are given as follows:\n\nComment 1. Maybe the contour plot of Fig.1 and Fig.3 can have a bar of labeled field.\n\nComment 2: This paper shows the performance of NeuralPCG method based on the Conjungate Gradient method, can this paper give some discussions on the performance of general preconditioning method based on deep learning method?\n\nComment3: This paper only shows a little reviews on the investigations of deep learning based preconditioner method and does not clarify the distinguished difference of this paper with other similar paper.\n\nComment4: The accuracy of classic finite element method often shows dependence of the numbers of mesh, does this method also show this dependency?", "strength_and_weaknesses": "Strength: \nThe NeuralPCG combine the advantage of classic iterative solver and the machine learning based solver. It can have a good accurate solutions compared with neural network based PDE solver and work faster than classic solver.\n\nWeaknesses: \nThe approach is now limited to linear PDEs with symmetric-positive-definite matrix, and the preconditioner’s sparsity is constrained to be the lower triangular sparsity of A.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject