Efficient Linear System Solver with Transformers

Authors: max  vladymyrov
johannes  von oswald
nolan  miller
mark  sandler
Venue: ICML_cc_2024_Workshop_AI4MATH
Published: 2024-06-01 09:13:02

Abstract: This paper investigates the potential of linear Transformers as solvers for systems of linear equations. We propose a novel approach where the Transformer encodes each equation as a separate token, allowing the model to process the system in a permutation-invariant manner. To enhance generalizability and reduce the parameter count, we introduce a block-wise re-parameterization technique for the attention weight matrices. This technique decouples the problem dimension from the model's parameter count, enabling the Transformer to effectively handle systems of varying sizes. Our experiments demonstrate the Transformer's competitive performance compared to established classical methods such as Conjugate Gradient, especially for systems with smaller sizes. We further explore the model's ability to extrapolate to larger systems, providing evidence for its potential as a versatile and efficient solver for linear equations.

Main paper: 
 introduction: 
 Solving linear systems of equations is a fundamental problem in numerous fields, including scientific computing, machine learning, and engineering. While traditional methods like Gaussian elimination and iterative solvers like Conjugate Gradient (Hestenes et al., 1952) are widely used, exploring alternative approaches holds the potential for more efficient and versatile solutions.Transformers, originally developed for natural language processing tasks (Vaswani et al., 2017) , have demonstrated a remarkable ability to capture complex relationships within sequential data. This ability extends beyond natural language processing, as evidenced by their successful application in solving a variety of problems, ranging from noisy linear regression and classification (Garg et al., 2022) to the traveling salesmen problem (Yang et al., 2023) and other domains (Mirchandani et al., 2023) . Transformers have shown promise in various scientific computing tasks beyond natural language processing (Li et al., 2020) .In this paper, we investigate the ability of Transformers to solve linear systems of equations of the form Ax = b, where A is a positive definite and symmetric matrix. This approach presents several compelling advantages:• Potential for Parallelism: Transformers' inherent parallelism in processing tokens could translate to efficient solvers, particularly for large-scale systems.• Extrapolation Capabilities: We explore the potential of Transformers to generalize beyond the size they are trained on, offering a more flexible tool.We propose a novel tokenization scheme that encodes each equation as a separate token, allowing the Transformer to process the system in a permutation-invariant manner. We demonstrate that even linear Transformers can solve linear systems of equations with high accuracy.Since the attention does not have memory, the whole linear system must be provided in-context. A naive encoding scheme would result in the number of model parameters scaling quadratically with the size of the linear system. To address this, we introduce a novel re-parameterization technique for the attention weight matrices that decouples the problem dimension from the model's parameter count. We empirically demonstrate that this re-parametrization does not compromise accuracy substantially while significantly reducing the number of parameters, requiring only dozens of parameters per layer. This approach enables the Transformer to effectively handle systems of varying sizes.Our experiments demonstrate that the proposed Transformer-based solver achieves comparable accuracy to 6-8 iterations of the Conjugate Gradient method, while demonstrating superior speed for systems with a small number of equations. Furthermore, we investigate the model's ability to extrapolate to larger systems, showcasing its potential as a versatile and efficient solver for linear equations. This research contributes to the understanding of Transformers' capabilities beyond traditional sequence processing tasks and opens avenues for further exploration of their applications in numerical problems. 
