Title: Neural Krylov Iteration for Accelerating Linear System Solving
Venue: NeurIPS 2024
Accepted: yes
Keywords: AI4PDE

Neural Operator

Scientific Computing

Krylov Subspace
Abstract: Solving large-scale sparse linear systems is essential in fields like mathematics, science, and engineering. Traditional numerical solvers, mainly based on the Krylov subspace iteration algorithm, suffer from the low-efficiency problem, which primarily arises from the less-than-ideal iteration. To tackle this problem, we propose a novel method, namely **Neur**al **K**rylov **It**era**t**ion (**NeurKItt**), for accelerating linear system solving.
Specifically, NeurKItt employs a neural operator to predict the invariant subspace of the linear system and then leverages the predicted subspace to accelerate linear system solving. To enhance the subspace prediction accuracy, we utilize QR decomposition for the neural operator outputs and introduce a novel projection loss function for training. NeurKItt benefits the solving by using the predicted subspace to guide the iteration process, significantly reducing the number of iterations.
We provide extensive experiments and comprehensive theoretical analyses to demonstrate the feasibility and efficiency of NeurKItt. In our main experiments, NeurKItt accelerates the solving of linear systems across various settings and datasets, achieving up to a 5.5× speedup in computation time and a 16.1× speedup in the number of iterations.

Neural Krylov Iteration for Accelerating Linear System Solving

Abstract

Solving large-scale sparse linear systems is essential in fields like mathematics, science, and engineering. Traditional numerical solvers, mainly based on the Krylov subspace iteration algorithm, suffer from the low-efficiency problem, which primarily arises from the less-than-ideal iteration. To tackle this problem, we propose a novel method, namely Neur al K rylov It era t ion ( NeurKItt ), for accelerating linear system solving. Specifically, NeurKItt employs a neural operator to predict the invariant subspace of the linear system and then leverages the predicted subspace to accelerate linear system solving. To enhance the subspace prediction accuracy, we utilize QR decomposition for the neural operator outputs and introduce a novel projection loss function for training. NeurKItt benefits the solving by using the predicted subspace to guide the iteration process, significantly reducing the number of iterations. We provide extensive experiments and comprehensive theoretical analyses to demonstrate the feasibility and efficiency of NeurKItt. In our main experiments, NeurKItt accelerates the solving of linear systems across various settings and datasets, achieving up to a 5.5× speedup in computation time and a 16.1× speedup in the number of iterations.

1 Introduction

Solving linear systems is the cornerstone of scientific computing, with applications in various fields including mathematics, science, and engineering[28]. Traditional solvers rely on the Krylov subspace iteration algorithm to tackle large-scale sparse linear systems[59]. It starts with a random initial vector, which corresponds to a one-dimensional linear subspace, and progressively expands this subspace to approximate the solution by iteratively minimizing the residual error.

However, the convergence speed and stability of the Krylov subspace iteration algorithms are significantly influenced by the dimensions and characteristics of the employed subspaces in the iterations[3]. Computational inefficiency and instability often arise in scenarios involving highdimensional problems or large matrices, especially when the matrices exhibit poor conditioning[9]. This inefficiency primarily arises from the less-than-ideal iteration, which leads to a higher number of iterations and consequently longer solving time. To address these challenges, our key insight is that prior knowledge about the linear system's invariant subspace benefits the Krylov subspace iteration by guiding the iteration process, which reduces the required number of iterations[13].

Motivated by this insight, we introduce Neur al K rylov It era t ion ( NeurKItt ), a novel method that leverages neural networks to accelerate linear system solving. NeurKItt comprises two modules:

0

10

20

30

40

Time (s)

10

11

10

9

10

7

10

5

10

3

Tolerance (2 norm)

NeurKItt None

NeurKItt BJacobi

NeurKItt SOR

NeurKItt ASM

NeurKItt ICC

NeurKItt ILU

GMRES None

GMRES SOR

GMRES ASM

GMRES ICC

GMRES ILU

0

5000

10000

15000

20000

25000

30000

35000

Average Iterations

10

11

10

9

10

7

10

5

10

3

Tolerance (2 norm)

NeurKItt None

NeurKItt BJacobi

NeurKItt SOR

NeurKItt ASM

NeurKItt ICC

NeurKItt ILU

GMRES None

GMRES SOR

GMRES ASM

GMRES ICC

GMRES ILU

the subspace prediction module and the acceleration module. The subspace prediction module uses a neural network to predict the invariant subspace of linear systems. Inspired by the concept that mapping a linear system to its invariant subspace can be viewed as an operator between two Hilbert spaces, we adopt the neural operator for subspace prediction. We also integrate thin QR decomposition and a projection loss function to optimize training, improving subspace prediction performance. For the acceleration module, it leverages the property that Krylov subspace iteration approximates an invariant subspace of the linear system. When partial information about this subspace is provided, our acceleration module utilizes it to reduce the iterations needed for precise solutions, effectively accelerating the process and addressing the challenges.

We provide comprehensive analyses to demonstrate NeurKItt's efficiency. Furthermore, extensive experiments conducted across various solver settings and different PDE problems validate the effectiveness of our approach. The results show that NeurKItt significantly accelerates linear systems solving, achieving up to a 5.5 × speedup. Both our theoretical analyses and experimental results collectively demonstrate the efficiency of NeuralKItt.

We summarize our contributions as follows:

To the best of our knowledge, our work is the first to apply the data-driven approach to optimize Krylov iteration algorithms for solving generic non-symmetric linear systems.

We introduce a novel strategy that predicts the invariant subspace of the linear system to accelerate Krylov iteration. To facilitate the subspace prediction, we design a projection loss for efficient training, in conjunction with QR decomposition for stable outputs.

Extensive experiments and theoretical analysis demonstrate that NeurKItt reduces the computational cost and the number of iterations for solving linear systems involving nonsymmetric matrices.

2 Related Work

2.1 Traditional Numerical Algorithms

In the field of computational mathematics, various algorithms have been devised to address the challenge of solving systems of linear equations. Among these, methods based on the Krylov subspace have garnered attention for their efficacy in handling large matrices[33, 43]. These methods notably mitigate computational complexity by seeking approximate solutions within a constrained, smaller subspace. Within this context, the Generalized Minimal Residual (GMRES) method [45, 14] plays a pivotal role, particularly in addressing non-symmetric matrices.

To enhance the efficiency of these Krylov subspace methods, a range of preconditioning methods are employed[9]. These methods aim to improve matrix conditioning, reduce iteration counts, and boost stability. Preconditioning varieties include Jacobi[46], Additive Schwarz Method (ASM)[64], and

Successive Over-Relaxation (SOR)[42]. They simplify the matrix structure for faster resolution, a significant advantage in large-scale problems.

Our methodology, NeurKItt, utilizes invariant subspaces to hasten convergence in Krylov subspaces, thereby expediting linear system resolutions. It is also engineered to integrate with existing preconditioning techniques, boosting its efficiency and applicability in complex computational scenarios.

2.2 Learning-based Acceleration Algorithms

Although learning-based linear system solvers require training, the training time is negligible compared to the time saved over millions of calls. Thus, there has been a surge in learning-based acceleration efforts in recent years. Research in using neural networks for accelerating linear system solving [30, 23, 60] use neural networks to optimize the Conjugate Gradient algorithm, thereby accelerating the solution of symmetric positive definite linear systems. [37] accelerated the solution of the Poisson equation. [20] focused on accelerating algorithm iterations by learning better algorithm parameters. Recent studies have utilized neural networks for matrix preconditioning. [17, 38, 53] have focused on improving algebraic multigrid preconditioning algorithms. [15, 50] have applied CNNs and machine learning, respectively, to optimize block Jacobi and ILU precondition. Furthermore, specialized preconditioning research is being conducted in areas like physics[2, 7], engineering[47], and climate science[1], demonstrating the breadth of these applications.

However, these studies often face limitations due to specific matrix properties or mainly focus on reducing low-precision solving costs, with fewer advancements in accelerating high-precision solving. In contrast, our approach, NeurKItt, addresses all these limitations. It employs subspace prediction to refine iterative methods universally across Krylov subspace algorithms. This approach not only reduces the initial iterations but also significantly improves the speed and stability of subsequent iterations.

2.3 Neural Operators

Numerical methods for solving PDEs commonly involves solving systems of linear equations, which require extensive computational resources. Recently, Neural operators (NOs), such as the Fourier Neural Operator (FNO)[31, 55, 12] and Deep Operator Network (DeepONet)[35, 36], have shown effectiveness in solving PDEs. Despite their effectiveness, NOs often grapple with challenges like diminished accuracy and stringent boundary condition prerequisites, which limit their applicability as standalone replacements for conventional algorithms in various scientific computing contexts[63, 19]. Given that mapping a matrix to its subspace is fundamentally an operator mapping task, NOs hold significant potential as tools for predicting matrix subspaces.

