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

TL;DR

This paper introduces an adaptive reinforcement learning approach to dynamically select the Krylov subspace dimension in GMRES, significantly accelerating convergence for solving large sparse linear systems.

Contribution

It proposes a novel single-life reinforcement learning framework to adaptively optimize Krylov subspace size during GMRES iterations, improving convergence speed over traditional methods.

Findings

Accelerates GMRES convergence by over an order of magnitude.

Adaptive Krylov exploration outperforms constant-restart GMRES.

Effective across various matrix sizes and sparsity levels.

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.