Optimization-Based Algebraic Multigrid Coarsening Using Reinforcement Learning

Large sparse linear systems of equations are ubiquitous in science and engineering , such as those arising from discretizations of partial differential equations . Algebraic multigrid ( AMG ) methods are one of the most common methods of solving such linear systems , with an extensive body of underlying mathematical theory . A system of linear equations defines a graph on the set of unknowns and each level of a multigrid solver requires the selection of an appropriate coarse graph along with restriction and interpolation operators that map to and from the coarse representation . The efficiency of the multigrid solver depends critically on this selection and many selection methods have been developed over the years . Recently , it has been demonstrated that it is possible to directly learn the AMG interpolation and restriction operators , given a coarse graph selection . In this paper , we consider the complementary problem of learning to coarsen graphs for a multigrid solver , a necessary step in developing fully learnable AMG methods . We propose a method using a reinforcement learning ( RL ) agent based on graph neural networks ( GNNs ) , which can learn to perform graph coarsening on small planar training graphs and then be applied to unstructured large planar graphs , assuming bounded node degree . We demonstrate that this method can produce better coarse graphs than existing algorithms , even as the graph size increases and other properties of the graph are varied . We also propose an efficient inference procedure for performing graph coarsening that results in linear time complexity in graph size . 1 Introduction . Algebraic multigrid ( AMG ) [ 1 , 2 , 3 ] is a widely used method for approximating the solution to large , sparse linear systems , particularly those arising from elliptic boundary value problems . AMG has proven to be effective for a variety of problems , including partial differential equations ( PDEs ) , image and video processing [ 4 , 5 ] , and sparse Markov chains [ 6 ] . The focus of this work is on developing reinforcement learning agents to construct optimal AMG methods for improved efficiency and robustness . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . AMG algorithms aim to solve a sparse linear system of the form Ax = b ( 1 ) by successively constructing coarser representations of the problem . Constructing an AMG method is effectively a graph coarsening problem , to define a coarse representation of Acxc = bc , and an interpolation problem , to define the transfer of solutions between the coarse and original ( or fine ) representation of the problem . AMG has strong theoretical support for model problems [ 2 , 7 , 8 ] , but it is important to recognize that the classical algorithm for AMG relies heavily on heuristics . Thus , there is a persistent need to both improve the general theoretical foundations of AMG [ 9 , 10 , 11 , 12 ] and to improve its efficiency , including by leveraging recent advances in machine learning . There are two main heuristics used within AMG , namely selecting the coarse nodes for a given matrix ( or graph ) and construction of the interpolation operator between coarse and fine graphs . Recently , there have been studies on developing machine learning techniques to construct the interpolation operator [ 13 , 14 ] . However , these methods require prior knowledge of the fine-coarse partitioning and the interpolation sparsity pattern . Hence , developing a learnable coarse-grid selection method is a necessary step in the development of fully learnable AMG algorithms . While much heuristic work has been done on the graph coarsening problem in this context ( e.g. , [ 15 , 16 ] ) , optimal partitioning into coarse and fine nodes is known to be NP hard [ 17 ] . Nevertheless , the ML-based coarsening scheme that we propose yields a strictly better coarsening for a class of AMG solvers on planar grids compared to previous algorithms . We expand on this notion in Appendix D , where we note the limitations of comparing directly to the observed convergence of AMG algorithms based only on heuristics . As we review in Section 3.1 , the graph of A in ( 1 ) can be partitioned into coarse and fine ( a.k.a . fineonly ) nodes ; this has previously been presented as an optimization problem [ 17 ] and accompanied with a greedy algorithm for approximating the solution . However , for computational feasibility , this approach is rather simple and is known to break down in some cases [ 18 ] . An alternative is to use combinatorial optimization approaches such as simulated annealing [ 18 ] . While this produces good coarsenings , it is infeasibly expensive for large graphs . In this study , we develop a reinforcement learning method to achieve scalable ( linear order in graph size ) and high-quality fine-coarse partitioning , focusing on discretizations of the 2D Poisson equation , −∆φ = f , ( 2 ) where ∆ is the Laplace operator and f ( x , y ) and φ ( x , y ) are real-valued functions . To the best of our knowledge , this work is the first to approach AMG coarsening using machine learning . Based on the optimization formulation of the coarsening problem introduced in [ 17 ] and outlined in more detail in Appendix D , we train a dueling double-DQN agent with topology adaptive convolution layers ( TAGCN ) [ 19 ] to yield a more efficient coarsening , as compared to previous studies . This optimization problem is tightly coupled , so that selecting one node to be in the coarse problem affects the eligibility of its neighbouring nodes for selection . Therefore , the process of coarsening is temporal in nature , and we model it as an RL problem , as discussed in Section 4 . By using the optimization framework [ 17 ] as our starting point , we are able to prove that our RL method produces a convergent multigrid solver , and by combining evaluation with a graph decomposition algorithm [ 20 , 21 ] we are able to coarsen efficiently . The main contribution of this paper is to propose an RL method to train a grid-coarsening agent for the 2D Poisson problem which : ( 1 ) can coarsen arbitrary unstructured planar grids , including those much larger than the training grids , assuming bounded node degree , ( 2 ) can coarsen a grid with computational expense that is linear in the grid size ( Theorem 2 ) , ( 3 ) is guaranteed to produce a convergent multigrid method ( Theorem 3 ) , and ( 4 ) produces superior coarse grids to existing heuristic methods ( Section 5 ) . 2 Related work . Machine learning for multigrid . While there is a long history of applying machine learning to multigrid algorithms [ 22 ] , the last two years have seen a particular focus on using neural networks to generate multigrid interpolation operators . Greenfield et al . [ 14 ] used a residual architecture and unsupervised learning to train a fully-connected neural network that could outperform conventional methods in generating the interpolation operator for diffusion problems on structured grids . In a later study [ 13 ] , the model was further improved by employing graph neural networks and extending to positive semi-definite A matrices . Given the coarse-fine splitting and the interpolation sparsity pattern , these neural networks are able to generate an interpolation operator which outperforms classical AMG methods in convergence factor . In other work , Schmitt et al . [ 23 ] used evolutionary computation methods to optimize a GMG algorithm . By learning from supervised data and using a U-Net architecture , Hsieh et al . [ 24 ] trained a convolutional network to improve a GMG algorithm for the structured Poisson problem . Katrutsa et al . [ 25 ] formulated the two-level V-cyle problem as a deep neural network and , using this formulation , they optimized the interpolation operator for GMG and evaluated the method on 1D structured grids . Reinforcement learning . Reinforcement learning ( RL ) algorithms have demonstrated success in challenging decision-making problems such as complex games [ 26 ] and robotics [ 27 ] . RL algorithms largely fall into two categories : value learning and policy optimization . Value learning algorithms such as Deep Q-Learning [ 28 ] attempt to learn the so-called Q function , a function that guides an agent to select the best action in a given state . In contrast , policy optimization algorithms such as Proximal Policy Optimization ( PPO ) [ 29 ] try to maximize the possibility of actions that lead to higher rewards . In this paper , we use Dueling Double DQN [ 30 , 31 ] . Graph neural networks . The natural modeling of many problems as graphs where local structure is important has led to the development of graph convolutional neural networks ( GCNNs ) . GCNNS are commonly categorized into two main types : spectral and spatial [ 32 ] . First introduced by Bruna et al . [ 33 ] , spectral methods define graph convolution as diagonal operators in the graph Laplacian eigenbasis . These methods , however , are highly domain specific and , once trained , they can not adapt to different graph structures . Moreover , due to the Laplacian matrix eigendecomposition , they are computationally expensive . However , a number of recent studies [ 34 , 35 ] have investigated how these limitations can be alleviated . Spatial methods , on the other hand , define graph convolution by locally propagating information . Message Passing Neural Networks ( MPNNs ) are a popular framework for spatial GCNNs [ 36 ] . A more general message passing framework was introduced by Battaglia et al . [ 37 ] . They introduced graph network blocks for learning relational data in the graph structure . Du et al . [ 19 ] introduced TAGCN , a GNN that is defined in the vertex domain of the graph . They have shown that the topologies of the learnable filters in TAGCN are adaptive to the topology of the graph on which convolution is performed . In this paper , we use a TAGCN for our RL agent . 3 Multigrid background . AMG algorithms infer a “ grid ” structure by using the graph of A in ( 1 ) . In the setup phase , a coarser grid with nc < n nodes is constructed , by selecting a subset of the nodes of the current grid . A full-rank interpolation operator , P ∈ Rn×nc , is created to map vectors from the coarse grid to the fine grid , and a coarse-grid operator , Ac ∈ Rnc×nc , is also created , often using the Galerkin product , Ac = P TAP . After setup , the AMG solve phase is executed by taking an initial guess for the solution , x ( 0 ) ∈ Rn , and performing a number of pre-relaxation ( iterative linear solver ) sweeps on the finest grid , often with Gauss-Seidel or weighted Jacobi , to damp high-energy errors in x ( 0 ) . Then , the residual is restricted to the coarse grid , and the error equation , in the form of ( 1 ) , is solved , and the solution is interpolation and added to the fine-level solution . Finally , a number of post-relaxation sweeps are performed to the corrected approximation . The AMG solve phase , which is analogous to that of geometric multigrid ( GMG ) [ 38 ] , is outlined in Algorithm 1 . Note that while we give the two-grid cycle only , the solution of the system at Step 5 can also be done approximately , recursing with the same algorithm to an arbitrary number of levels . Since the AMG solution phase is an iterative algorithm , we must assess both the cost per cycle ( measured by its grid complexity ) and the convergence factor ( measured by the reduction in the residual in each iteration ) in order to quantify efficiency . Of these , convergence is more thoroughly studied . From a theoretical perspective , typical convergence results rely on making assumptions on the coarsening and construction of the interpolation operator [ 7 , 8 , 9 , 10 , 11 ] to bound the per-cycle convergence factor . Alternately , heuristic or other techniques ( including ML approaches [ 13 , 14 ] ) are used to generate coarsenings and interpolation operators that may lead to good convergence . In contrast , grid complexity , defined as the sum of the grid sizes over all levels in the hierarchy divided by the size of the finest grid , measures the computational complexity of a cycle of an AMG Algorithm 1 Two-Level AMG Algorithm 1 : Input : Sparse matrix A ∈ Rn×n , right-hand side b ∈ Rn , initial guess x ( 0 ) ∈ Rn , interpolation matrix P ∈ Rn×nc , coarse-grid matrix Ac ∈ Rnc×nc , convergence tolerance δ , numbers of relaxation sweeps N1 , N2 ∈ N , and k = 0 . 2 : repeat : 3 : Perform N1 pre-relaxation sweeps on x ( k ) to obtain x̃ ( k ) 4 : Calculate the residual : r̃ ( k ) = b−Ax̃ ( k ) 5 : Project the residual to the coarse grid and solve : Ace ( k ) c = PT r̃ ( k ) . 6 : Interpolate and add the coarse-grid correction : x̂ ( k ) = x̃ ( k ) + Pe ( k ) c 7 : Perform N2 post-relaxation sweeps on x̂ ( k ) to get x ( k+1 ) 8 : k = k + 1 , compute r ( k+1 ) = b−Ax ( k+1 ) 9 : until : ||r ( k+1 ) || < δ algorithm , assuming that the dominant costs in the algorithm are proportional to the number of degrees of freedom on each level . Notably , however , a common theme in the theoretical approaches mentioned above [ 9 , 10 , 11 ] is that it is difficult to guarantee a fixed bound on the grid complexity while simultaneously bounding the resulting AMG convergence factor .

Summary: The paper suggests an RL framework for learning coarsening operators for algebraic multigrid (AMG). AMG is a prominent approach for solving large and sparse linear systems. The main idea of AMG is to solve the problem on several "grids" with decreasing resolution. Each grid resolution is responsible for a different error frequency.  There are two main design choices in AMG: (1) choosing a coarse grid, and (2) choosing an interpolation operator that transfers information between the grids. Previous studies have shown that learning methods are helpful for (2). The aim of this paper is to use learning for solving (1).