[論文レビュー] 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

ひとこと要約

本稿では、非構造的平面グリッド上の2次元ポアソン問題に対して、代数的多重グリッド（AMG）の粗化を最適化するための強化学習（RL）フレームワークを提案する。この手法は、小さなグリッド上で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困難な問題に対処すること。
• 学習可能でスケーラブルかつ理論的根拠を持つ粗化手法を開発し、学習済みグリッドを超えて一般化できること。
• 勾配アルゴリズムなどのヒューリスティック手法の限界を克服すること、特に大規模または複雑なグリッドでは性能が劣ること。
• 結果として得られる多重グリッドソルバーの理論的収束を保証するとともに、線形時間の推論を可能にすること。
• RLが収束係数の直接最適化なしに、古典的ヒューリスティクスをF-分数および収束挙動の面で上回ることを示すこと。

提案手法

• AMG粗化を、強化学習エージェントが段階的に粗いノードを選択する逐次意思決定問題として定式化し、マルコフ決定過程（MDP）としてモデル化する。
• 粗化品質に基づく密・疎・構造的報酬関数を最大化するように、トポロジーに適応するグラフ畳み込み層（TAGCN）を備えたデュアルダブル-DQNエージェントを訓練する。
• 粗化プロセスは[17]の最適化フレームワークに従い、補間品質と多重グリッド収束を保証する制約を定義する。
• 推論時に、ノードをトポロジカル順に処理できるように、グラフ分解アルゴリズムが用いられ、線形時間での粗化を実現する。
• エージェントは、例えば10×10から50×50の小さな平面グリッドで訓練され、ノード次数が有界な非構造的グリッド、大規模グリッド（最大120万ノード）へ一般化可能である。
• 本手法は、多重グリッドソルバーの収束を保証する（定理3）と同時に、グラフサイズに対して線形時間計算量を達成する（定理2）。

実験結果

リサーチクエスチョン

• RQ1 小さな平面グリッドで訓練された強化学習エージェントは、ノード次数が有界な大規模非構造的グリッドを、高品質な粗表現を維持したまま粗化できるか？
• RQ2 RLベースの粗化手法は、F-分数および多重グリッド収束の観点で、勾配アルゴリズムなどの古典的ヒューリスティクスを上回るか？
• RQ3 本手法は、グラフサイズに対して線形時間計算量を達成しながら、AMGソルバーの理論的収束保証を維持できるか？
• RQ4 高長方形比要素や変動するノード次数を含む、グリッドサイズおよび構造的複雑性の増加に伴い、RLエージェントの性能はどのように変化するか？
• RQ5 本手法は、現在の2次元ポアソン問題に特化しているが、その一般化性はどの程度に達するか？

主な発見

• 構造的グリッドでは、RL粗化手法が理論的上限値の96.4％に達するF-分数を達成し、勾配ベースライン（44.9％）を著しく上回っている。
• 構造的グリッドでは、RL手法が勾配法よりもF-分数で51.5％優れており、グリッドサイズが大きくなるほど性能が向上する。
• 非構造的グリッドでは、高長方形比（最大180）や変動するノード次数を有する多様なグリッドファミリーにおいて、RL手法が一貫して勾配法を上回っている。
• 評価時間はグリッドサイズに比例して増加し、理論的な線形時間計算量（定理2）を確認した。120万ノードのグリッドの粗化も効率的に実行可能であった。
• 収束係数の直接最適化を行っていないにもかかわらず、RLで粗化されたグリッドを用いた2レベルの多重グリッドソルバーは、勾配で粗化されたグリッドを用いたものと同等またはより優れた有効収束係数を示した。
• 本手法は、収束する多重グリッドソルバーを保証する（定理3）ため、大規模かつ未観測のグリッドへ一般化しても理論的頑健性を確保している。