Neural Krylov Iteration for Accelerating Linear System Solving

We thank the reviewer for the patience to read through our paper. We are pleased to re-introduce our work and respond to your comments as follows. We sincerely hope that our rebuttal could properly address your concerns. If so, we would deeply appreciate it if you could raise your score and your confidence. If not, please let us know your further concerns, and we will continue actively responding to your comments and improving our work.

---

Major Workings of Our Method

• We propose a novel method, namely NeurKItt, for accelerating linear systems solving. NeurKItt consists of two modules, one is the subspace prediction module, while the other is the acceleration module.
• The subspace prediction module employs the neural operator for predicting the invariant subspace $\hat{\mathcal{K}}$ of given linear systems $Ax=b$, as a matrix $X$ that contains eigenvectors corresponding to the smallest eigenvalues from $A$.
• The acceleration module uses the predicted subspace for accelerating the linear system solving. It is the Krylov subspace algorithm with predicted subspace involved for iteration.

---

Subspace Prediction Module

• We use FNO to predict the subspace of the matrix $A$ from the given linear system $Ax=b$. We give the introduction to FNO in Appendix C.
• The reason why we select the neural operator for subspace prediction is that the mapping from the matrix $A$ of linear system $Ax=b$ to its corresponding invariant subspace $\mathcal{\hat{K}}$ can be considered an operator, i.e., a mapping between two Hilbert spaces, thus solving such problems involves finding the corresponding operator.
• We employ the projection loss function for training the FNO and apply thin QR decomposition to the output of the FNO. Projection loss is designed for learning the subspace. The idea behind it is that the predicted subspace should contain the basis vectors in the target subspace. We employ the thin QR decomposition for orthogonalizing the FNO output. We provide the details about thin QR decomposition in Appendix E.

Why the smallest eigenvalues in line 163 on page 5?

• The FNO learns the subspace $\mathcal{S}$ that associates with the $n$ smallest eigenvalues, assuming $\mathcal{S}=\mathrm{span}\{ s_1, s_2, \ldots, s_n \}$, where $s_i$ is the eigenvector corresponding to the $i$-th smallest eigenvalue of the given matrix $A$, $i = 1, 2, \ldots, n$. We think such subspace $\mathcal{S}$ is easy to learn because Theorem 5.2 gives proof that if changes in a matrix $A$ 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. Theorem 5.2 can be found in Section 5.

---

Acceleration Module

• The acceleration module is the Krylov iteration algorithm with predicted subspace involved, inspired by Krylov recycling algorithms. The detailed pseudo-code about the acceleration module can be found in Appendix A.2. The key of the acceleration module is the Arnoldi relation $(I−C_kC^H_k )AV_{m−k} = V_{m−k+1}\underline{H}_{m−k}$, which takes the subspace into iteration.
• Here we give an intuitive idea about the acceleration module. Krylov subspace iteration solves the linear system problem by iteratively approximating the invariant subspace of matrix $A$ from the linear system $Ax=b$. When the subspace information about $A$ is given, we don't need to start the approximation from scratch, thus the acceleration module can improve the convergence speed.

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?

• $Y_k$ is the matrix of predicted subspace $\mathcal{\hat{K}}$ from the neural operator. In particular, we obtain $C_k$ and $U_k$ from the thin QR decomposition on $AY_k$, such that $AU_k=C_k$ and $C_k^{H}C_k=I_k$.

• Equation $C_k^{H}C_k=I_k$ allows us to find the basis of $\mathcal{\hat{K}}$, which enables us to remove the subspace $\mathcal{\hat{K}}$ in Equation (11) $(I−C_kC^H_k)AV_{m−k} = V_{m−k+1}\underline{H}_{m−k}$ by $(I-C_kC_k^H)$. This approach makes the Krylov subspace iteration more efficient because we only need to search the rest of the invariant subspace for matrix $A$.

• Equation $AU_k=C_k$ makes sure the matrix $\underline{H}_{m-k}$ is an upper Hessenberg matrix, which speedups the solving of the linear system in the later iteration procedure.

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?

