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 's 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-grin 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)
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 ,,,,,
5:
6: Coarsen node
7: end for
8:
end while
To train the agent, we use Dueling Double DQN 30
31 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.
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 /aal// 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:://s.tttttooooc//////s///sp
(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": "v3odeNuFb0J", "reviewer": "Reviewer_hVW8", "summary": "The authors propose using RL agent via a TAGCOnv network in DQN to select coarse points in the procedure of coarsening algebraic grids. Theoretical results are established for the integration of the RL into a convergent scheme, and the complexity of the algorithm is at O(n) via a Lloyd based aggregation to constrain the coarsening to an O(n) selection.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "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": "Overall the paper clearly proposes a useful improvement of the AMG algorithm. The novelty here is in combining RL other AMG technique which Im not familiar with in prior work. The authors give sufficient empirical validation to the advantage of the suggested combination. The results are clearly of interest to practitioners who use AMG solvers in a variety of settings from solving pdes to other general systems in 2D grids. The authros clearly (and honestly) state that the generalization to other problems is still unknown.\nThere are still some details missing that I would be grateful if the authors could address in the rebuttal:\n1.\tI assume that the Lloyd based aggregation is a k-means essentially, hence the question is how to determine K so that the partition will still be O(n) but the size of each cluster will be approximately similar to that of the other clusters to allow faithful sampling from each.\n2.\tIn this context, For example, I would expect the family of Aspect ratio to be fairly challenging. For some reason (could be a typo) the blue points for that family are not given in Fig. 3-right. Can you please elaborate on what type of grids, (and most importantly why) the Lloyd could be a limiting factor? and please report the aspect ration location in Fig 3-right.\n3.\tFor fig 1 what is the greedy dominance factor?\n4.\tIn fig 2 can we also see the greedy baseline?\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "no", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "daj-oI1fDhZ", "reviewer": "Reviewer_984W", "summary": "The authors try to use reinforcement learning to investigate large sparse linear systems of equations. Meanwhile, they introduce the graph model here. ", "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": "Using reinforcement learning to solve large sparse linear systems of equations is a very interesting problem. The authors introduce the graph and GNN as the tools to investigate them. However, the quality is not enough and the content is not clear. The experiments look very good, but the theories are not very clear. If the authors can show strong theories, I would like to change my grade. Furthermore, the authors still need to show the significance of their work to the readers. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Theories related to their experiments still need to improve. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "25"}, "extra_text": {}}, {"review_id": "GFFY3Kz_D_A", "reviewer": "Reviewer_oM5M", "summary": "An AMG coarsening algorithm is proposed based on reinforcement learning on graph neural networks. The coarsening is determined from the diagonal dominance criterium for AMG as proposed by MacLachlan and Saad. Numerical experiments on a wide variety of 2D meshes (structured and unstructured) for the Poisson problem show excellent performance compared to the simple greedy algorithm proposed by MacLachlan and Saad. Theoretical results are fairly limited but the approach is very promising algorithmically. However, more meshes (3D), elliptic operators (smooth and discontinuous coefficients), and better comparison to the state-of-the art should have been tried to truly assess the novelty and viability of the proposed algorithm.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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": "Questions / remarks:\n1) Can the proposed Reinforcement Learning (RL) solution be used as starting value for the NP hard problem (4) when solved as an integer programming? This would also allow to deduce how far from optimality the RL solution is.\n\n2) How is an optimal agent defined in Thm 1? As minimising the number of actions? It is worth point out that in practice one does not find an optimal agent.\n\n3) Sec 4.2 - line 204: Please give more details (in the appendix) about the \"random convex triangular mesh\" and its FEM discretisation. How is the mesh generated? Which FEM discretisation (order of elements)?\n\n4) Please explain briefly (in the appendix) how Lloyd aggregation works and why it would result in an $O(n)$ cost.\n\n5) Fig 3: Add other metrics like AMG grid complexity and convergence factor. Effective convergence factor combines, which is nice, but could be difficult to interpret.\n\n6) Fig 2: add the type of mesh in the legend.\n\n7) While I find the results of Poisson on unstructured meshes very impressive, more experiments should be done truly assess the novelty and viability of the proposed ideas: \n - 3D (un)structured \n - Variable coefficients (smooth and discontinuous) elliptic on 2D/3D\n - Helmholtz problems\n - Compare to state-of-the-art and not only to the greedy approach: standard AMG with heuristic coarsening and other ML optimized MG.\n - Maybe not needed but nice: actually solve the PDE and compare time, cost.\n\nIt is probably not feasible to perform all these tests, but I find doing at least a few of them necessary to make the submission acceptable for Neurips. It is recommended to put much more experiments in the appendix.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "NA", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4h"}, "extra_text": {}}, {"review_id": "5wG8tpfha0X", "reviewer": "Reviewer_nCB6", "summary": "The authors in the article have made a contribution to Algebraic Multigrid Coarsening. The AMG method is used to solve sparse linear systems. They proposed to replace heuristic approach used in the coarsening process by an Reinforcement Learning scheme. They have used Duel Double DQN with TAGCN to implement the RL based scheme. The paper describes in detail the problem formulation in terms of RL and presents results of the proposed algorithms. There are some important issues of concern regarding the submission.", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "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": " - The presented algorithm is only tested on 2D Poisson problem and they admit that they do not know how the algorithm will behave on other systems. This is a big hurdle in the acceptance of the paper. The contributions in the abstract and in section 1 are not formulated to support the view described earlier. A normal reader is prone to believe that the algorithm \"...can coarsen arbitrary unstructured grids, including those much larger than the training grid\" or is able to \"...learn to perform graph coarsening on small training graphs and then be applied to unstructured large graphs\"\n- It is assumed that the structural bias learned by RL during the training will not allow it to be used in different scenarios. So the algorithm may have shown good result for one set of problems, but it may need to be trained from scratch for any other type of problem. \n- Although, the authors claim that the algorithm is faster than any other technique, but there is no performance comparison present in the article. \n- The code provided by the authors is not useful. It is hard to know the dependencies and even after that effort one comes to know that there are missing files in the code.\n- For example, this line produces error in \"from DuelingNet2 import Net as Net_MPNN\" in MG_Agent.py\n- I do not agree with the authors that the complexity of Algorithm 2 is O(n). The Complexity of Algorithm 2 depends on the structure of matrix A. Finding arg max could itself have a complexity of O(n) in normal circumstances. So the whole algorithm may have the complexity of O(n^2).\n- It is not mentioned how the accuracy of the result is affected after grid coarsening. More coarse graph may also lead to bad results.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "NA", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: ## Original AC meta-review

This paper proposes a systematic way of computing optimal coarsening for algebraic multigrid (AMG) algorithms by using a) connection of this problem to graph algorithm b) formulation of the optimal selection as a game c) solution of the RL problem using graph neural networks. The resulting algorithm is compared against classical greedy algorithm and beats it on a wide range of meshes. One of the important contributions is the usage of LLoyd aggregation that allows the learned coarsening to be computed in linear complexity. Most of the reviewers agree that the approach is interesting algorithmically and promising in practice. Some concern was raised about theoretical contribution, but it is not so clear to me where this should go: the equivalence to RL task is shown, and how good the RL algorithms solves problem is quite difficult to show.

