AK-SLRL: Adaptive Krylov Subspace Exploration Using Single-Life Reinforcement Learning for Sparse Linear System

TL;DR

GMRES solves $A x = b$ by building a Krylov subspace under memory constraints, motivating restart strategies. This paper proposes AK-SLRL, which uses single-life reinforcement learning to adaptively choose the Krylov dimension $m$ online based on the residual $r_k$ with a SAC policy, without pretraining. The main contributions are a residual-driven state and reward design for online $m$ selection, demonstration of 5–30× faster convergence across 110 SPD matrices with $n>1000$ and $m \le 20$, and analysis of memory-time tradeoffs for large-scale problems. The results indicate substantial practical speedups, enabling faster and more scalable solvers for PDE-derived and other large sparse systems.</nobracket>

Abstract

This paper presents a single-life reinforcement learning (SLRL) approach to adaptively select the dimension of the Krylov subspace during the generalized minimal residual (GMRES) iteration. GMRES is an iterative algorithm for solving large and sparse linear systems of equations in the form of \(Ax = b\) which are mainly derived from partial differential equations (PDEs). The proposed framework uses RL to adjust the Krylov subspace dimension (m) in the GMRES (m) algorithm. This research demonstrates that altering the dimension of the Krylov subspace in an online setup using SLRL can accelerate the convergence of the GMRES algorithm by more than an order of magnitude. A comparison of different matrix sizes and sparsity levels is performed to demonstrate the effectiveness of adaptive Krylov subspace exploration using single-life RL (AK-SLRL). We compare AK-SLRL with constant-restart GMRES by applying the highest restart value used in AK-SLRL to the GMRES method. The results show that using an adjustable restart parameter with single-life soft-actor critic (SLSAC) and an experience replay buffer sized to half the matrix dimension converges significantly faster than the constant restart GMRES with higher values. Higher values of the restart parameter are equivalent to a higher number of Arnoldi iterations to construct an orthonormal basis for the Krylov subspace $ K_m(A, r_0) $. This process includes constructing $m$ orthonormal vectors and updating the Hessenberg matrix $H$. Therefore, lower values of $m$ result in reduced computation needed in GMRES minimization to solve the least-squares problem in the smaller Hessenberg matrix. The robustness of the result is validated through a wide range of matrix dimensions and sparsity. This paper contributes to the series of RL combinations with numerical solvers to achieve accelerated scientific computing.

Figures (4)

• Figure 1: SLRL architecture for Adaptive changes in Krylov subspace dimension
• Figure 2: Comparison of the Euclidean norm of the residuals for the constant restart GMRES and AK-SLRL algorithms
• Figure 3: Comparison of residual curves versus the number of Arnoldi steps taken for each method
• Figure 4: Residual curves for the initial iterations of the matrix for power network problem ($1138\_bus$)