3 Preliminaries

3.1 Krylov Subspace Iteration

In the realm of large-scale sparse linear systems, Krylov subspace methods are frequently employed as a standard solution strategy [46, 16]. At the heart of this methodology is the principle of utilizing the matrix involved in the linear system to iteratively construct a subspace. This subspace is crafted to approximate the actual solution space, with the iterative process continuing until the generated Krylov subspace sufficiently encompasses the real solution. Such an approach enables the iterative solutions within this confined subspace to progressively approximate the accurate solution vector x , thereby efficiently converging towards it.

Suppose we now solve the system

where A ∈ C n × n . The m -th Krylov subspace associated with the matrix A and the starting vector r ∈ C n as follows:

Generally, r = Ax 0 -b represents the initial residual, where x 0 is typically generated randomly or selected as an initial guess for the solution. The iterative process of the Krylov subspace leads to the Arnoldi relation:

where V m = ( v 1 , . . . , v m ) ∈ C n × m comprises unit vectors v i ∈ C n for i = 1 , . . . , m , mutually orthogonal, and V m +1 = ( V m , v m +1 ) extends this set. H m ∈ C ( m +1) × m is an upper Hessenberg matrix. H m ∈ C m × m is the first m rows of H m , distinct due to an extra element h m +1 ,m at the m +1 -th row, m -th column. e m is the m -th unit column vector, with its m -th element as 1 .

Krylov algorithms reduce the computational effort needed for large linear systems by creating a Krylov subspace. However, starting this process with a random vector can be time-consuming. It's suggested that understanding the matrix's invariant subspace before beginning the iterative process, a warm-starting approach, could notably decrease the required iterations and improve convergence speed and stability.

3.2 Neural Operators

NeurKItt utilizes the framework of neural operator architectures, notably the FNO [26, 31, 27]. We begin by defining a spatial dimension d ∈ N and a corresponding domain D ⊂ R d . Our focus is on approximating operators G : A ( D ; R d a ) →U ( D ; R d u ) . Here, a ∈ A ( D ; R d a ) and u ∈ U ( D ; R d u ) denote functions that map the domain D to R d a and R d u , respectively, where d a , d u ∈ N . Both spaces A ( D ; R d a ) and U ( D ; R d u ) are identified as Banach spaces.

In line with the definition provided by [32], a neural operator N : A ( D ; R d a ) → U ( D ; R d u ) is conceptualized as a mapping expressed by

for a designated depth L ∈ N . In this formulation, R : A ( D ; R d a ) → U ( D ; R d v ) , with d v ≥ d u , functions as a lifting operator, while Q : U ( D ; R d v ) → U ( D ; R d u ) serves as a local projection operator.

4 Method

NeurKItt ‡ utilizes a neural operator to predict the invariant subspace ˆ K of the linear system, and then uses it to accelerate the Krylov subspace iteration. Generally, NeurKItt is broadly divided into two components.

Subspace Prediction Module : Inspired by operator learning, we employ the neural operator capable of predicting subspaces at a low cost, with the help of QR decomposition. The aim is to precisely predict the invariant subspace ˆ K of the linear system.

Acceleration Module : Originating from the mathematical theories of Krylov iteration, we have developed a Krylov algorithm tailored to our problem. Utilizing the invariant subspace, this method expedites Krylov Subspace convergence and enhances stability.

4.1 Subspace Prediction Module

The mapping from the linear system to its corresponding invariant subspace can be considered an operator, i.e., a mapping between two Hilbert spaces. Solving such problems involves finding the corresponding solution operator. In our subspace prediction module, we choose the Fourier Neural Operator (FNO)[31] for subspace prediction. We give a brief introduction to FNO in Appendix C. Generally, for a linear system Ax = b derived from the parametric PDE problem, to predict its invariant subspace ˆ K , the input to FNO is the input function a ∈ R d a from the given PDE, where d a ∈ N . We provide a detailed discussion in Appendix B about how to build a linear system problem from a PDE problem, and what is the input function.

Our task is to learn the mapping between two Hilbert spaces G : R d a → C d × n using FNO. For FNO, the lifting transformation R first lifts the input a to a higher dimensional representation v 0 ∈ C d a × c , where c is the number of channels. Then we feed the v 0 to Fourier layers. After T Fourier layers forward, we have v T ∈ C d a × c from the last Fourier layer, which keeps the same shape as v 0 . The FNO's output X = Q ( v T ) is the projection of v T by the transformation Q : C d a × c → C d A × n . NeurKItt then uses QR decomposition to orthogonalize the matrix X , obtaining the predicted

=

Problem

Krylov Subspace Iteration

Neural Operator

Solution

Solution

Poisson Equation

Helmholtz Equation

...

subspace

(c1)

(c2)

(b)

(a)

NeurKItt

A,b

A

6

Y

K

K

K

A,b

A,b

K1

(a)

Poisson Equation

Helmholtz Equation

= Darcy Flow ... subspace ˆ K . We provide more details about how to predict the subspace given the 2D Darcy flow problem in Appendix C.

We define the one -sided distance [51] from the subspace Q to the subspace C as

where Π represents the projection operator for the associated space. Its mathematical interpretation corresponds to the largest principal angle between the two subspaces Q and C [3], and defining P Q as the spectral projector onto Q . The objective of the subspace prediction task is to approximate the invariant subspace of the matrix A with our predicted subspace ˆ K :

Neural Network

where θ represents the parameters of NO, and δ ( ˆ K , S ) is the distance between the two subspaces. Based on subsequent Theoretical Analysis 5.2, we choose S as the invariant subspace associated with the smallest n eigenvalues, assuming S = span S = span { s 1 , s 2 , . . . , s n } , where s i is the eigenvector corresponding to the i -th smallest eigenvalue of the given matrix A , i = 1 , 2 , . . . , n . To reduce computational complexity, considering the norm equivalence theorem in Banach spaces, we optimize using the following loss function:

This norm transforms the computation of subspace distance into the computation of the distance from a vector to a subspace. According to the properties of the Hilbert space, there exists a unique orthogonal decomposition:

We represent the projection operator for ˆ K by Q , with P Q = QQ ∗ , leading to the relationship:

Consequently, our projection loss function l ( ˆ K , S ) is defined as follows:

We train our subspace module by optimizing the above formula. Experiments have shown that as the projection loss decreases, the principal angle of our predicted subspace also decreases. The principal angle is a mathematical indicator derived in Theoretical Analysis 5.1 that influences the acceleration.

(b)

(c2)

warm start

subspace

4.2 Acceleration Module

For Krylov algorithms, the iteration count is a critical factor influencing computational load. Drawing inspiration from the Krylov recycling algorithms[40, 13, 6, 58, 41, 24], our approach NeurKItt utilizes the predicted invariant subspace ˆ K as the deflation space within the Krylov iteration framework. We provide the pseudocode of the acceleration module in Appendix A.2. The key procedure of our implementation is outlined as follows:

Considering the linear system (1), we obtain a k -dimensional invariant subspace ˆ K from the subspace prediction module. Then, NeurKItt computes matrices U k , C k ∈ C n × k from ˆ K and A such that AU k = C k and C H k C k = I k , the specific algorithm can be found in Appendix A.1. We can get the Arnoldi relation of our NeurKItt:

This suggests that when the k -dimensional invariant subspace is given, there is no need to start from scratch for constructing a new Krylov subspace K ( A , r ) . Building upon this foundation, K ( A , r ) can converge more rapidly to the subspace where the solution x lies. This can significantly reduce the dimensionality of the final Krylov subspace, leading to a marked decrease in the number of iterations and resulting in accelerated performance. Compared to NeurKItt, GMRES can be intuitively conceptualized as the special case of our NeurKItt, where k is initialized at zero[8, 39, 40].

Existing Krylov recycling algorithms share the same Arnoldi relation with NeurKItt. However, they can only accelerate the linear systems solving when there are strongly correlated linear systems provided, which is not common in scientific computing. For example, when solving PDEs with finite element software, matrices vary in size and element positions due to discretization methods and different grids, which makes the recycling invariant subspace from previous matrices impossible. Additionally, storing subspaces from previous solving requires significant storage when the previous linear system is large. Thus, these methods are not widely used in general scenarios. In contrast, NeurKItt, which leverages neural operators to predict subspaces, is not affected by the lack of strongly correlated linear systems.

5 Theoretical Analysis

5.1 Convergence Analysis

Let Q be an l -dimensional invariant subspace of matrix A , and let C = range( C k ) be a k -dimensional space ( k ≥ l ) selected to approximate Q . We refer to Theorem 3.1 in [40] for an in-depth convergence analysis under NeurKItt.