preliminaries: 
 In this section we introduce notations for linear Transformers, data representation, and the specific type of problem we consider.Given an input sequence e 1 , e 2 , ..., e N ∈ R D , a single head in a linear self-attention layer is typically parameterized by four matrices, key W K , query W Q , value W V and projection W P . The output of the non-causal layer at position i is e i + ∆e i where ∆e i is computed asEQUATIONEquivalently, we can use the parameters P = W P W V and Q = W K W Q , resulting in the equation:EQUATIONIn the case of multiple heads (P 1 , Q 1 ), (P 2 , Q 2 ), ..., (P h , Q h ), the effect is simply the summation of all heads:∆e i = H k=1 N j=1 (e j Q k e i )P k e j .(3)We define a linear Transformer as a multi-layer neural network composed of L linear self-attention layers parameterized by θ = {Q l k , P l k } k=1...H,l=1...L . To isolate the core mechanisms, we consider a simplified decoder-only architecture, excluding MLPs and LayerNorm components. This architecture has been used in previous work (von Oswald et al., 2023; Ahn et al., 2023; Vladymyrov et al., 2024) . For clarity, we focus on the single-head attention case and drop the index k.Our goal is to find a vector x ∈ R N that solves the system of N linear equations a i , x = b i for each i = 1 . . . N with a i ∈ R N , and b i ∈ R. In matrix form, this problem can be written as Ax = b, where A ∈ R N ×N and b ∈ R N .We focus on the case where A is a positive definite symmetric matrix with a fixed condition number κ. The condition number significantly impacts the convergence of iterative methods (Trefethen & Bau, 2022) . In this paper, we present preliminary results with κ = 5.x Linear Transformer(a 0 , b 0 , h 0 ) • • • (a N , b N , h N ) x 0 Figure 1.A Transformer receives an input embedding where each token corresponds to the parameters (ai, , bi) of a single equation within a linear system. A learned or predefined embedding hi can be added for improved performance. The initial solution estimate x0 is fed in as the final query token, and the final solution x is obtained after passing through the Transformer layers.The Transformer processes input data as a sequence of tokens. Each equation in the linear system is represented as a distinct token e i = (a i , b i , h i ) ∈ R D , where D = N + 1 + K and h i ∈ R K are optional embedding vectors. These embeddings can be either pre-defined or learned during model training. We can represent them collectively in a matrix form as H ∈ R N ×K . This tokenization scheme ensures that the model remains invariant to the order in which the equations are presented.Additionally, we append a query token e N +1 = (x 0 , 1 1+K ) to the sequence, where x 0 ∈ R d is an input placeholder from which the output will be read, initialized with all zeros. We constrain the attention mechanism to focus solely on the first N tokens of the sequence, ignoring the query token.We use superscript notation (e.g. a l i , b l i ) to denote the corresponding component of the i-th token in the Transformer's output at layer l. The initial layer corresponds to the input:(a 0 i , b 0 i ) = (a i , b i ).For a model with parameters θ, the prediction is obtained as a transformation of the x 0 component of the final token in the last layer:f θ ({e 1 , ..., e N }, e N +1 ) = x L N +1 .To train the linear Transformer using the tokens, we minimize the following mean squared error loss function using batches of different generated problems {A, b, x}:L(θ) = E A,b (f θ ({e 1 , ..., e N }, e N +1 ) -x) 2 . (4) 
re-parameterization of weight matrix: 
 Given that each token consists of three components A, b and H, we can equivalently rewrite the attention weight matrices by expanding the interaction between each component:  , Q =   Q A,A Q A,b Q A,H Q b,A Q b,b Q b,H Q H,A Q H,b Q H,H   .Assuming for simplicity that the dimension K of vectors h i is a multiple of N (i.e. K = lN ), H can be decomposed into l blocks H l ∈ R N ×N . Consequently, the components of P and Q controlling the interaction between A and all the H l will have a size N × N . Components handling the interaction with b will have sizes N × 1, 1 × N or 1 × 1.We propose the following block-wise re-parametrization the attention weights matrices. Each rectangular matrix is represented as µ1 N ×1 or µ1 N ×1 , while each square matrix is represented as µ 1 I N ×N + µ 2 1 N ×N . Here I is the identity matrix, and 1 is a matrix where all elements are ones.Our approach draws inspiration from low-rank approximations in numerical linear algebra (Halko et al., 2011) and is motivated by several factors: • Efficiency. utilizing a single scalar for each block significantly reduces per-block computation, enhancing the overall speed of the algorithm. Figure 3 . Runtime comparison of different linear system solvers on an A100 GPU, using batches of 1 000 systems with varying matrix sizes.. AA T represents the runtime for a batch of 1 000 matrix multiplications of matrix A provided for reference. BLOCK LINEAR TR is our proposed algorithm with 3 layers of linear Transformer using BLOCK encoding. EXACT (SYM) and EXACT (POS) are computed using jax.scipy.sparse.linalg.cg() solver assuming symmetric of positive definite matrix structure for A, respectively. CG is computed using jax.scipy.sparse.linalg.cg() solver for a given number of iterations.•• Generalization. This method decouples the problem dimension N from the number of model parameters. This allows us to apply the learned model to inputs of different sizes, not just the size it was trained on. We can also fine-tune model trained on one size for another.We refer to this new re-parameterization as BLOCK embedding, as opposed to FULL embedding where the full weight matrices are learned. Figure 2 shows the learned weights of the trained FULL and BLOCK embeddings for a 3-layer linear Transformer.  
