[论文解读] Optimization-Based Algebraic Multigrid Coarsening Using Reinforcement Learning

Ali Taghibakhshi, Scott MacLachlan+2 more|arXiv (Cornell University)|Jun 3, 2021

Elevator Systems and Control被引用 9

一句话总结

该论文提出了一种基于图神经网络的强化学习（RL）框架，用于在非结构化平面网格上优化二维泊松问题的代数多重网格（AMG）粗化过程。该方法在小型网格上训练RL智能体，以线性时间复杂度对更大、未见过的图进行粗化，实现了优越的粗网格质量——达到理论最优值的96.4%，在F-分数上比启发式方法（如贪心算法）高出51.5%，同时保持了多重网格的收敛性保证。

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.

研究动机与目标

• 为解决AMG粗化这一NP难问题，该问题对求解器的效率和鲁棒性至关重要。
• 开发一种可学习、可扩展且具有理论基础的粗化方法，使其能泛化至训练网格之外。
• 克服启发式粗化方法（如贪心算法）的局限性，后者在大规模或复杂网格上表现不佳。
• 在保证最终多重网格求解器理论收敛性的同时，实现线性时间推理。
• 证明强化学习可在不直接优化收敛因子的情况下，优于经典启发式方法在F-分数和收敛行为上的表现。

提出的方法

• 该方法将AMG粗化建模为一个序列决策问题，以马尔可夫决策过程（MDP）形式表示，其中RL智能体逐步选择粗化节点。
• 采用具有拓扑自适应图卷积层（TAGCN）的双DQN智能体进行训练，以最大化基于粗化质量的密集、稀疏且结构化的奖励函数。
• 粗化过程受到[17]中提出的优化框架指导，该框架通过约束确保插值质量和多重网格收敛性。
• 推理阶段使用图分解算法，通过按拓扑顺序处理节点，实现线性时间粗化。
• 智能体在小型平面网格（如10×10至50×50）上进行训练，并泛化至具有有界节点度的非结构化网格，包括最大达120万个节点的大规模网格。
• 该方法保证了最终多重网格求解器的收敛性（定理3），并实现了图规模的线性时间复杂度（定理2）。

实验结果

研究问题

• RQ1 在小型平面网格上训练的强化学习智能体，能否在保持高质量粗化表示的前提下，对具有有界节点度的大规模非结构化网格进行粗化？
• RQ2 基于RL的粗化方法是否在F-分数和多重网格收敛性方面优于经典启发式方法（如贪心算法）？
• RQ3 所提方法能否在保持AMG求解器理论收敛保证的同时，实现图规模的线性时间复杂度？
• RQ4 随着网格规模和结构复杂度（如高长宽比单元或可变节点度）的增加，该RL智能体的性能如何变化？
• RQ5 鉴于当前方法专门针对二维泊松问题，其泛化能力在多大程度上可扩展至其他问题？

主要发现

• 在结构化网格上，RL粗化方法的F-分数达到理论最优值的96.4%，显著优于仅达到理论最优值44.9%的贪心基线方法。
• 在结构化网格上，与贪心方法相比，RL方法在F-分数上提升了51.5%，且随着网格规模增大，性能进一步提升。
• 在非结构化网格上，该RL方法在各类网格族中均一致优于贪心方法，包括高长宽比（最高达180）和可变节点度的网格。
• 评估时间与网格规模呈线性关系，验证了理论上的线性时间复杂度（定理2），120万个节点的网格粗化过程高效完成。
• 尽管未直接针对收敛性进行优化，但使用RL粗化网格的两层多重网格求解器，其有效收敛因子与贪心方法相当或更优。
• 该方法保证生成收敛的多重网格求解器（定理3），确保在泛化至更大、未见过的网格时仍具备理论鲁棒性。