Optimization-Based Algebraic Multigrid Coarsening Using Reinforcement Learning

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.


Reviews and author responses:

 Regarding the network size, we note that the recent graph-based method for finding AMG interpolation operators (ICML 2020, reference [13]) also uses a small network size to achieve their results. Moreover, a wide range of robotics applications of reinforcement learning have small network sizes (OpenAI gym problems).  We agree that a full study of network size and training time for both this and related problems is missing from the literature, and will include this as a topic for future work.

We view this work as being part of the lineage of other ML-for-multigrid papers, which include [13, 14, 24, 25]. We happen to have used [17] as the particular AMG method in our RL algorithm, primarily because it is one of the few AMG approaches that offer a (limited) convergence guarantee.  Nonetheless, we believe that these ideas can (and should) be adapted to other AMG methods, and indeed to other key computational linear algebra problems on graphs. We’ve added a comment to the conclusions to discuss this.

According to our problem formulation, the reward is given at every time step to the agent, and the agent’s task is to minimize the expected sum of rewards, which is equivalent to minimizing the actions taken until termination in our formulation.

Another relevant (and much more recent) baseline is the coarse-grid selection using simulated annealing (SA) in [18]. We didn’t compare directly to this approach, because it is orders of magnitude more expensive (as discussed in [18]) and, in fact, can only be run on problems which are much smaller than those considered here.  We will emphasize this in the revised work.
In fact, we only train our model on graphs of sizes 30 to 120 (Sec. 4.2 line 204) and test it on a much wider range of planar graph sizes (Table 1). We have added a figure to show examples of our training grids. We will change the caption of Table 1 to emphasize these are all test grids, and we will mention “test grids” anywhere else they have been shown in the paper.


 1- The reviewer is correct that the convergence factor is a key attribute of AMG methods.  A second key attribute is the cost of the method, typically measured by the AMG grid or operator complexity.  The “effective convergence factor” used in our plots is one way of combining both cost and convergence into a single number, defined as the (measured) convergence factor of the algorithm raised to the inverse of the cost estimate, attempting to show convergence per unit work.  This calculation necessarily produces numbers that are larger than the measured convergence factor, since the exponential power is always less than 1.  We will include plots of the measured convergence factors for comparison in the appendix of the manuscript.

Broadly, AMG methods can be divided into those that provide a priori theoretical guarantees of convergence versus heuristic methods that often do better in practice but which cannot be shown to work in all cases. Our method falls into the first category, so we expect the convergence factors to be worse than state-of-the-art heuristic methods, but emphasize that the performance we measure is within the theoretical bounds from the algorithm. We will add details regarding the convergence theory from MacLachlan and Saad [17] to the appendix.

It is an excellent point that a consequence of achieving “better” (smaller) coarse grids for these problems is that the resulting convergence factors are worse than those seen for the greedy method, although both are bounded by the same theoretical convergence factor.  We attribute this to the fact that the bound assumes that the coarse grid satisfies certain minimal properties, and the coarse grids that we achieve are much closer to this assumption than those from the greedy algorithm.

2- We will add a plot of convergence factor versus grid size and include the theoretical upper bound in the revision. In doing so, we will emphasize that the aim of our approach is to find the lowest-cost coarsening that still guarantees the theoretical upper bound. 

3- This is an excellent point that we will return to in our future work. We will add the cost of solving the actual PDEs to the RL agent.

4- We determined the interpolation using the algorithm from MacLachlan and Saad [17], which is necessary for the convergence theory presented there. We will add this in the appendix, along with other theories for guaranteed convergence.

5- We will do the ablation study and show why TAGConv is well suited for our problem.


 Coarsening is a crucial step in building effective AMG solvers since it determines the sparsity of the next level(s) and plays an important role in determining the interpolation operator which directly impacts the convergence rate. Our paper introduces the first ML-based coarsening method for AMG and achieves a strictly better coarsening for a class of AMG solvers on any planar grid compared to previous algorithms. 

In our view, this paper follows on from [13] in ICML 2020, which introduces a learning-based method for determining AMG interpolation operators, where the authors used the coarsening and sparsity structure given by heuristic approaches from classical AMG methods. Employing the coarsening algorithm from our paper and the interpolation approach from [13], one would have all of the needed pieces to build a “fully learnable” AMG solver. Therefore, our paper opens the way to end-to-end AMG solutions.

To help clarify our work’s place in the field, we will add emphasize why coarsening is an important part of AMG and how this paper, for the first time, presents a learning-based method to achieve this and shows that it produces better results than known heuristic methods.


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

By illustrating examples of applications of their algorithms, the authors demonstrated how they could be used. Also, the algorithm has a theoretical guarantee of convergence. 

I 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. 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.

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).
  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.
Positive points:
-- 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.
-- 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.
-- The proposed solution with GNN and RL seems appropriate. 
-- The paper has a theoretical analysis.
-- The paper is well written and easy to follow. 
Negative points: 
-- I am not an expert in AMG but the evaluation seems lacking. Several questions:
In 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.
The paper discusses a theoretical convergence factor calculated on a two-grid problem. Can the authors add/discuss/evaluate the real convergence factor?
Would it be possible for the authors to add an experiment showing the convergence factor as a function of grid size/complexity?
Would 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.


Originality: The idea is original.  This paper tackles a sub-problem of AMG that wasn't previously addressed by learning tools.

Quality: This is a high-quality paper.

Clarity: very well written

Significance: 
By developing better coarsening strategies, many important applications can be significantly improved. The paper could lead to further research on the coarsening problem.

In 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.

Minor comments:
which interpolation was used?
L.89 GCNNS typo
Can you discuss the choice of TAGconv? An ablation study on the GNN architecture could be helpful.
 Limitations are briefly discussed

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

It'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.

Further comments and questions:
- the manuscript links to a non-anonymized github repository
- is the reward given at every timestep, or once at episode termination?
- the only baseline comparison is to the greedy algorithm [17] from 2007. Isn't there a more recent and relevant baseline to compare to?
- 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.  Yes.

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