[Paper Review] Accelerating GMRES with Deep Learning in Real-Time

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

Model Reduction and Neural Networks19 references4 citations

TL;DR

This paper proposes MLGMRES, a real-time deep learning framework that accelerates the GMRES iterative solver for large linear systems by learning problem-specific preconditioners in situ during simulation. Using a lightweight, non-local neural network (FluidNet), the method achieves ~2× speedup in time-to-solution without pre-training or large data requirements, demonstrating robust performance across varying problem sizes and complex source terms like multi-scale dipoles.

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.

Motivation & Objective

• To accelerate iterative linear solvers like GMRES in real-time scientific simulations without relying on pre-trained data or expert-provided preconditioners.
• To develop a machine learning framework that adapts in situ during simulation, learning effective preconditioners from streaming data as the problem evolves.
• To demonstrate that lightweight, non-local neural networks can generalize across problem sizes and complex source terms without scaling model parameters.
• To achieve significant speedup in time-to-solution for PDE-based simulations, particularly for the Poisson equation with challenging multi-scale features.
• To enable efficient use of HPC accelerators for AI-driven solver acceleration in real-time, without disrupting simulation workloads.

Proposed method

• Integrates a deep neural network (FluidNet) directly into the GMRES solver loop as an in situ preconditioner, updating in real-time during simulation.
• Uses a data curation strategy that retains only informative residual snapshots (via Algorithm 1), reducing training data volume to ~10% of time steps.
• Employs a convolutional neural network architecture with non-local attention mechanisms to capture long-range dependencies in the linear system's structure.
• Trains the network online using streaming residuals and solution updates, avoiding pre-training on large simulation datasets.
• Applies the learned preconditioner at each GMRES iteration to improve convergence rate and reduce time-to-solution.
• Uses a shared, fixed-size neural network across different problem resolutions, enabling scalability without increasing model complexity.

Experimental results

Research questions

• RQ1 Can a deep learning model learn effective preconditioners for GMRES in real-time, without prior training data or expert input?
• RQ2 How well do non-local neural network architectures generalize across varying problem sizes and complex source terms like multi-scale dipoles?
• RQ3 What is the impact of real-time, online training on solver convergence and time-to-solution in PDE simulations?
• RQ4 Can a lightweight, fixed-size neural network achieve consistent acceleration across different resolutions without retraining?
• RQ5 Does the use of streaming, in-situ data lead to faster convergence than standard GMRES with no preconditioning?

Key findings

• The MLGMRES framework achieves approximately a 2× speedup in time-to-solution for the Poisson equation across multiple test cases.
• The method significantly improves the convergence rate of GMRES, particularly evident in the fifth iteration error, which is reduced by about an order of magnitude.
• Neural networks capable of learning non-local relationships (e.g., FluidNet) outperform purely convolutional networks on complex, multi-scale problems.
• The same fixed-size neural network architecture achieves consistent acceleration across different problem resolutions, without requiring retraining or scaling parameters.
• Only about 1 in 10 time steps yielded data retained for training, demonstrating efficient data curation and low overhead.
• The framework operates entirely in real-time with no pre-training, enabling immediate deployment and adaptation during simulation.