• Supposing the original Arnoldi relation $AV=VH_0$ and the preconditioner matrix $P$, the preconditioned Arnoldi relation is $PAV=VH_1$. So, for NeurKItt, after the combination, the Arnoldi relation is $P(I-C_kC_k^H)AV=VH_2$. Here, $H_i$ is the upper Hessenberg matrices.
• In our experiments, the learned subspace for any preconditioning method is the subspace composed of eigenvectors corresponding to the smallest eigenvalues of the given matrix $A$.

---

We also provide a theoretical analysis of the acceleration module in Section 5.1, which gives the convergence analysis and explains our choice to predict the subspace rather than the solution. In Section 5.2, we further justify our choice of using the subspace composed of eigenvectors corresponding to the smallest eigenvalues for prediction.

If you are confused about any of the content in our paper, like what kind of obstacle hinders the understanding of our work, please let us know your further concerns, and we will continue to actively respond to your comments.

We thank the reviewer for the insightful and valuable comments. We respond to your comments as follows and sincerely hope that our rebuttal could properly address your concerns. If so, we would deeply appreciate it if you could raise your score and your confidence. If not, please let us know your further concerns, and we will continue actively responding to your comments and improving our work.

Weaknesses

Typos mentioned in Weakness 1-4

• Thank you for pointing out these typos. We will fix them in the future version.

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.

• Given a linear system $Ax=b$, these eigenvalues are derived from the same matrix $A$. Thank you for pointing out this potentially misleading statement. We will fix it in the future version. In particular, we will replace the sentence in lines 164-165 "where $s_i$ is the eigenvector corresponding to the $i$-th smallest eigenvalue, $i = 1, 2, \dots, n$." with "where $s_i$ is the eigenvector corresponding to the $i$-th smallest eigenvalue of the given matrix $A$, $i = 1, 2, \ldots, n$."

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.

• Thank you for pointing out this problem. Analyzing perturbations for non-Hermitian matrices involves pseudo-spectral methods, which makes it difficult to investigate. To simplify the analysis, Section 5.2 explores how to select an appropriate subspace in the Hermitian positive definite case. In practice, for Hermitian positive definite matrices, the smallest eigenvalues are compared following the natural order. But for non-Hermitian matrices, the eigenvalues will be sorted by comparing their modulus, such that: $|\lambda_1|\le|\lambda_2|\le|\lambda_3|\le|\lambda_4|\le\dots\le|\lambda_n|$ where $\lambda_i\in\mathbb{C}$, $i=1,2,\dots,n$, is the eigenvalue of a given non-Hermitian matrix $A$. We will add these details to Section 5.2 to avoid the discrepancy.

Questions

In 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.

• We are sorry for the misunderstanding caused by our improper citation. We fully adhere to academic ethics, and this issue was due to a writing oversight. Sentences in lines 237-248 provide the definitions and assumptions serving as the preliminaries for Theorem 5.2, which is mentioned in Theorem 5.2. To keep Theorem 5.2, which is cited properly, coherent with the original text, we reused the definitions and assumptions from reference [23], and this might potentially lead to disputes.
• Due to NeurIPS rebuttal restrictions, we are not allowed to submit a revised paper during rebuttal. We will make the quotes clear by adding the sentence " The following definitions and assumptions for Theorem 5.2 are from the reference [23]." at line 238 to indicate that content in lines 237-248 comes from the reference [23].

Dear Reviewer iSRc,

We are writing as the authors of the paper titled "Neural Krylov Iteration for Accelerating Linear System Solving" (ID: 18161). Thanks again for your valuable comments and constructive suggestions, which are of great help to improve the quality of our work. We have followed your suggestions to significantly enhance the quality of our work. As the deadline for the author-reviewer discussion period is approaching (due on Aug 13), we are looking forward to your further comments and/or questions.

We sincerely hope that our rebuttal has properly addressed your concerns. If so, we would deeply appreciate it if you could raise your scores. If not, please let us know your further concerns, and we will continue actively responding to your comments and improving our work.

Best regards,