Theorem 5.1. [40] Considering a space C , define V = range( V m -k +1 H m -k ) be the ( m -k ) -dimensional Krylov subspace generated by NeurKItt as in (11). Let r 0 ∈ C n , and r 1 = ( I -Π C ) r 0 . Then, for each Q such that δ ( Q , C ) < 1 ,

where γ = ∥ ( I -Π C ) P Q ∥ 2 and δ = δ ( Q , C ) .

The left-hand side is the residual norm subsequent to m -k iterations of NeurKItt, employing the predictive subspace C . In contrast, on the right-hand side, the initial term epitomizes the convergence of a deflated problem, given that all components within the subspace Q have been eradicated[39, 49]. The subsequent term on the right embodies a constant multiplied by the residual following m -k iterations of NeurKItt, when solving for r 1 . Supposed that the predictive space C encompass an invariant subspace Q , then we have δ = γ = 0 for the given Q , which ensures that the convergence rate of NeurKItt matches or surpasses that of the deflated problem. In most cases, ∥ P Q ∥ 2 is numerically stable and not large, therefore a reduced value of δ directly correlates with faster convergence in NeurKItt.

Now we compare two approaches to accelerating the linear systems solving[33]:

Providing an initial prediction for the solution, allows Krylov algorithms to start from this prediction. The solution involved might come from solutions to similar systems or predictions such as via neural networks[62].

Predicting the matrix invariant subspace, using this approximate subspace to speed up Krylov iterations.

For the first approach, particularly with large matrices derived from PDEs, the norm ∥ A ∥ 2 tends to be large. Our theoretical analysis in Appendix D suggests that FNO's direct solution prediction will not significantly accelerate the linear system solving, possibly only reducing a few Krylov subspace iterations. NeurKItt belongs to the second approach. This crucially speeds up the convergence of Krylov subspace iterations, markedly reducing total iterations and improving stability.

5.2 Subspace Property Analysis

The crucial questions about predicted subspace arise when accelerating Krylov subspace iterations: (1) What kind of subspace should be selected for acceleration? (2) Given the computational cost associated with employing neural networks, is it necessary to expend substantial computational resources to predict a highly accurate subspace?

Regarding the first question : As indicated by our convergence analysis 5.1, effectively accelerating the solution of linear systems only requires a predicted invariant subspace of matrix A . A matrix has many invariant subspaces, but which of these are easier to learn by the neural network[40]?

To investigate this problem, we simplify the specific scenario of the linear systems problem. The following definitions and assumptions for Theorem 5.2 are from the reference [24]. Specifically, We deal with a Hermitian positive definite matrix A and a corresponding Hermitian perturbation E , allowing A to have the eigendecomposition[24]:

where Q = [ Q 1 Q 2 Q 3 ] is an orthogonal matrix, Λ 1 = diag ( λ (1) 1 , . . . , λ (1) j 1 ) , Λ 2 and Λ 3 are defined in a similar fashion. Let the eigenvalues satisfy the following ordering * :

We consider the change in the invariant subspace range( Q 1 ) under a symmetric perturbation E of A . Let θ 1 ( · , · ) denote the largest canonical angle between two spaces. We do not require that || E || F be small, but we assume that the projection of E onto the subspace range( [ Q 1 Q 2 ] ) is small. We assume that || [ Q 1 Q 2 ] H E || F ≤ ϵ and that ϵ is small relative to λ (2) 1 -λ (1) j 1 . We also assume that η = || Q H 3 E || F is small relative to λ (3) 1 -λ (1) j 1 . Note that we do not need to assume that λ (2) j 2 -λ (1) j 1 is large. Also, let

Theorem 5.2. [24] Let A be Hermitian positive definite and have the eigendecomposition given in (13), and let E , ϵ , η , µ , and ˜ µ be defined as above. Then there exists a matrix ˜ Q 1 conforming to Q 1 such that range( ˜ Q 1 ) is a simple invariant subspace of A + E , and

The specific proof can be found in [24]. A similar bound applies to the perturbation of the eigenvalues associated with Q 1 . In the context of Theorem 5.1 and our NeurKItt, Q 1 corresponds to Q , whereas Q 2 and Q 3 can be chosen to fit the theorem.

This theorem shows that if changes in a matrix occur in the subspace corresponding to larger eigenvalues, the subspace associated with the smallest eigenvalues is minimally affected, as long

as the changes are smaller than the gap between the smallest and larger eigenvalues. To facilitate learning of matrix invariant subspaces, we use the subspace composed of eigenvectors corresponding to the smallest eigenvalues as the prediction target for our neural operator.

Regarding the second question : The analysis of the first question, along with the error convergence rate proved in Theorem 5.1, suggests that investing in substantial computational resources for predicting a high-precision subspace is unnecessary. As long as the predicted subspace shows a distinct correlation with the matrix invariant subspace, particularly when δ is small, significant acceleration can be achieved without resorting to high precision[13]. This principle underpins our development of a low-precision, cost-effective subspace prediction framework. Although the subspace prediction may not be precise, the final algorithm has achieved remarkable acceleration results.

6 Experiment

6.1 Experiment Settings

To comprehensively evaluate the performance of NeurKItt, we conducted extensive experiments. Our analysis centers on two primary performance metrics viewed through three perspectives. These tests are conducted on 3 datasets. Specifically, the three Perspectives are: (i) Matrix preconditioning techniques, spanning 7 standard methods. (ii) Accuracy criteria for linear system solutions, emphasizing 8 distinct tolerances. (iii) Different matrix sizes, considering 3 variations. For more details, please refer to the Appendix F.1.

Baselines. NeurKItt focuses on solving linear systems that involve large sparse non-symmetric matrices. The GMRES algorithm is widely used for non-symmetric large sparse linear system solving, which serves as the predominant solution. And we set it as the benchmark for our study. We use GMRES from PETSc (version 3.19.4).

Datasets. To investigate the algorithm's adaptability to various types of matrices, we examined three different linear equation challenges, each rooted in a PDE: 1. Helmholtz Equation [62]. 2. Darcy Flow Problem [31, 44, 27, 36]; 3. Non-uniform Heat Conduction Equation [48, 25, 4, 18]. For an in-depth exposition of the dataset and its generation, kindly refer to the Appendix F.4. For the runtime environment, refer to Appendix F.2.

Metrics. We adopt time speedup for "GMRES solving time / NeurKItt solving time" and iteration speedup for "GMRES iteration count / NeurKItt iteration count". Speedup over 1 denotes better NeurKItt performance. We also provide the average time spent and average iteration count in Appendix H. We use principal angle in Equation 5 to show how close the predicted subspace and the target subspace are.

6.2 Main Experiment

We compare NeurKItt with GMRES under various preconditioners and tolerance, and Table 1 presents main experimental results. More detailed experimental results about time and number of iterations can be found in Appendix H.

The results show that across all tolerances and preconditioning techniques, NeurKIt has been effective in accelerating the linear system solving and reducing the number of iterations required by the Krylov subspace method, achieving up to a 5.5× speedup in computation time and a 16.1× speedup in the number of iterations. These results demonstrate that the invariant subspace predicted by NeurKIt greatly enhances the convergence speed of the Krylov algorithms, thereby reducing the iteration count and accelerating the solving.

With the solution accuracy increasing, the acceleration does not significantly decrease. Though the acceleration effect of NeurKItt varies under different preconditioning, the minimum time speedup is no less than 1.67 in the Darcy flow problem, which suggests NeurKItt can be effectively combined with various preconditioning methods.

6.3 Generalization of Matrix Size

10000

40000

80000

120000

160000

1.0

1.5

2.0

2.5

3.0

3.5

4.0

4.5

Time Speedup

None,tol = 1e-4

None,tol = 1e-8

10000

40000

80000

120000

160000

4

6

8

10

12

Iteration Speedup

None,tol = 1e-4

None,tol = 1e-8

Matrix Size

Matrix Size

We conducted experiments on the Darcy Flow problem with varying matrix sizes under different preconditioning methods and tolerances. We provide the results in Figure 3. The experimental results show that our method consistently improves its acceleration effect as the matrix size increases, across different preconditioning methods and tolerances, which indicates that NeurKItt especially benefits solving large sparse linear systems.

6.4 Ablation Study

We conducted ablation experiments to further validate the performance of FNO and demonstrate the effectiveness of the designed projection loss. The results of the ablation experiments are shown in Table 2. The results demonstrate that FNO significantly outperforms MLP when using the same projection loss. This indicates that FNO is indeed capable of learning a better mapping from the parameter matrix to the subspace, resulting in predicted subspaces that are closer to the target subspaces.

In addition, we also attempted to utilize the Mean Squared Error (MSE) Loss for subspace prediction. The experimental results show that applying the MSE Loss on subspace learning

