Multipole Graph Neural Operator for Parametric Partial Differential Equations
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, Anima Anandkumar
Introduction
A wide class of important scientific applications involve numerical approximation of parametric PDEs. There has been immense research efforts in formulating and solving the governing PDEs for a variety of physical and biological phenomena ranging from the quantum to the cosmic scale. While this endeavor has been successful in producing solutions to real-life problems, major challenges remain. Solving complex PDE systems such as those arising in climate modeling, turbulent flow of fluids, and plastic deformation of solid materials requires considerable time, computational resources, and domain expertise. Producing accurate, efficient, and automated data-driven approximation schemes has the potential to significantly accelerate the rate of innovation in these fields. Machine learning based methods enable this since they are much faster to evaluate and require only observational data to train, in stark contrast to traditional Galerkin methods and classical reduced order models.
While deep learning approaches such as convolutional neural networks can be fast and powerful, they are usually restricted to a specific format or discretization. On the other hand, many problems can be naturally formulated on graphs. An emerging class of neural network architectures designed to operate on graph-structured data, Graph neural networks (GNNs), have gained popularity in this area. GNNs have seen numerous applications on tasks in imaging, natural language modeling, and the simulation of physical systems . In the latter case, graphs are typically used to model particles systems (the nodes) and the their interactions (the edges). Recently, GNNs have been directly used to learn solutions to PDEs by constructing graphs on the physical domain , and it is was further shown that GNNs can learn mesh-invariant solution operators . Since GNNs offer great flexibility in accurately representing solutions on any unstructured mesh, finding efficient algorithms is an important open problem.
The computational complexity of GNNs depends on the sparsity structure of the underlying graph, scaling with the number of edges which may grow quadratically with the number of nodes in fully connected regions . Therefore, to make computations feasible, GNNs make approximations using nearest neighbor connection graphs which ignore long-range correlations. Such approximations are not suitable in the context of approximating solution operators of parametric PDEs since they will not generalize under refinement of the discretization, as we demonstrate in Section 4. However, using fully connected graphs quickly becomes computationally infeasible. Indeed evaluation of the kernel matrices outlined in Section 2.1 is only possible for coarse discretizations due to both memory and computational constraints. Throughout this work, we aim to develop approximation techniques that help alleviate this issue.
To efficiently capture long-range interaction, multi-scale methods such as the classical fast multipole methods (FMM) have been developed . Based on the insight that long-range interaction are smooth, FMM decomposes the kernel matrix into different ranges and hierarchically imposes low-rank structures to the long-range components (hierarchical matrices). This decomposition can be viewed as a specific form of the multi-resolution matrix factorization of the kernel . However, the classical FMM requires nested grids as well as the explicit form of the PDEs. We generalize this idea to arbitrary graphs in the data-driven setting, so that the corresponding graph neural networks can learn discretization-invariant solution operators.
Inspired by the fast multipole method (FMM), we propose a novel hierarchical, and multi-scale graph structure which, when deployed with GNNs, captures global properties of the PDE solution operator with a linear time-complexity . As shown in Figure 1, starting with a nearest neighbor graph, instead of directly adding edges to connect every pair of nodes, we add inducing points which help facilitate long-range communication. The inducing points may be thought of as forming a new subgraph which models long-range correlations. By adding a small amount of inducing nodes to the original graph, we make computation more efficient. Repeating this process yields a hierarchy of new subgraphs, modeling correlations at different length scales.
We show that message passing through the inducing points is equivalent to imposing a low-rank structure on the corresponding kernel matrix, and recursively adding inducing points leads to multi-resolution matrix factorization of the kernel . We propose the graph V-cycle algorithm (figure 1) inspired by FMM, so that message passing through the V-cycle directly computes the multi-resolution matrix factorization. We show that the computational complexity of our construction is linear in the number of nodes, achieving the desired efficiency, and we demonstrate the linear complexity and competitive performance through experiments on Darcy flow , a linear second-order elliptic equation, and Burgers’ equation, which considered a stepping stone to Naiver-Stokes, is nonlinear, long-range correlated and more challenging. Our primary contributions are listed below.
We develop the multipole graph kernel neural network (MGKN) that can capture long-range correlations in graph-based data with a linear time complexity in the nodes of the graph.
We unify GNNs with multi-resolution matrix factorization through the V-cycle algorithm.
We verify, analytically and numerically, the linear time complexity of the proposed methodology.
We demonstrate numerically our method’s ability to capture global information by learning mesh-invariant solution operators to the Darcy flow and Burgers’ equations.
Operator Learning
We consider the problem of learning a mapping between two infinite-dimensional function spaces. Let and be separable Banach spaces and be the target mapping. Suppose we have observation of pairs of functions where is i.i.d. sampled from a measure supported on and , potentially with noise. The goal is to find an finite dimension approximation , parametrized by , such that
An important example of the preceding problem is learning the solution operator of a parametric PDE. Consider a differential operator depending on a parameter and the general PDE
Learning the Operator.
Learning the operator is typically much more challenging than finding the solution of a PDE for a single instance with the parameter . Most existing methods, ranging from classical finite elements and finite differences to modern deep learning approaches such as physics-informed neural networks (PINNs) aim at the latter task. And therefore they are computationally expensive when multiple evaluation is needed. This makes them impractical for applications such as inverse problem where we need to find the solutions for many different instances of the parameter. On the other hand, operator-based approaches directly approximate the operator and are therefore much cheaper and faster, offering tremendous computational savings when compared to traditional solvers.
PDE solvers: solve one instance of the PDE at a time; require the explicit form of the PDE; have a speed-accuracy trade-off based on resolution: slow on fine grids but less accurate on coarse grids.
Neural operator: learn a family of equations from data; don’t need the explicit knowledge of the equation; much faster to evaluate than any classical method; no training needed for new equations; error rate is consistent with resolution.
1 Graph Kernel Network (GKN)
Suppose that in (1) is uniformly elliptic then the Green’s representation formula implies
where is a Newtonian potential and is an operator defined by appropriate sums and compositions of the modified trace and co-normal derivative operators [42, Thm. 3.1.6]. We have turned the PDE (1) into the integral equation (2) which lends itself to an iterative approximation architecture.
Since is continuous for all points , it is sensible to model the action of the integral operator in (2) by a neural network with parameters . To that end, define the operator as the action of the kernel on :
Kernel convolution on graphs.
where is the neighborhood of , in this case, the entire discritized domain .
Domain of Integration.
Construction of fully connected graphs is memory intensive and can become computationally infeasible for fine discretizations i.e. when is large. To partially alleviate this, we can ignore the longest range kernel interactions as they have decayed the most and change the integration domain in (3) from to for some fixed radius . This is equivalent to imposing a sparse structure on the kernel matrix so that only entries around the diagonal are non-zero and results in the complexity .
Nyström approximation.
To further relieve computational complexity, we use Nyström approximation or the inducing points method by uniformly sampling nodes from the nodes discretization, which is to approximate the kernel matrix by a low-rank decomposition
where is the original kernel matrix and is the kernel matrix corresponding to the inducing points. and are transition matrices which could include restriction, prolongation, and interpolation. Nyström approximation further reduces the complexity to .
Multipole Graph Kernel Network (MGKN)
The fast multipole method (FMM) is a systematic approach of combining the aforementioned sparse and low-rank approximations while achieving linear complexity. The kernel matrix is decomposed into different ranges and a hierarchy of low-rank structures is imposed on the long-range components. We employ this idea to construct hierarchical, multi-scale graphs, without being constraint to particular forms of the kernel . We elucidate the workings of the FMM through matrix factorization.
The key to the fast multipole method’s linear complexity lies in the subdivision of the kernel matrix according to the range of interaction, as shown in Figure 2:
where corresponds to the shortest-range interaction, and corresponds to the longest-range interaction. While the uniform grids depicted in Figure 2 produce an orthogonal decomposition of the kernel, the decomposition may be generalized to arbitrary graphs by allowing overlap.
2 Multi-scale graphs.
The coarse graph representation can be understood as recursively applying an inducing points approximation: starting from a graph with nodes, we impose inducing points of size which all admit a low-rank kernel matrix decomposition of the form (6). The original kernel matrix is represented by a much smaller kernel matrix, denoted by . As shown in Figure (2), is full-rank but very sparse while is dense but low-rank. Such structure can be achieved by applying equation (6) recursively to equation (7), leading to the multi-resolution matrix factorization :
The complexity of the algorithm is measured in terms of the sparsity of as this is what affects all computations. The sparsity represents the complexity of the convolution, and it is equivalent to the number of evaluations of the kernel network . Each matrix in the decomposition (7) is represented by the kernel matrix corresponding to the appropriate sub-graph. Since the number of non-zero entries of each row in these matrices is constant, we obtain that the computational complexity is . By designing the sub-graphs so that decays fast enough, we can obtain linear complexity. For example, choose then . Combined with a Nyström approximation, we obtain complexity.
3 V-cycle Algorithm
We present a V-cycle algorithm (not to confused with multigrid methods), see Figure 1, for efficiently computing (8). It consists of two steps: the downward pass and the upward pass. Denote the representation in downward pass and upward pass by and respectively. In the downward step, the algorithm starts from the fine graph representation and updates it by applying a downward transition . In the upward step, the algorithm starts from the coarse presentation and updates it by applying an upward transition and the center kernel matrix . Notice that the one level downward and upward exactly computes , and a full -level v-cycle leads to the multi-resolution decomposition (8).
Employing (9)-(11), we use neural networks to approximate the kernel , and neural networks to approximate the transitions . Following the iterative architecture (4), we also introduce the linear operator , denoting it by for each corresponding resolution. Since it acts on a fixed resolution, we employ it only along with the kernel and not the transitions. At each time step , we perform a full V-cycle:
Upward Pass:
We initialize as and output . The algorithm unifies multi-resolution matrix decomposition with iterative graph kernel networks. Combined with a Nyström approximation it leads to computational complexity that can be implemented with message passing neural networks. Notice GKN is a specific case of V-cycle when .
Experiments
We demonstrate the linear complexity and competitive performance of multipole graph kernel network through experiments on Darcy equation and Burgers equation.
In this section, we show that MGKN has linear complexity and learns discretization invariant solutions by solving the steady-state of Darcy flow. In particular, we consider the 2-d PDE
and approximate the mapping which is non-linear despite the fact that (14) is a linear PDE. We model the coefficients as random piece-wise constant functions and generate data by solving (14) using a second-order finite difference scheme on a fine grid. Data of coarser resolutions are sub-sampled. See the supplements for further details. The code depends on Pytorch Geometric, also included in the supplements.
We use Nyström approximation by sampling nodes for each level. When changing the number of levels, we fix coarsest level , and let , , and . This set-up is one example that can obtain linear complexity. In general, any choice satisfying also works. We set width , iteration and kernel network as a three-layer neural network with width , so that coarser grids will have smaller kernel networks.
The left most plot in figure 3 shows that MGKN (blue line) achieves linear time complexity (the time to evaluate one equation) w.r.t. the number of nodes, while GKN (red line) has quadratic complexity (the solid line is interpolated; the dash line is extrapolated). Since the GPU memory used for backpropagation also scales with the number of edges, GKN is limited to on a single G-GPU while MGKN can scale to much higher resolutions . In other words, MGKN can be applicable for larger settings where GKN cannot.
Comparing with single-graph:
As shown in figure 3 (mid), adding multi-leveled graphs helps decrease the error. The MGKN depicted in blue bars starts from a fine sampling , and adding subgraphs, , , up to, . When , MGKN and GKN are equivalent. This experiment shows using multi-level graphs helps improve accuracy without increasing much of time-complexity.
Generalization to resolution:
The MGKN is discretization invariant, and therefore capable of super-resolution.We train with nodes sampled from a resolution mesh and test on nodes sampled from a resolution mesh. As shown in on the right of figure 3, MGKN achieves similar testing error on , independently of the training discretization.
Notice traditional PDE solvers such as FEM and FDM approximate a single function and therefore their error to the continuum decreases as resolution is increased. On the other hand, operator approximation is independent of the ways its data is discretized as long as all relevant information is resolved. Therefore, if we truly approximate an operator mapping, the error will be constant at any resolution, which many CNN-based methods fail to preserve.
2 Comparison with benchmarks
We compare the accuracy of our methodology with other deep learning methods as well as reduced order modeling techniques that are commonly used in practice. As a test bed, we use 2-d Darcy flow (14) and the 1-d viscous Burger’s equations:
with periodic boundary conditions. We consider mapping the initial condition to the solution at time one . Burger’s equation re-arranges low to mid range energies resulting in steep discontinuities that are dampened proportionately to the size of the viscosity . It acts as a simplified model for the Navier-Stokes equation. We sample initial conditions as Gaussian random fields and solve (15) via a split-step method on a fine mesh, sub-sampling other data as needed; two examples of and are shown in the middle and right of figure 4.
Figure 4 shows the relative test errors for a variety of methods on Burger’s (15) (left) as a function of the grid resolution. First notice that MGKN achieves a constant steady state test error, demonstrating that it has learned the true infinite-dimensional solution operator irrespective of the discretization. This is in contrast to the state-of-the-art fully convolution network (FCN) proposed in which has the property that what it learns is tied to a specific discretization. Indeed, we see that, in both cases, the error increases with the resolution since standard convolution layer are parametrized locally and therefore cannot capture the long-range correlations of the solution operator. Using linear spaces, the (PCA+NN) method proposed in utilizes deep learning to produce a fast, fully data-driven reduced order model. The graph convolution network (GCN) method follows ’s architecture, with naive nearest neighbor connection. It shows simple nearest-neighbor graph structures are insufficient. The graph kernel network (GKN) employs an architecture similar to (4) but, without the multi-level graph extension, it can be slow due the quadratic time complexity. For Burger’s equation, when linear spaces are no longer near-optimal, MGKN is the best performing method. This is a very encouraging result since many challenging applied problem are not well approximated by linear spaces and can therefore greatly benefit from non-linear approximation methods such as MGKN. For details on the other methods, see the supplements.
The benchmark of the 2-d Darcy equation is given in Table 4, where MGKN again achieves a competitive error rate. Notice in the 1-d Burgers’ equation (Table 5), we restrict the transition matrices to be restrictions and prolongations and thereby force the kernels to be orthogonal. In the 2-d Darcy equation (Table 4), we use general convolutions (10, 11) as the transition, allowing overlap of these kernels. We observe the orthogonal decomposition of kernel tends to have better performance.
Related Works
GNN and non-sparse graphs:
A multitude of techniques such as graph convolution, edge convolution, attention, and graph pooling, have been developed for improving GNNs . Because of GNNs’ flexibility, they can be used to model convolution with different discretization and geometry . Most of them, however, have been designed for sparse graphs and become computationally infeasible as the number of edges grow. The work proposes using a low-rank decomposition to address this issue. Works on multi-resolution graphs with a U-net like structure have also recently began to emerge and seen success in imaging, classification, and semi-supervised learning . Our work ties together many of these ideas and provides a principled way of designing multi-scale GNNs. All of these works focus on build multi-scale structure on a given graph. Our method, on the other hand, studies how to construct randomized graphs on the spatial domain for physics and applied math problems. We carefully craft the multi-level graph that corresponds to multi-resolution decomposition of the kernel matrix.
Multipole and multi-resolution methods:
The works propose a similar multipole expansion for solving parametric PDEs on structured grids. Our work generalizes on this idea by allowing for arbitrary discretizations through the use of GNNs. Multi-resolution matrix factorizations have been proposed in . We employ such ideas to build our approximation architecture.
Conclusion
We introduced the multipole graph kernel network (MGKN), a graph-based algorithm able to capture correlations in data at any length scale with a linear time complexity. Our work ties together ideas from graph neural networks, multi-scale modeling, and matrix factorization. Using a kernel integration architecture, we validate our methodology by showing that it can learn mesh-invariant solutions operators to parametric PDEs. Ideas in this work are not tied to our particular applications and can be used provide significant speed-up for processing densely connected graph data.
Acknowledgements
Z. Li gratefully acknowledges the financial support from the Kortschak Scholars Program. A. Anandkumar is supported in part by Bren endowed chair, LwLL grants, Beyond Limits, Raytheon, Microsoft, Google, Adobe faculty fellowships, and DE Logi grant. K. Bhattacharya, N. B. Kovachki, B. Liu and A. M. Stuart gratefully acknowledge the financial support of the Army Research Laboratory through the Cooperative Agreement Number W911NF-12-0022. Research was sponsored by the Army Research Laboratory and was accomplished under Cooperative Agreement Number W911NF-12-2-0022. The views and conclusions contained in this document are those of the authors and should not be interpreted as representing the official policies, either expressed or implied, of the Army Research Laboratory or the U.S. Government. The U.S. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notation herein.
Broader Impact
Many problems in science and engineering involve solving complex PDE systems repeatedly for different values of some parameters. Example arise in molecular dynamics, micro-mechanics, and turbulent flows. Often such systems exhibit multi-scale structure, requiring very fine discretizations in order to capture the phenomenon being modeled. As a consequence, traditional Galerkin methods are slow and inefficient, leading to tremendous amounts of resources being wasted on high performance computing clusters every day. Machine learning methods hold the key to revolutionizing many scientific disciplines by providing fast solvers that can work purely from data as accurate physical models may sometimes not be available. However traditional neural networks work between finite-dimensional spaces and can therefore only learn solutions tied to a specific discretizations. This is often an insurmountable limitation for practical applications and therefore the development of mesh-invariant neural networks is required. Graph neural networks offer a natural solution however their computational complexity can sometime render them ineffective. Our work solves this problem by proposing an algorithm with a linear time complexity that captures long-range correlations within the data and has potential applications far outside the scope of numerical solutions to PDEs. We bring together ideas from multi-scale modeling and multi-resolution decomposition to the graph neural network community.
References
Appendix
2 Experimental Details
The steady-state Darcy flow equation used in Section 4.1 takes the form
The viscous Burgers’ equation used in Section 4.2 takes the form
with periodic boundary conditions. We consider mapping the initial condition to the solution at time one . The initial condition is generated according to where with periodic boundary conditions. We set the viscosity to and solve the equation using a split step method where the heat equation part is solved exactly in Fourier space then the non-linear part is advanced, again in Fourier space, using a very fine forward Euler method. We solve on a spatial mesh with resolution and use this dataset to subsample other resolutions.
Linear complexity and comparison with GKN.
The first two block-rows in table 2 correspond to results of MGKN on the Darcy flow problem when using a different number of subgraphs with their respective number of nodes given by the vector . For example means a multi-graph of three levels with nodes respectively for each level. The third and fourth block-rows correspond to GKN, where, in the third block, we fix the domain of integration to be a ball with radius , and, in the fourth block, a ball with radius . The number corresponds to the number of nodes sampled in the Nyström approximation. All reported errors are relative errors. The training time corresponds to epoch using data pairs, while the testing time corresponds to evaluating the PDE for 100 new queries. The left and middle images of Figure 3 are constructed from this data.
Mesh invariance.
As shown in table 3, MGKN can be trained on data with resolution and be evaluated on with data with resolution . We train a MGKN model for each of the choices . The table demonstrates that we achieve consistently low error on any pair of train-test resolutions hence we learn an infinite-dimensional mapping that is resolution invariant. The right image in Figure 3 was generated using this data.
Comparison with benchmarks.
Table 4 and 5 show the performance of different methods on Darcy flow and Burgers’ equation respectively. The training size ; the testing size .
NN is a simple point-wise feedforward neural network. It is mesh-free, but performs badly due to lack of neighbor information. NN represents the baseline of a local map.
GCN, the graph convolution network, follows the architecture in , with naive nearest neighbor connection. Such nearest-neighbor graph structure has acceptable error for very coarse grid (). For common resolution , nearest-neighbor graph can only capture a near-local map, similar to NN. It shows simple nearest-neighbor graph structures are insufficient.
FCN is the state of the art neural network method based on Fully Convolution Network . It has a good performance for small grids . But fully convolution networks are mesh-dependent and therefore their error grows when moving to a larger grid.
PCA+NN is an instantiation of the methodology proposed in : using PCA as an autoencoder on both the input and output data and interpolating the latent spaces with a neural network. The method provably obtains mesh-independent error and can learn purely from data, however the solution can only be evaluated on the same mesh as the training data. Furthermore the method uses linear spaces, justifying its strong performance on diffusion dominated problems such as Darcy flow. This is further discussed below.
RBM is the reduced basis method , a classical reduced order modeling technique that is ubiquitous in applications . It approximates the solution operator within an linear class of basis functions, requiring data as well as the variational form the problem. Since the solution manifold of (14) exhibits fast decay of its Kolgomorov -width, linear spaces are near optimal hence it is not surprising that RBM is the best performing method . Compared to deep learning approaches, RMB is significantly slower as it requires numerical integration to form and then invert a linear system for every new parameter.
GKN stands for graph kernel network with and . It enjoys competitive performance against all other methods while being able to generalize to different mesh geometries. The drawback is its quadratic complexity, constrain GKN from a large radius. Therefore GKN has higher error rates on the Burgers equation where long-range correlation is not negligible.
MKGN is our new proposed method. MKGN has slight higher error on Darcy flow where linear spaces are near-optimal. but for the harder Burger’s equation, MGKN is the best performing method. This is a very encouraging result since many challenging applied problem are not well approximated by linear spaces and can therefore greatly benefit from non-linear approximation methods such as MGKN.