Authors

We thank the reviewer for the insightful and valuable comments. We respond to your comments as follows and sincerely hope that our rebuttal could properly address your concerns. If so, we would deeply appreciate it if you could raise your score and your confidence. If not, please let us know your further concerns, and we will continue actively responding to your comments and improving our work.

Weaknesses

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.

• We agree with the reviewer that training the neural network involves additional costs, and it is a shortcoming of any neural operator approach. But considering the benefits of acceleration over millions of linear system solving, such training costs could be negligible due to the time saved. This is one of the reasons why we claim the feasibility of NeurKItt in practice.

Questions

Are there any cases where GMRES (with no-preconditioners) fails to converge while NeurKItt can solve it because it is data-driven?

• There might be the case where GMRES fails but NeurKItt can solve.
• Specifically, NeurKItt consists of the subspace prediction module and acceleration module. The subspace prediction module is data-driven, while the acceleration module uses the predicted subspace to improve the Krylov iteration convergence speed. The case mentioned is due to the acceleration module that replaces the original Arnoldi relation with $(I-C_kC_k^H)AV_{m-k}=V_{m-k+1}\underline{H}_{m-k}$, which takes the invariant subspace into iteration. It can also be considered as a preconditioned linear system solving, which replaces the original linear system problem $Ax=b$ with $PAx=Pb$, where $P=(I-C_kC_k^H)$. This change directly improves the convergence speed while lowering the condition number of the given matrix $A$, which leads to the case that GMRES fails to but NeurKItt can solve. The data-driven module (Subspace Prediction Module) makes the improvement possible, but it's not the direct answer.

Is it necessary to learn the NO for each system? Or is it possible to learn a single network and reuse it?

• We'd like to give answers given different contexts.
• First, given the case that systems derived from different PDEs, we did not conduct relevant experiments. But we think it is possible to learn a single network and reuse it in the given context, while the network is pre-trained. Recent works like $1$ , $2$ , and $3$ , give experimental proofs that the neural operator can be a pre-train model, which can be reused for different PDEs with only fine-tuning.
• Second, given the case that systems derived from the same PDE but with different scales, like different sizes of the linear systems, we have conducted an experiment to show the potential of NeurKItt in Appendix I. In particular, we train the neural operator on linear systems with a fixed matrix size but predict linear systems with larger sizes, then use the predicted subspace to accelerate the solving for those linear systems with larger sizes. The experimental results show that NeurKItt successfully accelerates the solving, which indicates the potential to learn a single network and reuse it for the linear systems derived from the same PDE but with different scales.

(Maybe I missed this) Is it possible to include the training times for the neural networks?

• Yes, we have reported the training time for each dataset in Appendix H.2. In our experiments, all the training converges within 150 minutes, and keeping training has no additional performance improvement once reaching the convergence.

[1] Yang, Liu, et al. "In-context operator learning with data prompts for differential equation problems." Proceedings of the National Academy of Sciences 120.39 (2023): e2310142120.

[2] Hao, Zhongkai, et al. "DPOT: Auto-regressive denoising operator transformer for large-scale pde pre-training." arXiv preprint arXiv:2403.03542 (2024).

[3] Zhou, Anthony, et al. "Strategies for Pretraining Neural Operators." arXiv preprint arXiv:2406.08473 (2024).

Dear Reviewer XS9H,

We are writing as the authors of the paper titled "Neural Krylov Iteration for Accelerating Linear System Solving" (ID: 18161). Thanks again for your valuable comments and constructive suggestions, which are of great help to improve the quality of our work. We have followed your suggestions to significantly enhance the quality of our work. As the deadline for the author-reviewer discussion period is approaching (due on Aug 13), we are looking forward to your further comments and/or questions.

We sincerely hope that our rebuttal has properly addressed your concerns. If so, we would deeply appreciate it if you could raise your scores. If not, please let us know your further concerns, and we will continue actively responding to your comments and improving our work.

Best regards,

Authors

