NLAFormer: Transformers Learn Numerical Linear Algebra Operations

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 numerical performance of NLAFormer.