Table 2: Comparison of different settings on principal Angle (rad) for the Darcy flow problem, with a matrix size of 32,400. We replace the FNO with MLP for "w/o FNO" and replace the projection loss with MSE loss for "w/o Projection Loss".

is merely ineffective. Since we report the principal angle in radians, it implies that the predicted subspace ˆ K is merely orthogonal to the target subspace S . Such subspace ˆ K cannot be effectively utilized to accelerate linear system solving. Besides, it can be inferred that QR decomposition we have

employed plays an important role in subspace prediction. These results verifies that the projection loss and QR decomposition enable the model to better predict subspace.

7 Limitation and Conclusions

Limitation Our method only considers solving non-symmetric matrices, but for matrices with specific structures, we need to tail our method to achieve faster solving. Furthermore, we have not considered the potential impact of different preconditioning techniques on the subspace and have not optimized for different preconditioning techniques, though the influence caused by preconditioning has little impact based on the observation of our experiments.

Conclusions We propose NeurKItt, a novel method for accelerating linear systems solving through subspace prediction. To the best of our knowledge, this is the first attempt to apply a data-driven approach to optimize the Krylov subspace algorithm for non-symmetric linear systems. Specifically, we employ FNO to predict the invariant subspace and design a projection loss function for this task. By utilizing the predicted invariant subspace, we achieve accelerated Krylov subspace iterations. Theoretical analysis and extensive experiments demonstrate that NeurKItt effectively reduces computational costs and iteration counts.

References

Jan Ackmann, Peter D Düben, Tim N Palmer, and Piotr K Smolarkiewicz. Machine-learned preconditioners for linear solvers in geophysical fluid flows. arXiv preprint arXiv:2010.02866 , 2020.

Yael Azulay and Eran Treister. Multigrid-augmented deep learning preconditioners for the helmholtz equation. SIAM Journal on Scientific Computing , (0):S127-S151, 2022.

Christopher Beattie, Mark Embree, and John Rossi. Convergence of restarted krylov subspaces to invariant subspaces. SIAM Journal on Matrix Analysis and Applications , 25(4):1074-1109, 2004.

James V Beck, Ben Blackwell, and Charles R St Clair. Inverse heat conduction: Ill-posed problems . James Beck, 1985.

Michele Benzi, Carl D Meyer, and Miroslav Tma. A sparse approximate inverse preconditioner for the conjugate gradient method. SIAM Journal on Scientific Computing , 17(5):1135-1149, 1996.

Philipp Birken, Jurjen Duintjer Tebbens, Andreas Meister, and Miroslav Tma. Preconditioner updates applied to cfd model problems. Applied Numerical Mathematics , 58(11):1628-1641, 2008.

Salvatore Calì, Daniel C Hackett, Yin Lin, Phiala E Shanahan, and Brian Xiao. Neural-network preconditioners for solving the dirac equation in lattice gauge theory. Physical Review D , 107(3):034508, 2023.

Luiz Mariano Carvalho, Serge Gratton, Rafael Lago, and Xavier Vasseur. A flexible generalized conjugate residual method with inner orthogonalization and deflated restarting. SIAM Journal on Matrix Analysis and Applications , 32(4):1212-1235, 2011.

Ke Chen. Matrix preconditioning techniques and applications , volume 19. Cambridge University Press, 2005.

Lu Cheng and Kuan Xu. Solving time-dependent pdes with the ultraspherical spectral method, 2023.

Eric De Sturler. Truncation strategies for optimal krylov subspace methods. SIAM Journal on Numerical Analysis , 36(3):864-889, 1999.

Wenhan Gao, Ruichen Xu, Hong Wang, and Yi Liu. Coordinate transform fourier neural operators for symmetries in physical modelings. Transactions on Machine Learning Research , 2024.

André Gaul. Recycling krylov subspace methods for sequences of linear systems. 2014.

Gene H Golub and Charles F Van Loan. Matrix computations . JHU press, 2013.

Markus Götz and Hartwig Anzt. Machine learning-aided numerical linear algebra: Convolutional neural networks for the efficient preconditioner generation. In 2018 IEEE/ACM 9th Workshop on Latest Advances in Scalable Algorithms for Large-Scale Systems (scalA) , pages 49-56. IEEE, 2018.

Anne Greenbaum. Iterative methods for solving linear systems . SIAM, 1997.

Daniel Greenfeld, Meirav Galun, Ronen Basri, Irad Yavneh, and Ron Kimmel. Learning to optimize multigrid pde solvers. In International Conference on Machine Learning , pages 2415-2423. PMLR, 2019.

Stefan Güttel and Françoise Tisseur. The nonlinear eigenvalue problem. Acta Numerica , 26:1-94, 2017.

Zhongkai Hao, Songming Liu, Yichi Zhang, Chengyang Ying, Yao Feng, Hang Su, and Jun Zhu. Physics-informed machine learning: A survey on problems, methods and applications. arXiv preprint arXiv:2211.08064 , 2022.

Jun-Ting Hsieh, Shengjia Zhao, Stephan Eismann, Lucia Mirabella, and Stefano Ermon. Learning neural pde solvers with convergence guarantees. arXiv preprint arXiv:1906.01200 , 2019.

Thomas JR Hughes. The finite element method: linear static and dynamic finite element analysis . Courier Corporation, 2012.

Claes Johnson. Numerical solution of partial differential equations by the finite element method . Courier Corporation, 2012.

Ayano Kaneda, Osman Akar, Jingyu Chen, Victoria Alicia Trevino Kala, David Hyde, and Joseph Teran. A deep conjugate direction method for iteratively solving linear systems. In International Conference on Machine Learning , pages 15720-15736. PMLR, 2023.

Misha E Kilmer and Eric De Sturler. Recycling subspace information for diffuse optical tomography. SIAM Journal on Scientific Computing , 27(6):2140-2166, 2006.

Seid Koric and Diab W Abueidda. Data-driven and physics-informed deep learning operators for solution of heat conduction equation with parametric heat source. International Journal of Heat and Mass Transfer , 203:123809, 2023.

Nikola Kovachki, Samuel Lanthaler, and Siddhartha Mishra. On universal approximation and error bounds for fourier neural operators. The Journal of Machine Learning Research , 22(1):13237-13312, 2021.

Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces. arXiv preprint arXiv:2108.08481 , 2021.

Steven J Leon, Lisette G De Pillis, and Lisette G De Pillis. Linear algebra with applications . Pearson Prentice Hall Upper Saddle River, NJ, 2006.

Randall J LeVeque. Finite volume methods for hyperbolic problems , volume 31. Cambridge university press, 2002.

Yichen Li, Peter Yichen Chen, Wojciech Matusik, et al. Learning preconditioner for conjugate gradient pde solvers. arXiv preprint arXiv:2305.16432 , 2023.

Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint arXiv:2010.08895 , 2020.

Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Graph kernel network for partial differential equations. arXiv preprint arXiv:2003.03485 , 2020.

Jörg Liesen and Zdenek Strakos. Krylov subspace methods: principles and analysis . Numerical Mathematics and Scie, 2013.

Chih-Jen Lin and Jorge J Moré. Incomplete cholesky factorizations with limited memory. SIAM Journal on Scientific Computing , 21(1):24-45, 1999.

Lu Lu, Pengzhan Jin, and George Em Karniadakis. Deeponet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators. arXiv preprint arXiv:1910.03193 , 2019.

Lu Lu, Xuhui Meng, Shengze Cai, Zhiping Mao, Somdatta Goswami, Zhongqiang Zhang, and George Em Karniadakis. A comprehensive and fair comparison of two neural operators (with practical extensions) based on fair data. Computer Methods in Applied Mechanics and Engineering , 393:114778, 2022.

Kevin Luna, Katherine Klymko, and Johannes P Blaschke. Accelerating gmres with deep learning in real-time. arXiv preprint arXiv:2103.10975 , 2021.

Ilay Luz, Meirav Galun, Haggai Maron, Ronen Basri, and Irad Yavneh. Learning algebraic multigrid using graph neural networks. In International Conference on Machine Learning , pages 6489-6499. PMLR, 2020.

Ronald B Morgan. Gmres with deflated restarting. SIAM Journal on Scientific Computing , 24(1):20-37, 2002.

Michael L Parks, Eric De Sturler, Greg Mackey, Duane D Johnson, and Spandan Maiti. Recycling krylov subspaces for sequences of linear systems. SIAM Journal on Scientific Computing , 28(5):1651-1674, 2006.

Michael Lawrence Parks. The iterative solution of a sequence of linear systems arising from nonlinear finite element analysis, 2005.

William H Press. Numerical recipes 3rd edition: The art of scientific computing . Cambridge university press, 2007.