experiments: 
 Figure 3 compares the runtime required to compute the solution using different solvers on A100 GPU. Our proposed method (BLOCK LINEAR TR) requires only a few matrix multiplications to compute a forward pass through 3 layers of the linear Transformer. While the Conjugate Gradient (CG) method theoretically scales as O(N 2 ), which is asymptotically faster than matrix multiplication's O(N 2.8 ), the constant overhead associated with CG becomes significant for the small values of N considered here. Moreover, matrix multiplication is a more native operation for GPUs compared to the CG, which involves dot products and matrix-vector products. In this context, computing 3 layers of linear attention proves to be considerably faster than competing methods.Figure 4 compares the accuracy of our solver. We trained several variants of the 3-layer linear Transformer:• FULL, N = 9. This variant trains full weight matrices on problems with size N = 9 only. Due to the fixed 9 × 9 weight dimensions, this model cannot generalize to matrices of other sizes. However, it achieves the best performance we observed for problems of size N = 9 using linear Transformers.• BLOCK, N = 9. This variant trains weights using BLOCK encoding with 3 N ×N identity matrices as extra embeddings, also on N = 9 problems only. While the results can be applicable to problems of other sizes, the generalization qualiry is limited. Compared to the Conjugate Gradient baseline, our algorithm achieves an accuracy roughly equivalent to 6-8 iterations of CG, depending on the input size. 
conclusions: 
 This paper explored the novel application of linear Transformers for efficiently solving small systems of linear equations with symmetric and positive definite coefficient matrices. We demonstrated that by encoding each equation as a distinct token and implementing a block-wise reparameterization technique, Transformers could achieve accuracy comparable to 6-8 iterations of Conjugate Gradient, while being faster for small problem sizes. Our model exhibited the ability to generalize beyond its training data, effectively handling systems of varying sizes and unseen condition numbers.This research opens up exciting possibilities for utilizing Transformers in numerical tasks. Further investigation into architectural choices, training strategies, and the boundaries of generalization could lead to the development of even more efficient and adaptable solvers. This work reinforces the notion that Transformers, initially designed for natural language processing, hold remarkable potential as powerful tools across various scientific computing domains.While our approach shows promise, it has several limitations. Currently, we only handle positive definite symmetric matrices, which restricts the applicability of our method. Future work should explore extending this approach to general matrices, including non-symmetric and indefinite cases. Additionally, our method's performance on very large systems or highly ill-conditioned matrices needs further investigation.Scaling the approach to handle sparse matrices efficiently is another important direction for future research. Integrating our Transformer-based solver with classical methods, potentially as a preconditioner or in a hybrid algorithm, could leverage the strengths of both approaches and is an exciting area for further study. 
 
 Appendix: Google Research. Correspondence to: Max Vladymyrov <mxv@google.com>.The first AI for MATH Workshop at the 41st International Conference on Machine Learning, Vienna, Austria.Copyright 2024 by the author(s). 


Decision: This paper proposes a linear Transformer for solving linear systems equations. This idea and the proposed re-parameterization technique are novel according to Reviewers 6KUh and dVvN. The authors could further improve this paper by including a deeper theoretical analysis (Reviewer 6KUh), large-scale experiments (Reviewers 6KUh, dVvN), additional baseline comparisons (Reviewers tFdu, dVvN), additional evaluation metrics (Reviewer dVvN), and polish the writing (Reviewers tFdu, dVvN).

