NLAFormer: Transformers Learn Numerical Linear Algebra Operations
Zhantao Ma, Yihang Gao, Michael K. Ng

TL;DR
NLAFormer is a transformer-based model designed to learn and perform various numerical linear algebra operations efficiently, demonstrating the ability to learn algorithms like the conjugate gradient method for solving linear systems.
Contribution
This paper introduces NLAFormer, a novel transformer architecture that directly learns linear algebra operations, reducing complexity and enabling learning of classical algorithms.
Findings
NLAFormer successfully learns key linear algebra operations.
The model can approximate the conjugate gradient method.
Experiments show competitive numerical performance.
Abstract
Transformers are effective and efficient at modeling complex relationships and learning patterns from structured data in many applications. The main aim of this paper is to propose and design NLAFormer, which is a transformer-based architecture for learning numerical linear algebra operations: pointwise computation, shifting, transposition, inner product, matrix multiplication, and matrix-vector multiplication. Using a linear algebra argument, we demonstrate that transformers can express such operations. Moreover, the proposed approach discards the simulation of computer control flow adopted by the loop-transformer, significantly reducing both the input matrix size and the number of required layers. By assembling linear algebra operations, NLAFormer can learn the conjugate gradient method to solve symmetric positive definite linear systems. Experiments are conducted to illustrate the…
Peer Reviews
No public reviews on file for this paper yet. If you reviewed it on a platform where reviews are public (OpenReview, ICLR, NeurIPS, ICML), you can paste yours below so the community can read it here.
Videos
No videos yet. Explain this paper in a talk, walkthrough, or lecture? Add one.