Ouyuan Qin and Kuan Xu. Solving nonlinear odes with the ultraspherical spectral method. arXiv preprint arXiv:2306.17688 , 2023.

Md Ashiqur Rahman, Zachary E Ross, and Kamyar Azizzadenesheli. U-no: U-shaped neural operators. arXiv preprint arXiv:2204.11127 , 2022.

Youcef Saad and Martin H Schultz. Gmres: A generalized minimal residual algorithm for solving nonsymmetric linear systems. SIAM Journal on scientific and statistical computing , 7(3):856-869, 1986.

Yousef Saad. Iterative methods for sparse linear systems . SIAM, 2003.

Johannes Sappl, Laurent Seiler, Matthias Harders, and Wolfgang Rauch. Deep learning of preconditioners for conjugate gradient solvers in urban water related problems. arXiv preprint arXiv:1906.06925 , 2019.

Rishi Sharma, Amir Barati Farimani, Joe Gomes, Peter Eastman, and Vijay Pande. Weaklysupervised deep learning of heat transport via physics informed loss. arXiv preprint arXiv:1807.11374 , 2018.

Valeria Simoncini and Daniel B Szyld. On the occurrence of superlinear convergence of exact and inexact krylov subspace methods. SIAM review , 47(2):247-272, 2005.

Rita Stanaityte. ILU and Machine Learning Based Preconditioning for the Discretized Incompressible Navier-Stokes Equations . PhD thesis, University of Houston, 2020.

G. W. Stewart and Ji guang Sun. Matrix Perturbation Theory . Academic Press, 1990.

John C Strikwerda. Finite difference schemes and partial differential equations . SIAM, 2004.

Ali Taghibakhshi, Scott MacLachlan, Luke Olson, and Matthew West. Optimization-based algebraic multigrid coarsening using reinforcement learning. Advances in neural information processing systems , 34:12129-12140, 2021.

Andrea Toselli and Olof Widlund. Domain decomposition methods-algorithms and theory , volume 34. Springer Science & Business Media, 2004.

Alasdair Tran, Alexander Mathews, Lexing Xie, and Cheng Soon Ong. Factorized fourier neural operators. In The Eleventh International Conference on Learning Representations .

Lloyd N Trefethen. Spectral methods in MATLAB . SIAM, 2000.

Lloyd N Trefethen. Spectra and pseudospectra: the behavior of nonnormal matrices and operators. 2020.

Hong Wang, Zhongkai Hao, Jie Wang, Zijie Geng, Zhen Wang, Bin Li, and Feng Wu. Accelerating data generation for neural operators via krylov subspace recycling. arXiv preprint arXiv:2401.09516 , 2024.

David S Watkins. The matrix eigenvalue problem: GR and Krylov subspace methods . SIAM, 2007.

Cheng Yang, Xubo Yang, and Xiangyun Xiao. Data-driven projection method in fluid simulation. Computer Animation and Virtual Worlds , 27(3-4):415-424, 2016.

David Young. Iterative methods for solving partial difference equations of elliptic type. Transactions of the American Mathematical Society , 76(1):92-111, 1954.

Enrui Zhang, Adar Kahana, Eli Turkel, Rishikesh Ranade, Jay Pathak, and George Em Karniadakis. A hybrid iterative numerical transferable solver (hints) for pdes based on deep operator network and relaxation methods. arXiv preprint arXiv:2208.13273 , 2022.

Xuan Zhang, Limei Wang, Jacob Helwig, Youzhi Luo, Cong Fu, Yaochen Xie, Meng Liu, Yuchao Lin, Zhao Xu, Keqiang Yan, Keir Adams, Maurice Weiler, Xiner Li, Tianfan Fu, Yucheng Wang, Haiyang Yu, YuQing Xie, Xiang Fu, Alex Strasser, Shenglong Xu, Yi Liu, Yuanqi Du, Alexandra Saxton, Hongyi Ling, Hannah Lawrence, Hannes Stärk, Shurui Gui, Carl Edwards, Nicholas Gao, Adriana Ladera, Tailin Wu, Elyssa F. Hofgard, Aria Mansouri Tehrani, Rui Wang, Ameya Daigavane, Montgomery Bohde, Jerry Kurtin, Qian Huang, Tuong Phung, Minkai Xu, Chaitanya K. Joshi, Simon V. Mathis, Kamyar Azizzadenesheli, Ada Fang, Alán Aspuru-Guzik, Erik Bekkers, Michael Bronstein, Marinka Zitnik, Anima Anandkumar, Stefano Ermon, Pietro Liò, Rose Yu, Stephan Günnemann, Jure Leskovec, Heng Ji, Jimeng Sun, Regina Barzilay, Tommi Jaakkola, Connor W. Coley, Xiaoning Qian, Xiaofeng Qian, Tess Smidt, and Shuiwang Ji. Artificial intelligence for science in quantum, atomistic, and continuum systems. arXiv preprint arXiv:2307.08423 , 2023.

Xuejun Zhang. Multilevel schwarz methods. Numerische Mathematik , 63(1):521-539, 1992.

A Algorithmic Details

A.1 computes matrices U k and C k

NeurKItt can be modified to solve (1) by carrying over subspace ˆ K = span { y 1 , y 2 , . . . , y k } with respect to A as follows [11, 40]:

where the matrix Y k = [ y 1 , y 2 , . . . , y k ] . The matrices U k and C k to satisfy

A.2 NeurKItt Pseudocode

Algorithm 1: NeurKItt Krylov subspace iteration

1 Select m , the maximum size of the Krylov subspace. Define k as the dimension of the predicted invariant subspace ˆ K , where ˆ K = span { y 1 , y 2 , . . . , y k } .

2 Given matrix A , use subspace prediction module to predict the matrix Y k = [ y 1 , y 2 , . . . , y k ] and set tol as the convergence tolerance. Choose an initial guess x 0 . Calculate r 0 = b -Ax 0 , and initialize i = 1 .

3 Let [ Q , R ] be the reduced QR-factorization of AY k . 4 C k = Q 5 U k = Y k R -1 6 x 1 = x 0 + U k C k H r 0 7 r 1 = r 0 -C k C k H r 0 8 while ∥ r i ∥ 2 > tol do 9 i = i +1 10 Perform m -k Arnoldi steps with the linear operator ( I -C k C H k ) A , letting v 1 = r i -1 / ∥ r i -1 ∥ 2 and generating V m -k +1 , H m -k and B m -k . 11 Let D k be a diagonal scaling matrix such that ˜ U k = U k D k , where the columns of ˜ U k have unit norm. 12 ̂ V m = [ ˜ U k V m -k ] 13 ̂ W m +1 = [ C k V m -k +1 ] 14 G m = [ D k B m -k 0 ˜ H m -k ] 15 Solve min ∥ ̂ W H m +1 r i -1 -G m y ∥ 2 for y . 16 x i = x i -1 + ̂ V m y 17 r i = r i -1 -̂ W m +1 G m y 18 Compute the k eigenvectors ˜ z i of G H m G m ˜ z i = ˜ θ i G H m ̂ W H m +1 ̂ V m ˜ z i associated with smallest magnitude eigenvalues ˜ θ i and store in P k . 19 ˜ Y k = ̂ V m P k 20 Let [ Q , R ] be the reduced QR-factorization of G m P k . 21 C k = ̂ W m +1 Q 22 U k = ˜ Y k R -1

23 x i as the numerical solution, r i as the residual, i as the iteration count.

B From Partial Differential Equation to Linear System

We model a steady-state PDE problem as a corresponding system of linear equations using discrete numerical methods such as FDM, FEM, and FVM [52, 21, 22, 29, 10]. Numerical methods discretize a partial differential equation problem by mapping it from an infinite-dimensional function space to a finite-dimensional space, resulting in a system of linear equations.

We take the process of discretizing the two-dimensional inhomogeneous Helmholtz equation using the Finite Difference Method (FDM) as an example to illustrate how to transform a PDE into a system of linear equations:

This equation could be approximated as:

We use a 2 × 2 internal grid (i.e., N x = N y = 2 and ∆ x = ∆ y ), the unknowns u i,j can be arranged in row-major order as follows: u 1 , 1 , u 1 , 2 , u 2 , 1 , u 2 , 2 . For central differencing on a 2 × 2 grid, the vector b will contain the values of f i,j = f ( x i , y j ) and the linear equation system Ax = b can be expressed as:

In this example, we call the discretization of function k 2 ( x, y ) the input function, corresponding to a 2 × 2 grid. The function inputs vary from PDE problems.

C A Brief Introduction to Fourier Neural Operator

The Fourier Neural Operator (FNO) aims to map a continuous input function A to its corresponding continuous output function U within the Fourier domain. For end-to-end training on FNO, the function pairs ( A , U ) are discretized into instance pairs ( a, u ) . The purpose is to learn a mapping N between ( a, u ) , which can be expressed as:

The mapping N consists of several sequential steps: first, the input channel is lifted using R ; next, the mapping is performed through L Fourier layers {L 1 , L 2 , . . . , L L } ; finally, the output is projected back to the original channel using Q .

Q and R are pixel-wise transformations that can be implemented using models like MLP.

A typical Fourier layer consists of a pixel-wise linear transformation, characterized by weight W and bias b , along with an integral kernel operator K :

where σ is the nonlinear activation function. The integral kernel operator K composes three steps: Fast Fourier Transformation (FFT), spectral linear transformation, and inverse FFT.

In the subspace prediction task, we use projection layer Q mapping the last fno layer output to the target space. Each channel in the FNO output represents one of the basic vector of the predicted invariant subspace. We then apply QR decomposition to the output. Let the input function be a ∈ R d a × d a from the Darcy flow problem, where d a ∈ N is the resolution of the input function, which yields a linear system Ax = b for numerical solving, where matrix A ∈ R d A × d A . We aim to predict a subspace ˆ K , as a matrix X ∈ C d A × n , for matrix A in the linear system. First, the FNO's

lifting layer R transforms a to a higher-dimensional representation v 0 ∈ C d a × d a × c , while the 3rd dimension is channel dimension. After T Fourier layers forward, we have the v T ∈ C d a × d a × c . Because Fourier layers keep the shape unchanged, we apply a transformation Q to map the v T to the desired space, ˆ K = Q ( v T ) , with Q : C d a × d a × c → C d A × n .

In practice, transformation Q is a stack of transformation layers. It first flattens the first and second dimension of v T , obtaining q 0 ∈ C d 2 a × c . Then a fully-connected neural network (FNN) applies the mapping C d 2 a × c → C d 2 a × n to q 0 , obtaining q 1 ∈ C d 2 a × n . And another FNN applies the mapping C d 2 a × n → C d A × n to q 1 , obtaining the output X ∈ C d A × n . Finally, we apply QR decomposition to X for orthogonalizing, obtaining ˆ K = span { X } .

D Analyzing the Efficiency of Using FNO Predictions as Initial Solutions in Linear System Solvers

D.1 Theoretical Analysis

The central premise of this discussion hinges on the potential application of Neural Operators (NOs), such as the Fourier Neural Operator (FNO) [31], in addressing Partial Differential Equation (PDE) problems. Specifically, it explores the feasibility of utilizing these neural operators to predict solutions to PDE problems. Such predictions could serve as preliminary solution vectors for Krylov subspace methods in the resolution of corresponding linear systems. This approach holds promise for expediting the solution process.

However, the practicality of this concept, particularly in contexts where FNO is commonly discussed, warrants careful consideration. This is especially true in scenarios involving large matrix linear systems that emerge from Finite Difference Methods (FDM) [52], spectral methods [56], and similar approaches. A significant obstacle to the application of this method is the disparity between the type of relative error typically discussed in the context of neural operators and the relative error conventionally employed in computational mathematics. This divergence is further complicated by the impact of the norm of the PDE discretization matrix, which exhibits a tendency to increase with the refinement of the grid.

In the context of linear systems derived from PDEs, we consider the following standard formulation:

Within the framework of NOs, the prevalent metric for error assessment is the Mean Squared Error (MSE) of the solution function or solution vector x , designated as rtol NO :

where ˆ x NO represents the solution as predicted by the NO, while x denotes the precise solution.

In the domain of computational mathematics, the relative residual rtol is commonly utilized as the error metric for resolving linear systems:

where ˆ x represents the solution as predicted by the Krylov subspace method.

Utilizing orthogonal polynomial grid theory, specifically Theorem 30.1 in Trefethen's work [57], page 290, offers crucial insights for this analysis.

In this context, N denotes the count of discretization points in a one-dimensional Ordinary Differential Equation (ODE), and D N represents the discretization matrix constructed using Chebyshev orthogonal polynomial zeros as grid points, with the matrix dimension being N . This assertion holds similarly for other orthogonal polynomial grids, and a uniformly divided grid can be analogized to a grid defined by the zeros of Fourier polynomials.

The theorem suggests that with an increase in the number of grid points, resulting in a larger matrix size, the norm of the matrix ∥ A ∥ 2 associated with the discretized PDE also escalates. Then We consider the following equation:

Consequently, when ∥ A ∥ 2 is substantially large, there emerges a notable discrepancy between rtol NO and rtol . Generally, the predictive accuracy rtol NO of the FNO lies in the range of 10 -2 to 10 -4 . However, if the norm ∥ A ∥ 2 approaches the magnitude of 10 5 , the relative tolerance rtol as per computational mathematics algorithms may only range between 10 0 and 10 2 . In such instances, leveraging the FNO's predicted solution as the initial vector for solving the linear system may not substantially expedite the solution process.

D.2 Experimental Analysis

To empirically corroborate the theoretical framework presented, we delve into the Heat dataset F.4.3 utilized in this study, characterized by a matrix size of 52900. The experimental results underscore several noteworthy observations.

Employing the original FNO for training, we observed rtol NO ∼ 6 ∗ 10 -3 , which is in line with the benchmarks reported in the foundational FNO paper. However, because ∥ A ∥ 2 ∼ 1 . 2 ∗ 10 6 , it results in rtol ∼ 60 . This empirical finding is consistent with our theoretical predictions. In such a context, utilizing the FNO-derived solution as the initial vector for the linear system does not markedly expedite the resolution process.

This analysis implicitly suggests that utilizing predicted solutions as initial vectors in iterative algorithms mandates a high level of precision in the neural network model. Nonetheless, it is observed that in standard application scenarios, the impact on accelerating the process is marginal.

In contrast to this method, our algorithm NeurKItt, by predicting the invariant subspace, not only reduces the initial number of iterations but also, more importantly, substantially improves the convergence rate in subsequent iterations, As demonstrated in the convergence analysis 5.1. This marks a pivotal deviation from former techniques that focus on expediting the process by merely predicting initial solutions. Our theoretical analysis 5.2 shows that the predicted invariant subspace does not require high precision to accelerate the solution of linear systems and significantly reduce the number of iterations.

E Thin QR Decomposition

In our research, we utilize the Thin QR decomposition algorithm, a prevalent technique in numerical linear algebra [14]. The essence of Thin QR decomposition lies in its capacity to factorize a matrix into two distinct components: an orthogonal matrix Q and an upper triangular matrix R . This factorization is instrumental in diminishing both the computational complexity and the storage demands.

For a matrix A belonging to C n × m with the condition n ≥ m , the Thin QR decomposition is mathematically articulated as:

where Q is an n × m orthogonal matrix that complies with Q ∗ Q = I , and R is an m × m upper triangular matrix. In our study, the integration of QR decomposition within the neural network architecture is strategically employed to effectively resolve the orthogonal basis of a subspace.

F Details of Experimental Data and Setup

F.1 Specific parameters of the main experiment

Model Training and Predicting : We employ the 5 FNO layers with modes from {20, 32, 40} and the width from {32, 50, 64}. The learning rate is fix at 1 × 10 -3 while the batch size is selected from {16, 32}. The number of subspace dimensions is fixed at 10 in all experiments. For darcy flow and

heat datasets, we generate 8000 samples for training and 1600 samples for testing. For Helmholtz, we generate 1000 samples for training and 200 samples for testing.

Baseline : To achieve maximum efficiency, we initiated our process by creating linear equation systems of PDEs using either Python. These were subsequently processed within the C programming environment. For the GMRES solver, we employed the most recent version of PETSc, specifically 3.19.4.

Three Perspectives :

1. Precondition : The act of preconditioning is vital when tackling large matrix equations, as it significantly speeds up the resolution process and augments algorithmic stability. Given the variety of scenarios that may arise, we explored over 10 prevalent preconditioning methods, adapting them to suit different contexts.

2. Tolerance : The precision threshold of a solution directly influences both the number of iterations required and the overall computational time. Different algorithms manifest varying rates of convergence, and specific NOs hold unique tolerance standards. Our comprehensive evaluation encompassed a range of tolerances, with a focus on pinpointing 5-8 levels of optimal error precision for our analyses.

3. Matrix Dimensionality : The efficiency of an algorithm can vary with changes in matrix size. Furthermore, different NOs often require matrices of various sizes. Therefore, our investigation included the study of 5 distinct matrix sizes.

Two Performance Metrics : 1. Computational Time : This metric provides a clear and direct measure of an algorithm's performance. 2. Iteration Count : The count of iterations required reflects the algorithm's numerical sensitivity and stability, offering valuable insights into its performance characteristics.

F.2 Environment

To ensure consistency in our evaluations, all comparative experiments were conducted under uniform computing environments. Specifically, the environments used are detailed as follows:

CPU: Intel ® Xeon ® Gold 6246R CPU @ 3.40GHz

GPU: NVIDIA GeForce RTX 3090

F.3 Precondition

In this study, we implemented a series of preconditioning methods for our experiments, which are outlined below.

1. None: In scenarios where 'none' is selected for preconditioning, the linear system is addressed directly in its original form. This means that no preliminary transformations or adjustments are applied to the system before the iterative solving process.

Diagonal Preconditioning (Jacobi) [46]: This technique focuses solely on the diagonal components of the coefficient matrix. The preconditioner in this approach is essentially the inverse of these diagonal elements, offering a straightforward preconditioning solution.

Block Jacobi (BJacobi) [5]: A more advanced form of Jacobi preconditioning, Block Jacobi divides the coefficient matrix into smaller, manageable blocks, each representing a specific subdomain or a separate problem. These blocks are then preconditioned individually using their diagonal elements.

Successive Over-relaxation (SOR) [61]: SOR, a modification of the Gauss-Seidel method, incorporates a relaxation factor to enhance the process of convergence. This method reformulates the original problem by applying a weight, often improving the convergence rate for certain iterative algorithms.

Additive Schwarz Method (ASM) [54]: ASM operates on the principle of domain decomposition, partitioning the main problem into several smaller subdomains. Each subdomain's problem is solved independently, and the local solutions are then aggregated to form the overall solution.

6. Incomplete Cholesky (ICC) [34]: ICC is a preconditioning method based on the Cholesky decomposition, but drops certain off-diagonal elements during the decomposition, making it "incomplete". It's utilized for symmetric positive definite problems.

7. Incomplete LU (ILU) [46]: ILU is based on LU decomposition, but like ICC, drops certain off-diagonal elements during the decomposition. ILU can be applied to nonsymmetric problems.

F.4 Datasets

F.4.1 Helmholtz Equation

We consider a two-dimensional Helmholtz equation, which can be described by the following equation [62]:

where u ( x, y ) is the wave function, and k ( x, y ) denotes the spatially varying wavenumber, with f representing the source term, corresponding to origins of electromagnetic or acoustic waves. This equation is pivotal in several physical realms, notably in Acoustics, Electromagnetism, and Quantum Mechanics. Within these fields, the Helmholtz equation models phenomena such as wave propagation and vibration patterns.

In the context of the Helmholtz equation, the wavenumber k ( x, y ) is intimately connected to the frequency of the wave and the characteristics of the medium through which the wave propagates. In our experimental setup, the value of k ( x, y ) is determined using the GRF method [35] while the f ( x, y ) is the constant.

F.4.2 Darcy Flow problem

We consider two-dimensional Darcy flows, which can be described by the following equation [31, 44, 27, 36]:

where K is the permeability field, h is the pressure, and f is a source term which can be either a constant or a space-dependent function.

In our experiment, K ( x, y ) is derived using the GRF method [35]. The term f ( x, y ) is consistently set as a constant f .

F.4.3 Non-uniform Heat Conduction Equation

We consider a two-dimensional heat conduction equation in a non-uniform medium, which is formulated as follows [4, 18]:

where T ( x, y ) is the temperature, k ( x, y ) represents the spatially varying thermal conductivity, and q ( x, y ) T ( x, y ) models internal or external heat sources/sinks. The term f ( x, y ) acts as an external heat influence.

This equation is crucial for studying heat distribution in materials with varying properties, such as in the thermal management of electronic devices. Solving this equation helps predict temperature variations, essential for designing effective cooling or heating systems.

In our experiment, k ( x, y ) is derived using the GRF [35]. And q ( x, y ) is samples from U (0 , 1) . The term f ( x, y ) is consistently set as a constant f .

G Analysis of Hyperparameters

G.1 Hyperparameters of Fourier Neural Operator

Fourier Neural Operator has three key hyperparameters: model layers, mode and width for fourier layer. We conduct experiments to investigate the impacts of these hyperparameters. Results in Table 8 suggest that a proper combination of these hyper-parameters will improve the performance.

2

5

10

15

20

25

Subspace Dimension

0.75

1.00

1.25

1.50

1.75

2.00

2.25

2.50

2.75

Time Speedup

None,tol = 1e-5

None,tol = 1e-12

2

5

10

15

20

25

Subspace Dimension

2

3

4

5

6

7

8

9

10

Iteration Speedup

None,tol = 1e-5

None,tol = 1e-12

G.2 Impact of Invariant Subspace Dimension

We conduct extensive experiments investigate the acceleration impact of subspace dimension. Results in Figure 4 shows that with the subspace dimension increasing, iteration speedup increase as well.

However, the time speedup does not consistently increase with the subspace dimension. As shown in the left figure of Figure 4, the time speedup reaches its maximum when the subspace dimension is 8. When the subspace dimension increase from 8, the time speedup decreases. This may be influenced by the subspace size since the subspaces are dense matrices. Recalling the Arnoldi relation in Section 4.2, larger subspaces mean larger dense matrix multiplications. When the subspace dimension is sufficiently large, the additional time cost of dense matrix operations will offset the time savings from reduced iterations, thereby reducing the overall time speedup. This indicates that a larger subspace dimension is not necessarily better.

H Detailed experimental results

H.1 Time Efficiency Analysis for Subspace Prediction Module

Table 9 shows the time cost of the subspace prediction module for a single pass. The first row lists datasets with matrix side lengths. It's obvious that across various datasets, the computational cost of the subspace prediction module is extremely low, typically taking only a few milliseconds.

For comparison, in Darcy Flow problem, solving a linear system using GMRES without preconditioning at a tolerance of 1e-12 requires 34.8 seconds. While the neural network's inference cost is much more lower, which is negligible.

H.2 Training Time Analysis

0

5

10

15

20

25

30

0

2

4

6

8

10

Projection Loss

Helmholtz 62500

0

100

200

300

400

500

600

0

2

4

6

8

10

Projection Loss

Darcy 160000

Heat 90000

Time(min)

Time(min)

We train each dataset for 120 epochs, and list the training costs in Table 10. The total training time for the Helmholtz dataset is shorter compared to the other two datasets. This is because the problem is simple, allowing us to train on a dataset with fewer samples, and the matrix size involved in the problem is smaller than Darcy Flow and Heat problem.

As shown in the Figure 5, the training typically converges within 100 min. For the Helmholtz problem, it converges in 10 min; for the Darcy Flow dataset and Heat dataset, it converges less than 100 min. However, in scientific computing software, where solving linear systems is frequently required and often time-consuming, the time saved by our approach will add-up to pay off the training costs.

H.3 Darcy Flow problem

H.4 Non-uniform Heat Conduction Equation

H.5 Helmholtz Equation

I Do we need to learn the neural network for matrices of a fixed size?

An important question is whether a model trained on a fixed-size dataset can be used to predict other matrices under the same PDE problem, which is common in scientific computing. In this section, we attempt to address this question using Darcy Flow dataset. First, we downsample the input function of the larger Darcy Flow dataset to generate new Darcy Flow datasets according to the discretization method mentioned in B. Specifically, we downsample the original dataset with a matrix size of 158,404 to obtain Darcy Flow datasets with matrix sizes of 39,204 and 9,604.

Next, we use the subspace prediction module trained on the Darcy Flow dataset with a matrix size of 9,604 to predict the subspaces of larger datasets. The input to the module is the input function from the larger matrix dataset, which is downsampled to the same size as the input function of the 9,604 matrix dataset, and the output is the predicted subspace of the original matrix dataset size, recovered using cubic spline interpolation.

In our experiments, the principal angles obtained through this method are generally small. For example, the model trained on the 9,604 matrix dataset predicted a principal angle of 0.07 for the test set of the 39,204 matrix dataset and 0.08 for the test set of the 158,404 matrix dataset. The experimental results in Table 14 are based on the acceleration results using the predicted invariant subspaces.

These results indicate that our method can scale across different matrix sizes for the given parametric PDE, which demonstrates that NeurKItt is practical. Additionally, using this approach for subspace prediction significantly reduces the required training time.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The main claims in the abstract and introduction accurately reflect the contribution and scope of the paper.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We refer to Section 7 for the discussion of our limitations.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We have cite the specific paper with the assumptions and proofs for the Theorems used in our paper.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide the pseudocode for our algorithm in Appendix A, and provide the details about experimental settings and hyperparameter of our model in Appendix F. The dataset generation method is provided in Appendix F.4.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We refer to Section 4 for code repo.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide details about our experiments and settings in Appendix F

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We provide the stand deviation of average time cost and average iteration count in Appendix H.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide the details about experiments compute recources in Appendix F.2. Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: Contents in our paper focus on accelerating the linear system solving, which could have potential help to the industry or other field involving linear system solving.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: Our data and model focus on linear system solving, especially PDE problems, which have no high risk for misuse.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We have checked this and confirmed the paper poses no such risks.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We provide the instruction in supplemental material about our code and data. Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: SUMMARY:
The authors use a neural operator approach to generate subspaces that are used for the acceleration of Krylov subspace methods for several partial differential equations setups.