Reviews: review: paper_summary: paper_summary: The paper presents an innovative approach to solving systems of linear equations using linear Transformers. The authors propose encoding each equation as a separate token, leveraging the permutation-invariant nature of self-attention mechanisms. They introduce a novel re-parameterization technique that decouples the problem dimension from the model’s parameter count, enabling efficient handling of systems of varying sizes. Experimental results demonstrate competitive performance compared to classical methods such as Conjugate Gradient, particularly for smaller systems, and the potential for extrapolation to larger systems.
strength_weakness: strength_weakness: 1. The paper focuses on solving linear systems with positive definite and symmetric matrices, which is a specific and somewhat limited scope.
2. The paper focuses on empirical results, but a deeper theoretical analysis of the proposed method's computational complexity and convergence properties would add significant value.
3. The paper demonstrates the method's efficiency for smaller systems, but its performance and scalability for very large systems remain unclear. Additional experiments with larger systems would strengthen the claims.
questions: questions: 1. Why choose Linear Transformer over Standard Transformer?
paper_summary: paper_summary: This paper studies the application of linear Transformers for solving a small systems of linear equations with symmetric and positive definite coefficient matrices. To reduce the number of parameters, a novel re-parameterization technique for the attention weight matrices is introduced. While this is an interesting research, its applicability in practice is questionable.
strength_weakness: strength_weakness: 1. The literature review is very limited. In AI/ML literature, there are several recent works on applying Transformer for solving/studying dynamical systems and linear systems. For instance:

@article{geneva2022transformers,
  title={Transformers for modeling physical systems},
  author={Geneva, Nicholas and Zabaras, Nicholas},
  journal={Neural Networks},
  volume={146},
  pages={272--289},
  year={2022},
  publisher={Elsevier}
}

@article{charton2021linear,
  title={Linear algebra with transformers},
  author={Charton, Fran{\c{c}}ois},
  journal={arXiv preprint arXiv:2112.01898},
  year={2021}
}

@incollection{arslan2021machine,
  title={Machine Learning Algorithms for Solving Linear Systems of Equations},
  author={Arslan, Hilal and Bozyigit, Fatma},
  booktitle={Developing Mathematical Literacy in the Context of the Fourth Industrial Revolution},
  pages={155--184},
  year={2021},
  publisher={IGI Global}
}  
Also, there are also works on other machine learning methods to solve linear system in the AI/ML literature. However, the author did not review/compare these works in the same field/direction. 

2. The representation of this paper can be improved. For instance, line 109: 'We can this new re-parametrization BLOCK embedding...'
3. Other machine learning-based approach for solving linear system should be included in addition to conjugate gradient method for a better comparison. 
4. The applicability of the proposed approach is limited:
1). How do the end-users control the trade-off between precision and computation time? 
2). The authors mentioned that "Although the Conjugate Gradient (CG) solution scales as O(N2), which is asymptotically faster than matrix multiplication’s O(N2.8), the constant overhead of CG becomes significant for the small N that we consider here." What would be a good N for applying proposed methods in practice?
questions: questions: More discussion on figure 2 should be added. In a given layer, the structures in top and bottom are quite different. What are the main points that the authors are trying to state with figure 2?
paper_summary: paper_summary: This paper articulates a novel linear transformer architecture to solve linear systems. Equations are encoded as permutation-invariant tokens and re-parameterized to decouple the problem dimension from the model’s parameter count. Experiments show that the proposed method has a competitive performance against the SOTA linear solvers.
strength_weakness: strength_weakness: - Experiments are hand-waving. Only small-scale problems (up to a problem size of 100) are tested, while numerical solvers typically need to deal with large-scale problems (say, $Ax = b$ where $A$ is a sparse matrix of size 1000000x1000000). 
- The authors claim a "competitive performance", yet no quantitative results are reported in the experiments section (e.g., time cost, final accuracy achieved). 
- The baseline is not convincing. In practice, CG solvers are typically preconditioned with multigrid (MG) preconditioners. A sole CG iterator does not fully represent the SOTA. 
- The authors evaluate the MSE error during training and testing. This is OK for intermediate steps, however, in practice, the final precision of a numerical solution is evaluated by the relative residual of the solution, $r(x) = ||Ax - b|| / ||b||$. I would like to see the quantitative data on residual errors. 
- This manuscript is formatted with the NeurIPS 2024 template, please re-format it with the ICML template. Also, is this manuscript also submitted to NeurIPS 2024?
questions: questions: See "Reasons To Reject" above. I would recommend an acceptance after I see additional experimental data on the issues detailed above.