[论文解读] 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

一句话总结

该论文提出MLGMRES，一种实时深度学习框架，通过在仿真过程中现场学习问题特定的预条件子，加速大规模线性系统的GMRES迭代求解器。采用轻量级、非局部神经网络（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的仿真中实现显著的时间求解速度提升，特别是针对具有挑战性的多尺度特征的泊松方程。
• 实现实时HPC加速器在AI驱动求解器加速中的高效利用，且不干扰仿真工作负载。

提出的方法

• 将深度神经网络（FluidNet）直接集成到GMRES求解器循环中作为现场预条件子，在仿真过程中实时更新。
• 采用数据筛选策略，仅保留具有信息量的残差快照（通过算法1），将训练数据量减少至时间步长的约10%。
• 采用具有非局部注意力机制的卷积神经网络架构，以捕捉线性系统结构中的长程依赖关系。
• 使用流式残差和解更新数据在线训练网络，避免在大规模仿真数据集上进行预训练。
• 在每次GMRES迭代中应用学习到的预条件子，以提高收敛速度并减少时间求解。
• 在不同问题分辨率间使用共享的、固定大小的神经网络，实现无需增加模型复杂度的可扩展性。

实验结果

研究问题

• RQ1 深度学习模型是否能在无先前训练数据或专家输入的情况下，实时学习GMRES的有效预条件子？
• RQ2 非局部神经网络架构在不同问题规模和复杂源项（如多尺度偶极子）上的泛化能力如何？
• RQ3 实时在线训练对PDE仿真中求解器收敛性和时间求解的影响如何？
• RQ4 轻量级、固定大小的神经网络是否能在无需重新训练的情况下，实现不同分辨率下的稳定加速？
• RQ5 使用流式、现场数据是否能实现比无预条件子的标准GMRES更快的收敛速度？

主要发现

• MLGMRES框架在多个测试案例中，对泊松方程实现了约2倍的时间求解速度提升。
• 该方法显著提升了GMRES的收敛速度，尤其在第五次迭代的误差上，降低了约一个数量级。
• 具备学习非局部关系能力的神经网络（如FluidNet）在复杂多尺度问题上优于纯卷积网络。
• 相同的固定大小神经网络架构在不同问题分辨率下均实现了稳定的加速，无需重新训练或调整参数。
• 仅约1/10的时间步长的数据被保留在训练中，证明了高效的数据筛选和低开销。
• 该框架完全在实时环境中运行，无需预训练，支持仿真过程中的即时部署和自适应。