STRENGTHS:
The authors are interested in a very relevant problem of computational science and engineering of solving linear systems of equations. The attempt of providing a convergence analysis for the proposed method.

WEAKNESSES:
The details of the method remain hidden to me and it seems difficult to grasp the details of the method and many details and even the major workings of the method are unclear.

QUESTIONS:
Why the smallest eigenvalues in line 163 on page 5?
I assume that Y_k in Algorithm 1 is the predicted subspace from the neural network, but then explain the derivation used for equation (11)? How is this motivated?
How is the method of the authors combined with preconditioning? The details are not given but the choice of the preconditioner is crucial and how would this depend on the learned subspace the method provides?

LIMITATIONS:
The method is limited to linear problems but this is a very broad and general class. If the method was better explained and the details clear and convincing there would be no limitations for these kinds of problems.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-06-24

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe manuscript introduces a novel method, referred to as NeurKItt, which combines neural network techniques with Krylov subspace methods to accelerate the solution of linear systems derived from partial differential equations (PDEs). The core innovation of NeurKItt lies in its capability to predict invariant subspaces associated with the matrices defining these linear systems, thereby significantly reducing the number of iterations required for convergence.\n\nSTRENGTHS:\n1. **Originality:**\nThe manuscript introduces a novel approach by integrating neural network techniques with Krylov subspace methods to predict invariant subspaces of linear systems, significantly enhancing the efficiency of these traditional methods.\n\n2. **Quality:**\nThe manuscript effectively targets the critical issues of computational inefficiency and instability in the context of high-dimensional and poorly conditioned linear systems. The choice of problem is highly relevant to both academic research and practical applications in scientific computing, making the work valuable to a wide audience.\n\n3. **Clarity:** \nThe manuscript is well-structured, presenting complex concepts in an understandable manner. It effectively communicates the challenges of existing methods and how NeurKItt addresses these issues, providing clear explanations and logical progression from problem statement to solution. However, improvements in typographical accuracy and some clarifications in methodological descriptions could further enhance clarity.\n\n4. **Significance:**\nThe significance of this work lies in its potential impact on the efficiency of solving large-scale linear systems, which are crucial in various scientific and engineering applications. By reducing computational time and resource consumption, NeurKItt could significantly benefit fields reliant on large-scale computations, making this contribution highly relevant to both academic research and industry applications.\n\nWEAKNESSES:\n1. The sentence \"Research in using neural networks for accelerating linear system solving\" on line 75 appears abruptly terminated and lacks a verb or continuation that would complete the thought.\n\n2. On line 81, the term \"precondition\" is likely a typographical error. The correct term, given the context, could be \"preconditioning\".\n\n3. There appears to be a typographical error in Equation (8), where the variable $z$ should presumably be $y$.\n\n4. Equation (13) on page 7 starting with an unnecessary 's' before the equal sign seems to be a typo.\n\n5. In the manuscript, Section 4.1 discusses the use of Fourier Neural Operators (FNO) for parametric PDE problems and introduces the subsequent subspace prediction concept. It employs a subspace $S$ defined as the span of vectors $\\{s_1, s_2, \\cdots, s_n\\}$. The manuscript mentions that $s_i$ is the eigenvector corresponding to the $i$-th smallest eigenvalue, but it does not specify whether these eigenvalues are of the matrix  or another related matrix. This omission could lead to confusion about the origin and relevance of these eigenvectors, particularly in how they are integrated into the model and influence the outcome of the subspace prediction. \n\n6. In Section 4.1, the manuscript proposes using invariant subspaces associated with the smallest eigenvalues to train the subspace prediction module. However, an oversight in the theoretical analysis (Section 5.2) arises from applying this concept to potentially non-Hermitian matrices. The section (section 5.2) is based on Hermitian positive definite matrices, whose eigenvalues are all real numbers that can be orderly compared. Non-Hermitian matrices, on the other hand, may exhibit complex eigenvalues, complicating or nullifying the notion of \"smallest\" eigenvalues due to their lack of a natural order. This discrepancy raises concerns about the applicability of the method to non-Hermitian matrices, which are prevalent in many practical applications.\n\nQUESTIONS:\nIn reviewing the theoretical analysis provided between lines 237 and 248, I observed a notable similarity in both the phrasing and the mathematical formulations with those presented in reference [23]. Could the authors clarify the extent of originality in these sections? It is crucial for academic integrity to distinguish clearly between direct quotes and original analysis.\n\nLIMITATIONS:\nThe authors discussed limitations in Section 7 Limitation and Conclusions.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a data-driven approach to accelerate solving linear systems. Linear Systems are widespread in scientific computing applications like solving partial differential equations (PDEs), nonlinear systems, etc. so this method has potential for major impact. The proposed method builds upon the idea of neural operators that learn mapping between function spaces.\n\nWhile most prior methods have attempted to accelerate solving linear systems by predicting a better initial guess, this work instead predicts the matrix invariant subspace. It uses this subspace to accelerate the Krylov Iterations. The method is validated on linear systems originating from PDEs and achieves speedups of around $5.5\\times$ in computation time.\n\nSTRENGTHS:\n1. The presented method is novel, and to the best of my knowledge there doesn't exist a method that uses a similar approach\n2. Experiments are well designed and use strong baseline methods (GMRES from PETSc). Speedup in computation time and iteration count over such a strong baseline presents a strong case for the proposed method.\n\nWEAKNESSES:\n1. The Neural Operators need to be trained on each individual problem. So while there is an improved convergence speed, there is an increased overall time when the training time is taken into account. However, this is a pretty common shortcoming across neural operator approaches and hence I haven't used this shortcoming as part of my overall scoring.\n\nQUESTIONS:\n1. Are there any cases where GMRES (with no-preconditioners) fails to converge while NeurKItt can solve it because it is data-driven?\n2. Is it necessary to learn the NO for each system? Or is it possible to learn a single network and reuse it?\n3. (Maybe I missed this) Is it possible to include the training times for the neural networks?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper develops a new method, NeurKItt, for accelerating the solution of non-symmetric linear systems. NeurKItt constructs an approximate invariant subspace of non-symmetric matrix A using the Fourier Neural Operator. This invariant subspace is then used as the initial subspace within GMRES. The paper provides theoretical support for the proposed method and empirical results showing the benefits of NeurKItt.\n\nSTRENGTHS:\n1.) The paper develops an interesting new method for solving non-symmetric linear systems based on operator learning. In particular, I have not seen the Fourier Neural Operator used in this way before. I believe the idea has the potential to be quite useful and could have valuable implications for other problems.\n\n2.)  The authors provide theoretical analysis supporting the method.\n\n3.) The method performs well in practice, which is what is most important. NeurKItt seems to yield significant reductioning in the number of GMRES iterations, and also yields good wall-clock time speed-ups (which is particularly impressive as these matrices are very sparse). So I think the method has the potential to be very useful in practice. \n\nOverall, the paper proposes an exciting new idea that appears to work quite well and has the potential to be quite useful.\nMoreover, the paper addresses an important hard problem: improving the solutions of non-symmetric linear systems. This setting is quite challenging, as the non-symmetry makes things difficult. Often, methods for accelerating GMRES are highly problem-dependent, so a method that can lead to generic improvements is significant. I'm happy to recommend its acceptance to NeurIPS provided an issue given below in the weaknesses section is properly addressed.\n\nWEAKNESSES:\nI'd say the main weakness of the paper is its presentation. In particular, a significant issue is that the paper lacks a precise description of how FNO is implemented to learn the invariant subspace. The appendix only briefly describes FNO but does not explain how the paper applies it. The current discussion in lines 148-151 is vague and unclear.  In particular, the output of FNO is the output of an operator on a function, not a subspace. So, how are you getting the predictive subspace from the output of the FNO? I have ideas of how this is done, but rather than guessing, I would like the authors to explain the precise procedure adequately. This should then be added to the paper, ideally in the main body, but at least the supplement with clear pointers in the main paper of where to find the details. \n\nI'm willing to consider raising my score if the authors address this point well. If they don't, I will lower my score, as I cannot recommend acceptance of a paper for which a significant part of the methodology is unclear.  \n\nAside from this, the authors should do a careful spell-check, as there are many typos throughout, which is somewhat distracting.\n\nQUESTIONS:\nSee weaknesses.\n\nLIMITATIONS:\nThe authors have adequately addressed the current limitations of the method.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-13"}