We thank the reviewer for the insightful and valuable comments. We respond to your comments as follows and sincerely hope that our rebuttal could properly address your concerns. If so, we would deeply appreciate it if you could raise your score. If not, please let us know your further concerns, and we will continue actively responding to your comments and improving our work.

Weaknesses

How are you getting the predictive subspace from 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.

• Here, we use the 2D Darcy flow problem to demonstrate how the Fourier Neural Operator (FNO) predicts the subspace given the input function.

  • Let the input function be $a\in\mathbb{R}^{d_a\times d_a}$ from the 2D Darcy flow problem, where $d_a\in\mathbb{N}$ is the resolution of the input function, which yields a linear system $Ax=b$ for numerical solving, where matrix $A\in\mathbb{R}^{d_A\times d_A}$. We aim to predict a subspace $\mathcal{\hat{K}}$, as a matrix $X\in\mathbb{C}^{d_A\times n}$, for matrix $A$ in the linear system. First, the FNO's lifting layer $\mathcal{R}$ transforms $a$ to a higher-dimensional representation $v_0\in\mathbb{C}^{d_a\times d_a\times c}$, while the 3rd dimension is channel dimension. After $T$ Fourier layers forward, we have the $v_T\in\mathbb{C}^{d_a\times d_a\times c}$. Because Fourier layers keep the shape unchanged, we apply a transformation $Q$ to map the $v_T$ to the desired space, $\mathcal{\hat{K}}=Q(v_T)$, with $Q:\mathbb{C}^{d_a\times d_a\times c}\rightarrow\mathbb{C}^{d_A\times 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\in\mathbb{C}^{d_a^2\times c}$. Then a fully-connected neural network (FNN) applies the mapping $\mathbb{C}^{d_a^2\times c}\rightarrow\mathbb{C}^{d_a^2\times n}$ to $q_0$, obtaining $q_1\in\mathbb{C}^{d_a^2\times n}$. And another FNN applies the mapping $\mathbb{C}^{d_a^2\times n}\rightarrow\mathbb{C}^{d_A\times n}$ to $q_1$, obtaining the output $X\in\mathbb{C}^{d_A\times n}$. Finally, we apply QR decomposition to $X$ for orthogonalizing, obtaining $\mathcal{\hat{K}}=\mathrm{span}\{X\}$.

• To address the problems identified in the reviewers' comments, we will make several changes to our paper. First, we will replace lines 148-151 with more precise content to clarify the subspace prediction task. The revisions include the problem setup and the detailed description of how to predict subspace by FNO. Second, we will add the 2D Darcy flow example mentioned above to Appendix C to give an image of how to predict the subspace by the FNO. The sentences quoted below are the revised version of lines 148-151:

  • "Generally, for a linear system $Ax=b$ derived from the parametric PDE problem, to predict its invariant subspace $\mathcal{\hat{K}}$, the input to FNO is the input function $a\in\mathbb{R}^{d_a}$ from the given PDE, where $d_a\in\mathbb{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 space $\mathcal{G}:\mathbb{R}^{d_a}\rightarrow \mathbb{C}^{d\times n}$ using FNO.

  • For FNO, the lifting transformation $\mathcal{R}$ first lifts the input $a$ to a higher dimensional representation $v_0\in\mathbb{C}^{d_a\times 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\in\mathbb{C}^{d_a\times n}$ 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:\mathbb{C}^{d_a\times c}\rightarrow\mathbb{C}^{d_A\times n}$. NeurKItt then uses QR decomposition to orthogonalize the matrix $X$, obtaining the predicted subspace $\mathcal{\hat{K}}$. We provide more details about how to predict the subspace given the 2D Darcy flow problem in Appendix C."

• Due to NeurIPS rebuttal restrictions, we cannot submit a revised version during the rebuttal. We will incorporate these modifications into our paper in the future version. If these modifications do not adequately address the issues raised, please let us know your further concerns, and we will continue to actively respond to your comments.

Aside from this, the authors should do a careful spell-check, as there are many typos throughtout, which is somewhat distracting.

• Thank you for bringing this to our attention. We will conduct a thorough spell-check and correct any typos in the future version.