[論文レビュー] Accelerating GMRES with Deep Learning in Real-Time

Kevin Luna, Katherine Klymko+1 more|arXiv (Cornell University)|Mar 19, 2021

Model Reduction and Neural Networks参考文献 19被引用数 4

ひとこと要約

本論文では、シミュレーション中に現地で問題固有のプリコンディショナを学習することで、大規模な線形方程式系に対するGMRES反復解法の速度を向上させるリアルタイムのディープラーニングフレームワーク、MLGMRESを提案する。軽量で非局所的なニューラルネットワーク（FluidNet）を用いることで、事前学習や大規模なデータ要件が不要でありながら、時間解法までの速度が約2倍に向上し、さまざまな問題サイズや多スケールのデュポールのような複雑な源項においても頑健な性能を示す。

ABSTRACT

GMRES is a powerful numerical solver used to find solutions to extremely large systems of linear equations. These systems of equations appear in many applications in science and engineering. Here we demonstrate a real-time machine learning algorithm that can be used to accelerate the time-to-solution for GMRES. Our framework is novel in that is integrates the deep learning algorithm in an in situ fashion: the AI-accelerator gradually learns how to optimizes the time to solution without requiring user input (such as a pre-trained data set). We describe how our algorithm collects data and optimizes GMRES. We demonstrate our algorithm by implementing an accelerated (MLGMRES) solver in Python. We then use MLGMRES to accelerate a solver for the Poisson equation -- a class of linear problems that appears in may applications. Informed by the properties of formal solutions to the Poisson equation, we test the performance of different neural networks. Our key takeaway is that networks which are capable of learning non-local relationships perform well, without needing to be scaled with the input problem size, making them good candidates for the extremely large problems encountered in high-performance computing. For the inputs studied, our method provides a roughly 2$ imes$ acceleration.

研究の動機と目的

• 事前学習済みデータやエキスパート提供のプリコンディショナに依存せずに、リアルタイムの科学的シミュレーションにおける反復的線形ソルバー（例：GMRES）の高速化を図ること。
• シミュレーション中に現地で適応する機械学習フレームワークの開発を目的とし、問題の進化に伴ってストリーミングされるデータから効果的なプリコンディショナを学習すること。
• 軽量で非局所的なニューラルネットワークが、モデルパラメータのスケーリングなしに、問題サイズや複雑な源項にわたって一般化できることを示すこと。
• PDEに基づくシミュレーションにおいて、特に挑戦的な多スケール特徴を有するポアソン方程式に対して、時間解法までの速度を顕著に向上させること。
• シミュレーションワークロードを乱さずに、AI駆動のソルバー高速化にHPCアクセラレータを効率的に活用できること。

提案手法

• 深層ニューラルネットワーク（FluidNet）をGMRESソルバーループ内に直接統合し、シミュレーション中にリアルタイムで更新されるイン・サイトプリコンディショナとして使用する。
• Algorithm 1により、情報量の多い残差スナップショットのみを保持するデータキュレーション戦略を採用し、トレーニングデータ量を時間ステップの約10%に削減する。
• 線形方程式系の構造における長距離依存関係を捉えるために、非局所アテンション機構を備えた畳み込みニューラルネットワークアーキテクチャを採用する。
• 大規模なシミュレーションデータセットでの事前学習を回避し、ストリーミング残差と解の更新を用いてオンラインでネットワークをトレーニングする。
• 各GMRES反復において学習されたプリコンディショナを適用し、収束速度の向上と時間解法までの短縮を実現する。
• 異なる問題解像度間で共通して使用可能な、固定サイズのニューラルネットワークを採用することで、モデル複雑度の増加なしにスケーラビリティを実現する。

実験結果

リサーチクエスチョン

• RQ1 ディープラーニングモデルは、事前学習データやエキスパートの入力なしに、リアルタイムでGMRESの有効なプリコンディショナを学習できるか？
• RQ2 非局所的なニューラルネットワークアーキテクチャは、さまざまな問題サイズや多スケールのデュポールのような複雑な源項にわたって、どれほど一般化できるか？
• RQ3 PDEシミュレーションにおけるリアルタイムでオンライン学習が、ソルバーの収束性と時間解法までの影響に及ぼす影響は何か？
• RQ4 軽量で固定サイズのニューラルネットワークは、再トレーニングなしに、さまざまな解像度で一貫した加速を達成できるか？
• RQ5 ストリーミングでイン・サイトのデータを使用することで、プリコンディショナなしの標準的なGMRESよりも収束が速くなるか？

主な発見

• MLGMRESフレームワークは、複数のテストケースにおいてポアソン方程式の時間解法までの速度を約2倍に向上させた。
• 本手法はGMRESの収束速度を顕著に改善し、特に第5反復目の誤差が約1桁低下した。
• 非局所的関係を学習可能なニューラルネットワーク（例：FluidNet）は、複雑で多スケールな問題において、純粋な畳み込みネットワークを上回る性能を示した。
• 同じ固定サイズのニューラルネットワークアーキテクチャが、再トレーニングやパrameterスケーリングなしに、さまざまな問題解像度で一貫した加速を達成した。
• トレーニングに使用されたデータは、時間ステップの約10分の1にまで削減された。これは、効率的なデータキュレーションと低オーバーヘッドを示している。
• フレームワークは完全にリアルタイムで動作し、事前学習が不要であるため、シミュレーション中に即座にデプロイされ、適応可能である。