Linear algebra with transformers

1 INTRODUCTION . Since their introduction by Vaswani et al . ( 2017 ) , transformers , originally designed for machine translation , were applied to various problems , from text generation ( Radford et al. , 2018 ; 2019 ) to image processing ( Carion et al. , 2020 ) and speech recognition ( Dong et al. , 2018 ) where they soon achieved state-of-the-art performance ( Dosovitskiy et al. , 2021 ; Wang et al. , 2020b ) . In mathematics , transformers were used for symbolic integration ( Lample & Charton , 2019 ) , theorem proving ( Polu & Sutskever , 2020 ) , formal logic ( Hahn et al. , 2021 ) , SAT solving ( Shi et al. , 2021 ) , symbolic regression ( Biggio et al. , 2021 ) and dynamical systems ( Charton et al. , 2020 ) . All these problems pertain to symbolic mathematics , or involve a large amount of symbolic computation . When working on these tasks , transformers manipulate mathematical symbols , just like words in natural language . But mathematics are not limited to symbol manipulation : many practical applications involve numerical calculations , either exact ( e.g . arithmetic ) or approximate ( e.g . function evaluation , numerical solutions of equations ) . The use of transformers for numerical computation has been less studied , and many early experiments with arithmetic have proved disappointing ( Nogueira et al. , 2021 ) . This is , nevertheless , an important question : most problems in mathematics and science involve both symbolic and numerical computations . If we want transformers to solve these problems end-to-end , they need to be able to perform numerical calculations with high accuracy . In this paper , we train transformers to compute solutions of problems of linear algebra , which serve as fundamental building blocks in many scientific problems : basic operations on matrices , matrix inversion , eigenvalue and singular value decompositions . We introduce and discuss four encodings to represent problems and solutions as sequences that transformers can process , and train small transformers ( up to 6 layers , 10 to 50 million trainable parameters ) over generated datasets of random matrices . Trained models can compute approximate solutions to these problems ( to a few percents of their L1 norm ) with over 90 % accuracy ( 99 % in most cases ) . We also show that they can generalize out of their training distribution , and be retrained to extrapolate to larger problems than the ones they were trained on . We believe these results pave the way for using transformers as end to end solvers for problems of mathematics and science . After introducing the problems of linear algebra we are studying and presenting the encodings we use to represent them as sequences that can be used by our models , we discuss data generation , architecture and experimental settings . Then , we present our experiments on nine different problems , and discuss out-of-distribution generalization and few shot learning for eigenvalue computation . Finally , we discuss our results and future directions for research , and present related works . 2 PROBLEMS AND DATASETS . Let M and N be m× n matrices and V ∈ Rm . We study nine problems of linear algebra : • matrix transposition : find MT , a n×m matrix , • matrix addition : find M +N , a m× n matrix , • matrix-vector multiplication : find MTV , a vector in Rn , • matrix multiplication : find MTN , a n× n matrix , • eigenvalues : M symmetric , find its n ( real ) eigenvalues , sorted in descending order , • eigenvectors : M symmetric , find D diagonal and Q orthogonal such that M = QTDQ , set as a ( n+ 1 ) × n matrix , with ( sorted ) eigenvalues in its first line , • singular values : find the n eigenvalues of MTM , sorted in descending order , • singular value decomposition : find orthogonal U , V and diagonal S such that M = USV , set as a ( m+ n+ 1 ) ×min ( m , n ) matrix , • inversion : M square and invertible , find its inverse P , such that MP = PM = Id . These problems range from basic operations on individual coefficients of the input matrices ( transposition and addition ) , to computations involving several arithmetic operations over many coefficients ( multiplication ) , and complex nonlinear transformations involving the whole matrix , with cubic complexity ( decompositions and inversion ) . For each problem , we generate datasets of pairs of matrices ( I , O ) , by sampling random input matrices I ( see section 2.2 ) , and computing the output O with a linear algebra package ( NumPy linalg ) . When a problem has several input or output matrices , they are concatenated into one ( for instance , the two m× n operands of the addition task are concatenated into onem×2nmatrix I ) . All coefficients in I andO are set in base ten floating-point representation , and rounded to three significant digits in the mantissa . 2.1 ENCODING MATRICES AS SEQUENCES . The input and output to our problems are matrices . To be processed by transformers , they need to be converted into sequences of tokens . We encode a m×n matrix by first coding its dimensions as two symbolic tokens ( Vm and Vn ) , followed by its mn coefficients , encoded as sequences . Through this paper , we will use four encoding schemes for matrix coefficients : P10 , P1000 , B1999 , and FP15 . In base 10 positional encoding ( P10 ) , a number is represented as a sequence of five tokens : one sign token ( + or - ) , 3 digits ( from 0 to 9 ) for the mantissa , and a symbolic token ( from E-100 to E+100 ) for the exponent . For instance 3.14 will be represented as 314.10−2 , and encoded as [ + , 3 , 1 , 4 , E-2 ] . P1000 ( positional base 1000 ) provides a more compact representation by encoding the mantissa as a single token ( from 0 to 999 ) , and representing a number as the triplet ( sign , mantissa , exponent ) . B1999 ( balanced base 1999 ) pushes this one step further by encoding together the sign and mantissa ( from -999 to 999 ) . Finally , FP15 encodes each floating point number x = m10b as a unique token FPm/b . Table 1 provides examples of these encodings . Additional details and examples can be found in Appendix A . Selecting an encoding is a trade-off . Long encodings ( P10 , P1000 ) embed knowledge about numbers that the model can leverage ( e.g . numbers can be compared using their signs and exponents , addition and multiplication can be learned by memorizing small tables ) . Compact encodings use a larger vocabulary ( harder to learn ) , but generate shorter sequences that facilitate training with transformers . In P10 , a 20 × 20 matrix is a sequence of 2002 tokens , close to the practical limit of transformer implementations that use a quadratic attention mechanism . In FP15 , it is only 402 tokens long . 2.2 RANDOM MATRIX GENERATION . In most of our experiments , we train models over datasets of random matrices with uniformly distributed coefficients in [ −A , A ] ( with A = 10 ) . Occasionally , we sample gaussian coefficients with the same standard deviation ( σ = A/ √ 3 ) . In the symmetric case , these matrices are known as Wigner matrices . Their eigenvalues have a centered distribution with standard deviation σ = √ ns , where s is the standard deviation of the coefficients ( Mehta , 2004 ) . As n increases , this distribution converges to the semi-circle law ( p ( λ ) = √ 4σ2 − λ2/2πσ2 ) for all coefficient distributions with bounded variance . If the coefficients are gaussian , the associated eigenvectors are uniformly distributed over the unit sphere . When investigating out-of-distribution generalization for the eigenvalue problem , we will need to generate random symmetric matrices with different distributions of their eigenvalues ( corresponding to random matrices with non iid coefficients ) . To this effect , we randomly sample symmetric matrices M , with gaussian coefficients , and compute their eigenvalue decomposition M = PDPT , with P the orthogonal matrix of eigenvectors ( uniformly distributed over the unit sphere since the coefficients are gaussian ) . We then replace D , the diagonal matrix of eigenvalues of M , with a diagonal D′ sampled from another distribution . Finally , we recompute M ′ = PD′PT , a symmetric matrix ( because P is orthogonal ) with eigenvalues distributed as we choose , and eigenvectors uniformly distributed over the unit sphere . 3 MODELS AND EXPERIMENTAL SETTINGS . We use the standard transformer architecture introduced in Vaswani et al . ( 2017 ) , with an encoder and a decoder connected by a cross-attention mechanism . Most of our models have 512 dimensions , 8 attention heads and up to 6 layers . We experiment with different number of layers and attention heads in the encoder and decoder . All training is supervised , and minimizes the cross-entropy between model prediction and the correct solution . We use the Adam optimizer ( Kingma & Ba , 2014 ) with a learning rate of 10−4 , an initial warm-up phase of 10,000 steps and cosine scheduling ( Loshchilov & Hutter , 2016 ) . All training data is generated on the fly , in batches of 64 examples . Every 300,000 examples , 10,000 random problems are generated and used to evaluate the model . When evaluating , we consider that a predicted sequence seqP is a correct solution to the problem ( I , O ) ( I and O the input and output matrices ) if it can be decoded as a valid matrix P ( several matrices for singular and eigen decomposition ) that approximates the correct solution to a given tolerance τ ( τ ∈ { 5 , 2 , 1 , 0.5 % } ) . For addition , transposition , multiplication , eigen and singular values we check that P verifies ‖P − O‖ < τ‖O‖ ( with ‖A‖ = ∑ i , j |ai , j | , for A = ( ai , j ) , i.e . L1 norm ) . For eigenvalue decomposition , we check that the solution ( Q , D ) predicted by the model can reconstruct the input matrix , i.e . ‖QTDQ− I‖ < τ‖I‖ . For singular value decomposition , we check that ‖USV − I‖ < τ‖I‖ . For matrix inversion , we check that ‖PI − Id‖ < τ‖Id‖ = τ . The choice of the L1 norm is important : norms like L2 and L∞ will favor models that correctly predict the largest coefficients in the solution . For eigen and singular value problems , this amounts to predicting the largest values , an easier problem than the one we want to solve . We consider different tolerances for different problems . Since we round numbers to three significant digits , 0.5 % is the best we can hope . In fact , a number xwith mantissa 1.00 is subjected to a maximal rounding error of 0.5 % ( x ∈ ] 1.005 , 0.995 ] ) , which may accumulate when several ( rounded ) numbers are summed , and increase again when nonlinear operations are considered . When discussing results , we consider tolerances of 0 % for transposition , which involves no arithmetic , 1 % for basic matrix operations ( addition , multiplication ) , and 2 or 5 % for non linear operations ( decomposition , inversion ) , but we usually provide results for all tolerance levels . Most of our experiments focus on 5× 5 square matrices , or rectangular matrices with as many coefficients ( e.g . 6× 4 , 2× 13 ) . This helps when comparing encodings : for larger dimensions , varying sequence lengths make comparisons difficult . We also study scaled-up versions of the problems ( from 8 × 8 to 15 × 15 ) , and datasets with matrices of variable dimensions ( 5-10 or 5-15 ) . In this paper , we limit ourselves to problem that can be solved using small models ( with up to 6 layers ) . Scaling to larger problems , and leveraging deeper architectures is left for future research .

Summary: The authors train generic, dense transformers to perform several standard linear algebra computations, ranging from simple tasks like transposition to complex nonlinear tasks such as matrix inversion. They restrict themselves to relatively small matrices due to the practical limits of the dense, quadratic attention mechanism. The main result of the paper is that transformers can perform fairly well on all tasks, meaning that they can usually produce outputs that are correct upto relatively small tolerance. The paper also shows that some forms of out-of-distribution generalization are possible, and that this phenomenon is sensitive to the details of the training distribution.