Optimization-Based Algebraic Multigrid Coarsening Using Reinforcement Learning
Authors: Ali Taghibakhshi
Scott MacLachlan
Luke Olson
Matthew West
Venue: NeurIPS 2021 Poster (2021)
Keywords: Algebraic Multigrid
Reinforcement Learning
Graph Partitioning

Abstract: 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.

Optimization-Based Algebraic Multigrid Coarsening Using Reinforcement Learning
Ali Taghibakhshi
Scott MacLachlan
Mechanical Science and Engineering University of Illinois at Arbana-Champaal Urbana, IL 61801, USA
Mathematics and Statistics
Memorial University of Newfoundland and Labrador
alit2@illinois edu
John's ss NI, Canada aanhcllchaannmmnn
Luke Olson
Matthew West
Computer Science
Mechanical Science and Engineering University of Illinois Urbana-Chann Champaign
University of Illinois at Urbana-Chaman Shbama-Champaign
Urbana, IL 61801, USA lukeo@illinois edu
Urbana, IL 61801, USA mwest@illinois edu
Abstract
Large sparse linear systems of equations are ubiquitous in science and engineering, such as those arising from discretizations partial diffcrential equations. Alge- braic multigrid (AMG) methods such linear systems, with an extensive body of underlying mathematical theory. A system of linear equations defines graph 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 represen- tation. 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 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 reinforcement learning (RL) agent based on graph neural networks (GNNs), which can learn 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.
of the most common methods of solving
Introduction
widely used method for approximating the solution large, sparse linear systems, particularly those arising from elliptic boundary value problems. AMG has proven be effective for 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 construct optimal AMG methods for improved efficiency and robustness.
Algebraic multigrid (AMG) [1 2
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

AMG algorithms aim to solve sparse linear system of the form
Ax =
(I)
by successively constructing coarser representations of the problem. Constructing an AMG method is effectively a graph coarsening problem, to define coarse representation of Agt. = be, 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] 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 99 10 12 and to improve its efficiency, including by leveraging recent advances in machine learning. There are two main heuristies 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-gria carse-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] I6)), optimal partitioning into coarse and fine nodes is known be NP hard [7]. Nevertheless, the ML-based coarsening scheme that we propose yields strictly better coarsening for class of AMG solvers on planar grids compared to previous algorithms. We expand on this notion Appendix D, where we note the limitations of comparing directly to the observed convergence of AMG algorithms based only on heuristics.
As review in Section 3.1 the graph of A in (1) can be partitioned into coarse and fine (a.k.a. fine- only) 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 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 A is the Laplace operator and f(x, y) and o(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 [17 and outlined in more detail in Appendix D, we train dueling double-DDQ agent with topology adaptive convolution layers (TAGCN) [9] 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 temporal in nature, and model it as RL problem, discussed in Section 4 By using the optimization framework [17] as our starting point, we are able to prove that our RL method produces convergent multigrid solver, and by combining evaluation with graph decomposition we are able to coarsen efficiently.
algorithm 20 21
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 grid with computational expense that linear 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
2 Related work
Machine learning for multigrid. While there 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 al. [14] used a residual architecture and unsupervised learning to train fully-connected neural network that could outperform conventional

---

methods in generating the interpolation operator for diffusion problems on structured grids. In later study [33, 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. 124] trained convolutional network to improve 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 ID structured grids.
Reinforcement learning. Reinforcement learning (RL) algorithms have demonstrated success in challenging decision-makee problems such as complex games (26) and robotics 227] 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 311.
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. [3), spectral methods define graph convolution as diagonal operators in the graph Laplacian eigenbasis. These methods, however, are highly domain specific and, once trained, they cannot adapt to different graph structures. Moreover, due to the Laplacian matrix eigendecomposstton, they are computationally expensive. However, a number of recent studies [34 (35] have investigated how these limitations can be alleviated. Spatial methods, 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, 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 TAGCN for our RL agent.
3 Multigrid background
AMG algorithms infer a "grid" structure by using the graph of A in [ In the setup phase, a coarser grid with ne < nodes is constructed, by selecting subset of the nodes of the current grid. A full-rank interpolation operator, P E Raxx,, is created to map vectors from the coarse grid to the fine grid, and a voarse-grid operator, A R=* is also created, often using the Galerkin product, A. = pT AP. After setup, the AMG solve phase is executed by taking an initial guess for the solution, (0) E TR.", and performing a number of pre-relaxation (iterative linear solver) sweeps on the finest grid, often with Gauss-Seide or weighted Jacobi, to damp high-energy errors in (0()) Then, the residual is restricted to the coarse grid, and the error equation, in the form of [ 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) (8), is outlined in A Algorithm| I Note that while we give the two-grid cycle only, the solution of the system at Step 5|an 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] 9 10 II to bound the per-cycle convergence factor. Alternately, heuristic or other techniques (including ML approaches [3 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 cycle of AMG
3

---

Algorithm Two-Level AMG Algorithm
1: Input: Sparse matrix A E R=xx, right-hand side b E R", initial guess x(() E I", interpolation matrix P R coarse-grid matrix A E Rn.x*e convergence tolerance 6, numbers of relaxation sweeps N2 E N. and 0.
2: repeat:
3: Perform N1 pre-relaxation sweeps on ((k) to obtain .(aa
4: Calculate the residual: f(k) = 6 b- Aa(k)
5: Project the residual to the coarse grid and solve: Agec (R) p())
6: Interpolate and add the coarse-grid correction: 2(6) =(()) + pee
7: Perform N2 post-relaxation sweeps on ((4) to get (1++1)
8: k = k + 1, compute (k++)) = ---- (r(++1)
9: until: ||++
algorithm, assuming that the dominant costs in the algorithm are proportional to the number of degrees of freedom each level. Notably, however, common theme in the theoretical approaches mentioned above [O, [T is that is difficult to guarantee a fixed bound on the grid complexity while simultaneously bounding the resulting AMG convergence factor.
3.1 Grid coarsening and greedy coarsening
Classical AMG coarsening is based on heuristics to generate maximal independent subsets of the graph on each level [2] While often effective, there no guarantees either the grid complexities of the generated hierarchies or the convergence properties of the resulting algorithm. In an attempt to address this, MacLachlan and Saad [17] proposed optimization form of the coarsening problem, building on the work of Macl achllnn et al. (9], which provides rigorous bound on the two-level AMG convergence factor depending properties of the grid partitioning. Define binary variable indicating if a node is fine or coarse:
f = G if
ifie
where F and C' denote the sets of fine and coarse nodes, respectively. The key element in the theory of [9] is suitable diagonal dominance of the submatrix of A restricted to the F set. This is encapsulated by the following optimization problem, which defines the largest F that satisfies the dominance criterion:
maximize
fe
(4a)
subject to |asal 2 0
+ laygl for all such that fa 1.
(4b)
JEF
Here, 0 is a dominance parameter that ties directly to the convergence theory in 19 (we use 0 0.56). Thus, solving 4 offers the lowest (two-grid) complexity to achieve a solver with a certain guaranteed convergence factor per iteration. In [17), MacLachlan and Saad proved that this optimization problem NP-hard and proposed a greedy algorithm for approximating its solution. In [18) it shown that the greedy algorithm is ineffective in solving 4 in some cases, and simulated annealing approach (SA) is used instead. While that approach is effective, it has a high cost, due to the inefficiency of SA. Here, aim to apply tools from ML retain the effectiveness of the SA approach, but at lower cost. Our proposed method is linear in problem size, while SA is less efficient in comparison.
4 Reinforcement learning for multigrid coarsening
While grid coarsening could thought as a function that maps a fine grid to coarse grid, is natural to formulate it as iterative process that selects coarse nodes one by one, so that each decision can take into account the locations of previously selected coarse nodes. We thus start by reformulating the optimization problem (4 as a reinforcement learning problem.

---

4.1 Environment
To define the reinforcement learning environment, we start from the symmetric sparse linear system (1) with an x n matrix A. We define a graph G (V. E) that has node i E V for each variable index 1,....7 and an edge (i.j) E E if A., * 0 for i + j- Consider two binary variables fi E {0, 1] and e2 E {0, 1} for each node E V. Here f. indicates whether node i is fine node, as
in 3 and 24 indicates whether node violates the diagonal dominance constraint 44b). Recall that
F {i ft 1} is the set of fine nodes and C {i fi = 0} is the set of coarse nodes. The state space is now defined by S = {(fi, 2/)] i = I,.... n which is 2n binary variables, two for each node. The initial state So consists of all fine nodes, so ft = 1 for all i, and 2ú is determined by:
1 if ft 1 and |aal < 8 EEjE dnyl. 0 otherwise.
At each time step the agent chooses one violating fine node convert into a coarse node. The action
space is thus A = {i B/ = 1}. Given an action a chosen in state s, the next state s' changes S% = O, so node a becomes coarse, and va recomputed for all nodes by (5). Note that u/ = i for all nodes that are not adjacent to the new coarse node a. The reward r(s)
|C| is the negative of the number of coarse nodes and the environment terminates when there are no more actions to take (that is, when no nodes are violating).
Theorem 1. An optimal agent for the environment described above will exactly solve the optimization problem
Proof. Because each time step of the environment adds exactly one new coarse node, the total payoff function is p = (n + 1)/2, where n. the number of coarse nodes at termination. The environment terminates at the first time that the constraint (4b) is satisfied and thus an optimal agent will minimize n. subject to (4b), therefore minimizing the number of steps. Equivalently, the number of fine nodes will be maximized, which is (4a).
4.2 Agent and training
We want to Iearn an agent that is able to coarsen graphs for many different Ar = b systems, which will be of different sizes and different graph topologies. We thus use graph neural network agent that can adapt to arbitrary graph structures and a dueling architecture [1] to enable graph-size-invant invariaat advantage estimation. Specifically, of 4-size filters and 32 hidden units. This means that the agent output at each node depends on information from nodes up to 12 hops away the graph (3 layers x 4-size). We evaluate the agent on the graph G associated with the matrix A from ① using fa and vn as node features and the elements of the A matrix as edge features. TAGConv networks perform graph convolutions in the vertex domain and learn filters that are adaptive to the topology of the graph. They inherit the convolution properties of CNNs on rectangular grids but perform graph convolutions in the sense of graph signal processing,
use 3 TAGConv [19 layers for the agent, each consisting
The agent has two separate output heads in dueling architecture 31.. one to estimate the state value (ss) and the other to estimate the state-action advantage A(s, a). (In this paragraph, we temporarily overload the notation A to refer the advantage function rather than the matrix from ( and V to refer to the state value and not the vertex set of G.) These two output heads each consist of one TAGConv layer and the lower two TAGConv layers shared, with final averaging layer for the V head. The state-action value function (s, a) is, thus, the sum of the outputs from the two heads. While ducling architectures have benefits for training, they are also valuable for graph problems where the Q and V values are dependent on the size of the graph but A is not, as is the case for multigrid coarsening. This size dependence arises because Q and V are estimating the number of coarse nodes needed for the graph, which scales with the graph size. In contrast, the advantage A is only estimating the differential impact of coarsening a particular node, which will be small number that depends only on the local graph structure and does not increase with the global graph size. Using a dueling architecture, we can train on small graphs, where we will learn V and A accurately for a certain size of graph, and then evaluate the agent on large graphs where V will no longer be accurate but will continue to be correct. Since we only need the output of for evaluation, this provides a scale-in invariant solution to the reinforcement learning problem.
5

---

Algorithm 2 Evaluation Algorithm
1: Use Lloyd aggregation decompose the node into subdomains {V1, V2,.., Vk}
2: while Constraint (4b) is not satisfied do
3: Evaluate the advantage TAGCN network obtain the advantage A, for each node
4: for k to K such that V& contains at least one node with = 1 do argmax Aj ,,,,, 6: Coarsen node
5:
7: end for
8:
end while
To train the agent, we use Dueling Double DQN 30
Each episode consists of different random convex triangular mesh in 2D with approximately 30 120 nodes. Meshes are generated by selecting number of uniformly random points in 2D, taking their convex hull, and meshing the interior using pygmsh [39 (see Figure 1 as an example). On each mesh, we compute the piecewise linear finite-elemer discretization of the Laplacian. A, corresponding to the PDE (2), and perform rollout of the environment in Section 4.1 All training and testing problems used dominance factor of f = 0.56. We train for 10k episodes of DDQN with learning rate of 10-3 replay buffer of size x 103, and batch size of 32. the beginning of the training, the exploration value the DDQN algorithm is set to 1 and decayed at rate of 1.25 10-4 to minimum value of 10-2 The target network weights are updated after each episodes, and the frozen network is replaced by the target network after every 10 learning steps. See Appendix B for discussion of the choices made for the network and its architecture.
31
Figure I: Examples training grids with sizes ranging from approximately 30 to 120 nodes.
4.3 Evaluation
To coarsen a graph, we can simply apply the trained agent to the graph, coarsen the fine node with the highest advantage output, and repeat. However, a direct implementation of this will incur cost of ( for a graph with nodes, because each step to coarsen node requires an O(n) cost to evaluate the TAGCN network and there are ((n) coarse nodes required. Thus, follow 18 and use a localization strategy to reduce cost.
To perform graph coarsening with (n)) total cost, combine graph decomposition via Lloyd aggregation 20 21 with the trained agent. See Appendix for more detailed discussion of Lloyd aggregation. We first use the Lloyd aggregation algorithm decompose the node set into disjoint union of subdomains. Then, at each step, we evaluate the advantage TAGCN network on the entire graph and coarsen one node per subdomain (if there any available). By decomposing into O subdomains of bounded size, this method terminates in (1) steps and results in an (nn) cost for the entire coarsening procedure. Algorithm details this method, Theorem 2 proves that has O(n) cost, and Figure 3 (right) shows numerical verificttion of the linear cost scaling. Furthermore, Theorem uses the existing AMG theoretical results [17 prove that Algorithm necessarily results convergent multigrid method.
Theorem 2.
For grids with n nodes having n-independent bounded node degree and n-independent bounded subdomain size when using Lloyd aggregation with fixed number of cycles, the time complexity of Algorithm is O(n).
Proof. Since the Lloyd subdomain size is bounded and the number of Lloyd aggregation cycles are fixed, generating the subdomain decomposition has (n) time complexity [20]. Evaluating each

---

RL (F-fraction 0.71)
Greedy (F-fraction 0.66)
RL (F-frnction 0.73)
Greedy (F (F-fraction 0,60)
AGRIORII
Figure 2: Example coarsenings of meshes from the "Different Size" (left) and "Structured" (right)
families of test grids, comparing RL and greedy coarsening algorithms for dominance factor 0.56.
TAGCN layer consists of y = Eí=1 GeTy + bln, where is the number of node features, b is a learnable bias, Ge E Rax= is the graph filter, and Ve E In are the node features. Moreover, the graph filter is polynomial in the adjacency matrix M of the graph: G2 = Σj= M where J is a constant and Gr.j are the filter polynomial coefficients. The bounded node degree of the graph implies that M is sparse and the action M has O(n) cost, and thus the full TAGCN evaluation also has (n) cost. Because the Lloyd aggregate size is bounded, we have K (n) and so the for-loop in Algorithm 2 has (n) cost. Finally, the bounded Lloyd subdomain size also implies that the number of while-loop iterations in Algorithm iis independent of n, and thus the total cost is On)).
Theorem 3. Algorithm is guaranteed to terminate and satisfy constraint (4b) Moreover, if matrix A in linear system ① is diagonally dominant, the resulting two-grid multigrid cycle (Algorithm) guaranteed to converge with convergence factor bounded independently of n.
Proof. Since every iteration of Algorithm coarsens at least one node, the algorithm must terminate since the graph is finite. Moreover, as long as there are fine nodes violating constraint (4b). the algorithm does not terminate, and so constraint (4b) must be satisfied at termination. Finally, from Theorems 3 and in MacLachlan and Saad [17], diagonal dominance of A and (4b) are sufficient to guarantee uniform convergence of the two-grid cycle.
5 Numerical Experiments
To test the performance of the RL method for multigrid coarsening, we use six different families of test grids, as summarized in Table []] Each of these families explores different type of grid structure and has a parameter to vary this structural element. For all tests, we used the agent trained following Section 4.2 and evaluated using Algorithm 2 with mean Lloyd aggregate size of 400 nodes and 1000 Lloyd cycles. The code was implemented using PyTorch Geometric 40), PyAMG [41], and NetworkX [42]. All computation was performed using CPU-only on an 8-core 19 MacBook Pro. The training time was approximately 5 hours and the average evaluation time was approximately 20 seconds per test grid. Example grid coarsenings are shown in Figure 2
Evaluation metrics: F-fraction and effective convergence faccorr The primary metric for coarse grid quality that we are optimizing is the number of fine nodes, which should maximized as in (4a). To quantify this, use the F-fraction (higher is better), which is the ratio of the number of fine nodes to the total number of nodes. As a secondary metric, we use the effective convergence factor (lower is better) of the two-level multigrid method formed from the coarse grid using the reduction based AMG interpolation operator from 9 17]. The effective convergence factor is defined as the per-cycle AMG convergence factor raised to the power of one the AMG grid complexity, which is measure of the convergence factor per unit of computational work.
Comparison methods: Greedy and theoretical bounds. To assess the performance of our RL coarsening method, we compare with the simple greedy method of MacLachlan and Saad [17]. which progressively coarsens the grid by greedily selecting the node with the lowest value of /oal// Djer (ıipl at each step (see (4b) for reference). While the greedy method provides a good
All code and data for this paper is wtthttts:///gttit http:p//....tttttooc///////////s
(MIT licensed).

---

Family
Example
F-Fraction vs, Attribute
0.8
RL
F-Fraction
0.7
Structured: A family of 18 rectangular structured grids with different grid sizes.
0.6
Gzeedy
0.5
20000
40000
Grid
HL
F-Braction
Graded Mesh: family of 12 unstruc- tured grids with different convex shapes, and graded meshes, i.c., smoothly vary- ing mesh size across the domain.
0.700
0.675
Greedy
0.650
10000
2000
Grid Size
F-Fraction
RL
Wide Valence: A family of 12 unstruc- tured convex grids with the same size and different average node degree stan- dard deviation.
0.70
Greedy
0.65
1.0
1.5
Avg Node Degree STD
0.72
Different Size: family of 42 unstruc- tured convex grids with grid size ranging from about 60 to 52 000 nodes. The aver- age node degree standard deviation and mesh aspect ratio constant.
F-Prection
0.70
0.68
Greedy
0.06
20000
40000
Grid Size
Aspect Ratio: A family of 12 unstruc- tured convex grids with the same size and different average mesh aspect ratio, varying from about 2 to 180.
F-Fraction
0.00
RL
Greedy
0.55
100
Aspoct Ratio
0.70
RtI.
Non-Convex: family of 12 unstruc- tured non-convex grids. The grids are generated by removing geometrical shapes from main geometry and mesh- ing the remainder.
T-Feoction
0.68
Greedy
0.66
10000
20000
30000
Grid Size
Table 1: Families of test grids used for numerical experiments, showing F-fractions (higher is better).

---

x10-3
Time / Grid Size (seconds)
0.8
1.0
F-Praction
Upper bound
0.7
RL
0.5
0.6
Greedy
0.5
0.0
0
20000
40000
0.0
0.5
1.0
Grid Size
Grid Size
x10
Figure 3: Left: F-fraction for the RL method and comparison methods (higher is better) for the Structured family of grids. Right: Evaluation time divided by grid size, showing linear scaling in grid size.
baseline, we can also derive a theoretical bound for the structured grid family, as described below in Theorem
Theorem 4. For the Poisson problem 2 discretized with 5-point finite difference stencil structured grid of size Dg x By. the constraint (4b) implies that the F-fraction, f. is bounded by:
f
(6)
T2
Tly
Rg
Proof. To satisfy constraint (4b) for the 5-point stencil, every pentomino (a five cell structure obtained after removing corners of 3 grid) within the grid must have at least one coarse node. Given an n2 x Ny rectangular grid, there are at least (n, (n, 2) pentominoes entirely lying within the grid, there are at least this many coarse nodes. On the other hand, every coarse node appears in at most five different pentominoes. Letting ne be the number of coarse nodes, this implies n./ (n. gives (6).
(nz - (m, 2) < 5n.- Rearranging this inequality and using f
Structured grid coarsening results. The left panel in Figure 3 shows the performance of the RL method compared to the greedy baseline and the theoretical upper bound. For large grids, the RL method achieves 96.4% of the theoretical upper bound, which is 51.5% higher than the greedy method. Most significantly, RL performs better on all grid sizes, with improving performance the grid size increases. The right panel Figure] confirms that the time to coarsen is linear in the grid size, for structured grids with up to 1.2 million nodes.
Unstructured grid coarsening results. For the unstructured grid families, the F-fractions obtained by the RL and greedy methods are shown in the right- column in Table D Here, we see that the RL method always substantially outperforms the greedy method for every grid family, and for all parameter variations. These families include very diverse and challenging grids, including those with very wide valence (some nodes with high degree and others with low degree) and those with high aspect ratio elements (up to 180). This data summarized in the left panel of Figure 4 The right panel of this figure shows the effective convergence factors of the resulting two-level multigrid methods. The methods were not directly optimizing for this metric, but we nonetheless see that the RL and greedy grids have generally similar convergence factors. The main exceptions are for cases where the greedy method has dramatically lower F-fraction than the RL method, such as the structured grid family. Further details on these numerical results are given in Appendix A.
6 Conclusion
In this paper, we introduce a reinforcement learning (RL) method utilizing graph convolution neural networks (GCNN) to approach the NP-hard problem of grid coarsening for algebraic multigrid methods (AMG). We use our method for the 2D Poisson problem wide variety of unstructured

---

Rt. Efffective Convergence Factoor
1.0
0.8
Incressing Grid Size
Structured Grids
F-fruction
0.9
Graded Meshes
0.7 Aspect
Increasing
Wide Valence
0.8
Different Size
Ratio
Aspect Ratio
0.6
Non convex
0.7
45 Degree Line
0.5
0.5
0.7
0.6
1.0
Greody F-fraction
Groedy Effective Convergence Factor
Figure 4: Comparison between RL and greedy algorithms on the F-fraction (higher is better) and effective convergence (lower is better) metrics, for all families of test grids.
grids. Moreover, we demonstrate that our approach strictly outperforms the existing heuristic method [17. while still preserving the theoretical guarantees for convergence of the resulting multigrid solver.
Our RL agent is trained on small grids and uses a ducling advantage and value function decomposition prove that evaluation algorithm scales linearly with grid size. The main limitations of the current work is that the agent specialized to single PDE (Poisson's equation) 2D grids and it is unclear how readily it will generalize to other problems. In future work, we propose to extend the approach considered here more general families of AMG algorithms (and, indeed, to other problems and algorithms in numerical linear algebra), aiming to achieve improved efficiency in the resulting AMG solvers.
to apply to arbitrarily large grids without performance degradation. Furthermore,
graph-decomposition-bee
Funding Transparency Statement
Funding in direct support of this work: NSERC grant RGPIN- 2019-05692 SM. The authors have no competing interests to declare.
References
[1] A. Brandt, S. F. McCormick, and J. W. Ruge. Algebraic multigrid (AMG) for sparse matrix equations. D. Evans, editor, Sparsity and Its Applications. Cambridge University Press, Cambridge, 1984.
[2] John W Ruge and Klaus Stüben. Algebraic multigrid. In Multigrid methods, pages 73-130. SIAM, 1987.
[3] Klaus Stüben. Algebraic multigrid (AMG): experiences and comparisons. Appl. Math. Comput., 13(3-4):4194555 1983.
[4] D. Chen, S. MacLachlan, and M. Kilmer. Iterative parameter choice and algebraic multigrid for anisotropic diffusion denoising. SIAM Sci. Comput., 33:2972- 2994, 201
[5] Abhijit Kundu, Vibhav Vineet, and Vladlen Koltun. Feature space optimization for semantic video segmentation. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2016.
[6] H. De Sterek, Thomas A. Manteuffel, Stephen F. McCormick, Quoc Nguyen, and John Ruge. Multilevel adaptive aggregation for Markov chains, with application web ranking. SIAM J. Sci. Comput., 30(5):2235 (003)222222222 2008.
[7] A. Brandt. Algebraic multigrid theory: the symmetric case. Appl. Math. Comput., 19(1-4):23-56
1986.
[8] Mandel. Algebraic study of multigrid methods for symmetrie, definite problems. Appl. Math. Comput., 25:39-56, 1988.
10

---

[9] S. MacLachlan, T. Manteuffel, and S. McCormick. Adaptive reduction- baase AMG. Numer Linear Alg. Appl., 13:599-620, 2006.
[10] Artem Napov and Yvan Notay. An algebraic multigrid method with guaranteed convergence rate. SIAM Journal on Scientific Computing,
2012.
[11] J. Brannick, Y. Chen, J. Kraus, and I.. Zikatanov. Algebraic multilevel preconditioners for the graph Laplacian based on matching in graphs. SIAM Journal on Numerical Analysis, 5163):185-18277, 2013.
[12] Scott P. MacLachlan and Luke N. Olson. Theoretical bounds for algebraic multigrid perfor- mance: review and analysis. Numerical Linear Algebra with Applications, 21(2):194-220, 2014.
[13] Ila Luz, Meirav Galun, Haggai Maron, Ronen Basri, and Irad Yavneh. Learning algebraic multigrid using graph neural networks. International Conference on Machine Learning. pages 6489-6499 6499. PMLR, 2020.
[14] Daniel Greenfeld, Meirav Galun, Ronen Basri, Irad Yavneh, and Ron Kimmel. Learning to optimize multigrid PDE solvers. In International Conference on Machine Learning, pages 2415-2423. PMLR, 2019.
[15] Andrew Cleary, Robert D Falgout, Jim E Jones, et al. Coarse-grid selection for parallel algebraic multigrid. In International Symposium on Solving Irregularly Structured Problems in Parallel, pages 104-115. Springer, 1998.
[16] David M. Alber and Luke N. Olson. Parallel coarse-grid selection. Numerical Linear Algebra with Applications, 14(8):611-643, 2007.
[17] Scott MacLachlan and Yousef Saad. A greedy strategy for coarse-grid selection. SIAM Journal on Scientific Computing, 29 2007.
[18] U. Zaman, S. P. MacLachlan, L. Olson, and M. West. Coarse-grie grid selection using simulated annealing. arXiv preprint arXiv:2/055/33280 2021.
[19] Jian Du, Shanghang Zhang, Guanhang Wu, José MF Moura, and Soummya Kar. Topology adaptive graph convolutional networks. arXiv preprint 1710. 10370, 2017.
[20] W. N. Bell. Algebraic multigrid for discrete differential formss Ph.D. thesis, University of Illinois at Jrbana-Champaln 2008.
[21] Stuart Lloyd. Least squares quantization in PCM. IEEE transactions information theory, 28(2):129-137, 1982.
[22] C. W. Oosterlee and R. Wienands. A genetic search for optimal multigrid components within a Fourier analysis setting. SIAM Journal on Scientific Computing. 24(3):924-94) 2003.
[23] Jonas Schmitt, Sebastian Kuckuk, and Harald Köstler. Optimizing geometric multigrid methods with evolutionary computation. arviv preprint arXiv: 1900002790 2019.
[24] Jun-Ting Hsieh, Shengjia Zhao, Stephan Fismann, Lucia Mirabella, and Stefano Ermon, Learn- ing neural PDE solvers with convergence guarantees. arXiv preprint arXiv: 0et..0906000000 2019.
[25] Alexandr Katrutsa, Talgat Daulhaev, and Ivan Oseledets. Black-box learning of multigrid parameters. Journal of Computational and Applied Mathematics, 368:1 12524, 2020.
[26] Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy Lillicrap, and David Silver. Mastering Atari, Go, chess and shogi by planning with a learned model. Nature, 588:604 609, 2020.
[27] Julian Ibarz, Jie Tan, Chelsea Finn, Mrinal Kalakrishnan, Peter Pastor, and Sergey Levine. How to train your robot with deep reinforcement learning: lessons we have learned. The International Journal of Robotics Research, 40(4-5):698-72) 2021.
[28] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 518(7540):529-533 2015.
[29] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 17X1v:1707.06337, 2017.
11

---

[30] Hado Van Hasselt, Arthur Guez. and David Silver. Deep reinforcement learning with double Q-learning. Proceedings of the AAAI Conference on Artificial Intelligence, volume 30, 2016.
[31] Ziyu Wang. Tom Schaul, Matteo Hessel, Hado Hasselt, Marc Lanctot, and Nando Freitas. Dueling network architectures for deep reinforcement learning. In International conference on machine learning, pages 1995-2003. PMLR, 2016.
[32] Zonghan Wu, Shirui Pan, Fengwen Chen, Guodong Long, Chengqi Zhang, and S Yu Philip. A comprehensive survey on graph neural networks. IEEE transactions on neural networks and learning systems, 2020.
[33] Joan Bruna, Wojciech Zaremba, Arthur Szlam, and Yann LeCun. Spectral networks and locally connected networks on graphs. arXiv preprint arXiv: /3/.: 12.2203, 2013.
[34] Michaël Defferrard, Xavier Bresson, and Pierre Vandergheynst. Convolutional neural networks on graphs with fast localized spectral filtering, arXiv preprint arXiv:1606.09377 2016.
[35] Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. arXiv preprint arXiv:1609.02880 2016.
[36] Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International Conference on Machine Learning, pages 1263-1272. PMLR, 2017.
[37] Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzale Vinicius Zambaldi. Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint 1.06v:8806./// 01261, 2018.
[38] William L Briggs, Van Emden Henson, and Steve F McCormick. A multigrid tutorial. SIAM, 2000.
[39] Nico Schlömer. pygmsh: A Python frontend for Gmsh, 2021. (GPL-3.0 License).
[40] Matthias Fey and Jan E. Lenssen. Fast graph representation learning with PyTorch Geometric. In ICLR Workshop on Representation Learning on Graphs and Manifolds, 2019. (code is MIT licensed).
[41] L. N. Olson and B. Schroder. PyAMG: Algebraic multigrid solvers in Python, 2018. Release 4.0 (MIT licensed).
[42] Aric Hagberg, Picter Swart, and Daniel S Chult. Exploring network structure, dynamics, and function using networkx. Technical report, Los Alamos National Lab. (LANL), Los Alamos, NM (United States), 2008. (code is BSD licensed).
12


Reviews:
[{"review_id": "tn3g2OelfK", "reviewer": "Reviewer_xtZu", "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).\n", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "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": "Provide a full review of the submission, including its originality, quality, clarity, and significance. See https://neurips.cc/Conferences/2021/Reviewer-Guidelines for guidance on questions to address in your review, and /faq for how to incorporate Markdown and LaTeX into your review.\nPositive points:\n-- The topic is important. Solving huge sparse linear systems is a fundamental challenge in numerical linear algebra that comes up in numerous scientific fields and applications.\n-- The paper addresses a specific subproblem that has not been addressed in the deep learning literature: choosing a coarse grid. There are several complications to this problem: for example, the selection process is not differentiable, and it is often computationally hard.\n-- The proposed solution with GNN and RL seems appropriate. \n-- The paper has a theoretical analysis.\n-- The paper is well written and easy to follow. \nNegative points: \n-- I am not an expert in AMG but the evaluation seems lacking. Several questions:\nIn my understanding, the convergence factor is the main way to compare AMG methods. It seems like the convergence factor of the proposed method is not competitive.\nThe paper discusses a theoretical convergence factor calculated on a two-grid problem. Can the authors add/discuss/evaluate the real convergence factor?\nWould it be possible for the authors to add an experiment showing the convergence factor as a function of grid size/complexity?\nWould it be interesting to include convergence related losses in RL? Currently, the RL optimization only takes the sparsity of the coarse grid and the diagonal dominancy.\n\n\nOriginality: The idea is original.  This paper tackles a sub-problem of AMG that wasn't previously addressed by learning tools.\n\nQuality: This is a high-quality paper.\n\nClarity: very well written\n\nSignificance: \nBy developing better coarsening strategies, many important applications can be significantly improved. The paper could lead to further research on the coarsening problem.\n\nIn summary, I like the idea and the paper. The only reservation I have is the actual performance of the suggested method - it appears the convergence factor isn't better than the previous methods (fig 3). I would appreciate it if the authors could address my concern.\n\nMinor comments:\nwhich interpolation was used?\nL.89 GCNNS typo\nCan you discuss the choice of TAGconv? An ablation study on the GNN architecture could be helpful.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations are briefly discussed", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "PPYHGj-Fukk", "reviewer": "Reviewer_5a9n", "summary": "Algebraic multigrid method is widely used in a variety of problems including partial differential equations, image and video processing, and sparse Markov chains. This paper is the first approach of involving reinforcement learning to solving AMG coarsening and achieves superior results both theoretically and empirically.\n", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "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": "This is an interesting paper that incorporates reinforcement learning into a theoretical problem. As a novel approach, this paper stands out as the first of its kind.\n\nBy illustrating examples of applications of their algorithms, the authors demonstrated how they could be used. Also, the algorithm has a theoretical guarantee of convergence. \n\nI am not familiar with AMG-related algorithms. It is possible that I do not fully understand this work's place in this field. Nevertheless, the paper is well-written, and I like the idea it presents.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors mentioned in their paper that the main limitations is that the agents are specialized to a single PDE on 2D grids, the generalization to other problems is unclear.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "4wDnbzUGtvu", "reviewer": "Reviewer_qGDk", "summary": "The paper presents an algorithm for course node selection for algebraic multigrid methods on the poisson problem. It formulates the optimization problem as an RL setting, and solves it with DDQN.", "questions": "", "limitations": "", "rating": 4, "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": "On the ML side, there are few surprises: the paper proposes to solve Eq. 5 using RL (Dueling DDQN), and since this is a graph problem, a GNN is used for the value network. The setup (node/edge encoding, RL formulation) seems reasonable for this problem. The only thing that surprised me was that both network size (3 layer 32-width node function) and training time (5h single-actor on a laptop CPU) were magnitudes smaller than other RL problems studied in the field. Whatever coarsening heuristic & value function is learned, it's likely not very complex.\n\nIt's quite possible that the paper solves a real problem for its niche; but without being a domain expert it's hard for me to judge its significance, and I'm not sure how relevant this work is for the wider ML community. I've worked with AMG before, and haven't come across this formulation for graph coarsening; the wider NeurIPS audience will likely have even less connection to this problem. A cursory look at the cited related work like this particular problem formulation was introduced by [17] and mostly iterated on by the same research group. As the ML component to this method is fairly standard, maybe it'd make more sense to publish this work in a more specialized journal or conference.\n\nFurther comments and questions:\n- the manuscript links to a non-anonymized github repository\n- is the reward given at every timestep, or once at episode termination?\n- the only baseline comparison is to the greedy algorithm [17] from 2007. Isn't there a more recent and relevant baseline to compare to?\n- the paper mentions that it's able to apply the learned policy to bigger meshes than seen in training, but there are no details on train/test split for the results shown. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"ethical_concerns": "None.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}]

Author rebuttal:

Meta-review: There have been a few papers in the literature on using ML to "learn" interpolation operators in multigrid solvers, but to the best of my knowledge this is the first paper that attempts to learn the coarsening strategy itself. Admittedly, this work is still a little exploratory: the objective is a simple theoretical bound (which theoretically only applies to the two-grid case), the examples are very simple, there is not much discussion of wall clock time performance up to a given accuracy on real problems with real high performing AMG solvers, little exploration of the architecture of the learning agent itself with e.g. ablation studies. Nonetheless, the paper deserves attention because it has the potential to open up a refreshingly new direction of research. The reviewers seem to agree on this, as well as the fact that the paper is clearly written and does not attempt to exaggerate its contributions. AMG is a very important numerical method, even if it has heuristic elements, so if ML-based coarsening and interpolation strategies can improve it, that might have large impact.
Decision: Accept (Poster)