For the experiments, as it was noted in the reviews, the resulting convergence is only shown for 2D Laplace problem, which can be considered as a minus for the paper. However, the framework that are proposed do not depend on the type of the PDE, so this would have been a nice addition, but not change the main message of the paper.

To summarize, this is one of the first works where RL approach seems to provide a practical tool for constructing AMG precondition and can be important for the field.

## SAC meta-review (the AC's recommendation was borderline)

This paper had quite opposite recommendations from the reviewers so I had to check more carefully how to calibrate them with respect to the ones on other accepted / rejected papers. I also read the paper. While I agree that this paper presents an interesting line of research that should be presented to the machine learning community, the current submission did not appear ready yet for publication at NeurIPS. In particular, as mentioned by some reviewers, the authors should make their paper more self-contained in terms of many non-standard topics in the machine learning community (like Lloyd's aggregation, which is not described at all even in an appendix). The work should also be better related to the current machine learning literature: for example, RL has already been used to solve several combinatorial algorithms (see [A],[B] e.g.), how does their approach to solve (4) fit and compare with this literature? (this was not mentioned at all in the current submission) Finally, comparing only with the greedy algorithm of [17] appears limited for such a niche area, especially since the authors already mentioned that their approach should be more effective than simulated annealing [18] in lines 147-148, but yet do not demonstrate this at all empirically or theoretically. Including such a comparison would improve the significance of the work.

I encourage the authors to take these comments in consideration for a resubmission.

[A] Bengio, Y., Lodi, A., & Prouvost, A. (2021). Machine learning for combinatorial optimization: a methodological tour d’horizon. European Journal of Operational Research, 290(2), 405-421. (on arXiv since 2018)
[B] Mazyavkina, N., Sviridov, S., Ivanov, S., & Burnaev, E. (2021). Reinforcement learning for combinatorial optimization: A survey. Computers & Operations Research, 105400. (on arXiv since 2020)
Decision: Accept (Poster)