Linear algebra with transformers

Abstract: Most applications of transformers to mathematics, from integration to theorem proving, focus on symbolic computation. In this paper, we show that transformers can be trained to perform numerical calculations with high accuracy. We consider problems of linear algebra:  matrix transposition, addition, multiplication, eigenvalues and vectors, singular value decomposition, and inversion. Training small transformers (up to six layers) over datasets of random matrices, we achieve high accuracies (over 90%) on all problems. We also show that trained models can generalize out of their training distribution, and that out-of-domain accuracy can be greatly improved by working from more diverse datasets (in particular, by training from matrices with non-independent and identically distributed coefficients). Finally, we show that few-shot learning can be leveraged to retrain models to solve larger problems.



Reviews and author responses:

 Thank you for your comments. We believe the co-training results have to do with the relative difficulty of tasks. Without co-training, it takes less than 50 epochs (with 300k examples per epoch) for the model to learn eigenvalues, over 200 for eigenvectors, and more than 300 for inversion (see figure 2 page 15). It is less than 50 for the basic tasks (TADM). 

With co-training, in the TADMEFI case (all tasks in equal proportion), the model quickly learns the easy tasks that constitute the majority of the training data, and this seems to impact negatively the hard tasks (F and I), which need more examples to get started. The situation is a little better in the EFI case: without the easiest tasks mobilising the majority of samples,  F can be learned to some accuracy (20%). But the easier tasks still dominate. 

A solution might just be to balance the training sets by oversampling the harder cases. Another possibility would be to adjust task frequencies dynamically, by sampling tasks in inverse proportion of model accuracy. All tasks would then begin with the same frequency, but easier tasks would be less considered as the model learn them. We are experimenting with both approaches, and will report on them. 




 I thank the authors for their response, for adjusting their framing, for including additional experiment data, and for expanding on why they found the asymmetric generalization results surprising. Together these address the main concerns I raised. I also appreciate the new co-training experiments, although I do not know what to make of the low scores on F and I -- did training loss go low on these tasks and yet it still failed on the test set, or did the models just not have the capacity for all the tasks?

 I thank the authors for their response, which addressed some of the points raised in my initial review, and I welcome the additional experiments with noisy supervision and different recurrent models. 
I still find the significance of the findings limited and I wish that the authors proposed novel model classes (tailored for the specific problems) rather than reporting the performance of a generic architecture. Nonetheless, this paper is well written and could help future research on solving algebraic problems with neural networks. I'm raising my score accordingly.


This paper consider the problem of approximating algebraic computations over matrices using transformers.  
Experiments with different encodings are presented, investigating the use of transformers for approximating a number of algebraic operations. While I found the paper well written, I didn't find it very impactful. The authors are not proposing a novel technique for addressing the problem, but only report some experiments with different encodings for the matrices and a standard transformer architecture. I was hoping to get more insights after reading this article, such as what architectural choices are beneficial for each specific problem and why. Instead, the proposed approach is training an off-the-shelf model with randomly distributed examples, and the experiments do not consider alternative techniques.

In my opinion, the motivation for this approach is also lacking, since the reported experiments only consider very small problems that can be solved exactly. I wish that the paper considered instead cases that need to be approximated, or possibly prove that indeed transformers trained on smaller problems can generalise to much higher dimensions. Since the learning algorithm has access to an oracle (numpy) that can provide exact supervision, an interesting problem is how to select the most informative input instances to be solved during training. Another aspect that is worth investigating in my opinion is how to cope with noise in the training data, allowing the training of transformers with an approximate oracle. - There isn't a novel contribution besides some experiments with an off-the-shelf model
- The motivation is not supported by the experiments, which only report results in low-dimensional settings
- The experiments do not consider any other technique besides transformers 

 Thank you very much for your review and comments

**Framing**

*I found the following claim to be problematic: 'Our results on out-of-distribution generalization provide justification to the idea that models trained over random data can be used to solve "real world" problems'.*

We agree with your comment, and remove the last part of our conclusion (which included this sentence). On the motivation and practicality of our work, we believe that, since linear algebra is ubiquitous in problems of science, demonstrating that transformers can solve such problems is important. If it were impossible to train transformers on such problems, their applicability to problems of science would be very limited. 

**Sparse data reporting**

We have added a lot of additional results in the appendix. Appendix B discusses the impact of various norms for evaluation. Appendix C1 provides learning curves for loss and accuracy, and C2 discusses the impact of model size (dimension and depth). We present alternative architectures (LSTM, GRU and Universal/shard layers Transformers) in Appendix D and experiments with noisy data in section E1. In the main paper, we indicated parameters in the legend of our tables.

**Out of distribution findings seem unsurprising**

The result we found most surprising is the fact that models trained on matrices with Laplace-distributed eigenvalues could generalize to Wigner matrices (while the reverse was not true). Most previous works on the subject considered that the key to out-of-distribution generalization was to train on large and diverse datasets, and it seems very natural to train models on large groups of “typical” examples. Now Wigner matrices form a large class of very typical matrices, which play a role similar to the normal distribution in statistics. Matrices with Laplace eigenvalues, on the other hand, are a very specific family. 
Another factor that contributed to our surprise is that many recent works (see for instance Welling et al. https://arxiv.org/abs/2109.13986) have shown that many transformer models only generalize close to their training distribution. Our result appear like a counter-example.

**Co-training?**

Thank you very much for this suggestion. We launched co-training experiments, and added a section in Appendix (E.2). The experiments are still running, and the full results should be available in a few days (and will be added to the paper). The results so far indicate that

- arithmetic tasks (add, transpose, multiply) can be co-trained at no cost (transpose and add achieve 100% accuracy after 600k examples, transpose+add+vector multiplication in 3 million.
- “full cotraining” (addition, products, eigen decomposition and inversion) seem possible, as of now 60% of test examples are accurately solved, which suggests that most (if not all) of the tasks are being learned.

To illustrate, one of the models currently trained reports, after 30 epochs (9 million examples) the following test accuracies on the different tasks:

- Translation : 100%
- Addition: 100%
- Dot product (vector matrix multiplication): 73%
- Matrix multiplication : 99%
- Eigenvalues: 2% 
- Eigenvectors and inversion 0%

**edit 11/22** We have added new results on co-training, confirming that basic operations can be learned to 100% accuracy. A more precise discussion can be found in appendix E2 of the updated paper.

 Hey everyone,

After reading the reviews and the author response, my impression is that the main concern is the *motivation* behind solving linear algebra with learning approaches to begin with. I agree with reviewers TnqM and Jby6 that this is very unlikely to ever be used in practice, because there obviously are more efficient (and more precise) algorithms for these problems. I still find the paper very insightful: In discussions on solving mathematical problems with learning methods there are frequent claims that ML models could never learn this or that. This paper represents a very interesting datapoint on the limits of ML models that could help to guide future research, which I think is a good motivation to accept the paper. The authors make some good additional arguments in the response to reviewer Jby6.

I'd love to hear the opinions of reviewers TnqM and Jby6 on this argument.

 We thank the reviewers for their comments and insightful feedback. We replied to their questions individually, and updated the paper by adding new material which roughly doubled its size. 

Here are the major updates, and the key lessons from additional experiments.
**11/21 edit : final results on co-training**
 
As suggested by reviewer LVVb, we edited the end of the conclusion and the reference to “real world” problems.

We added many sections to the appendix (11 pages overall)

- Appendix B : A discussion on norms for evaluation, justifying our choice of L1, with learning curves for all tasks (because L2 norm overestimates accuracy of eigen and singular value predicition).
- Appendix C1: Learning curves for loss and accuracy, with a discussion on training speed, attention heads, and learning rates and noise.
- Appendix C2 : Experiments with larger models for additions, products, and eigenvalues. 

    *Key lesson:* Increasing embedding dimension helps, transformer depth not so much. 
- Appendix C3 : Model performance on various Wigner, non-Wigner, and hybrid datasets
- Appendix D1 : Experiments with other architectures (LSTM and GRU), for additions and eigenvalues

     *Key lesson:* LSTM and GRU have good performance on eigenvalues, but bad on addition.
- Appendix D2 : Experiments with shared layer transformers (Universal Transformers), for addition and eigenvalues

     *Key lesson:* shared layers transformers are a possible alternative to our current models, which uses less parameters (but not necessary less memory)
- Appendix E1 : Experiments with noisy data on the addition and eigenvalue tasks.

     *Key lesson:* noisy data can be learned, so long signal to noise ratio does not exceed tolerance on accuracy
- Appendix E2 : Co-training   

     *Key lesson:* co-training is possible on basic matrix operations (transposition, addition, multiplication).
- Appendix F : Number of parameters, formulas and size of the models we use
- Appendix G: Empirical distributions of eigenvalues of Wigner matrices, for different dimensions and distributions of coefficients

 Thank you very much for your comments and suggestions. 

**Novelty and motivation**

*There isn't a novel contribution besides some experiments with an off-the-shelf model*

We believe that demonstrating that transformers can perform tasks from linear algebra with high accuracy is a novel contribution. Many previous works have noted the difficulty of training neural networks and especially sequential models on arithmetic tasks. We believe this is important: the solutions to many problems of science use linear algebra, proving that transformers (or other deep models) can solve them is a necessary step if we want to use deep networks for science. 


**Noisy data**

*Another aspect that is worth investigating in my opinion is how to cope with noise in the training data, allowing the training of transformers with an approximate oracle.*

Thank you very much for this suggestion. We added experiments with noisy data, for addition and eigenvalue calculation (Appendix E1). They prove that transformers can learn from noisy data, so long the signal to noise ratio (measured as the ratio of the standard deviation of the noise to the standard deviation of the matrix coefficients) does not exceed the tolerance we use to evaluate models (e.g. models trained on data with 1% error can only predict the noiseless solution with 1% error). We also note that, because of the Wigner laws, models trained on the eigenvalue problem are slightly more robust than models trained on addition.

**Other techniques**

*The experiments do not consider any other technique besides transformers*

Appendix D was added, which provides results on different architectures. In D1, we show that LSTM and GRU can perform as well as transformers on some problems (eigenvalues), but also fail badly on others (addition), confirming previous results on the difficulty to train seq2seq models on arithmetic tasks (this comforts our claim to novelty on this aspect). In section D2, we experiment with Universal Transformers, models with only one layer that is looped through several times. We show that they are a potential alternative to the models discussed in the paper. Appendix C also provides additional experimental results, but using transformers.


 Thank you very much for your comments and suggestions.

**L1 distances**
 
*I would have loved to see more than just L1 distances*

We added a new section, Appendix B, that provides learning curves for L1, L2 and L-infinite distances, and a discussion of their relative merit. The main findings are that whereas L1 and L2 are equivalent on basic matrix algebra, they are not for the eigen, singular decomposition and inversion problems. The distribution of eigenvalues of random matrices is very constrained (it has a fixed variance, and eigenvalues tend to “repulse each other”, and spread all over the interval), and this makes large absolute values easy to predict. The net result is that bad models tend to have high accuracy according to the L2 norm. A similar situation seems to prevail for matrix inversion. This makes L1 preferable to L2. L-infinite is a possible alternative, but the fact that it is computed on just one element from the solution makes it less robust than L1.


**Wider training and evaluation distributions**

*I would have loved to see a study that includes a wider variety of training and evaluation data.*

We added results on hybrid or non-Wigner samples (Appendix C3). There is also a new section on noisy data (appendix E1).

**Larger, and different, models**

*The models used in this paper have sometimes rather odd (small) hyper parameters. E.g., for many experiments the models have only 2 layers. I’d love to see larger models and see if they improve the results (and the exact number of parameters).*

We added a lot of experimental results in the Appendix. There is a specific section on larger models (Appendix C2), which shows that for addition, product and eigenvalues, increasing embedding dimension helps learning better and faster, while adding layers brings no improvement. This is a significant departure from mainstream ideas on transformers, which usually recommend deeper architectures. 

We also tried shared layer transformers (variants of Universal Transformers) in appendix D2. We show that they achieve comparable results with fewer parameters (as they are, basically, transformers with only one layer that is iterated through).

Finally, in appendix D1, we have added comparisons with other seq2seq architectures (LSTM and GRU). 

**Number of parameters**

Appendix F now provides an exact formula for computing the number of parameters of the transformer architecture we use, and gives the size of some of the models used in this paper.

(Note that for the models using the FP15 encoding, the number of parameters is probably inflated, because many possible values of the encoding are never seen in our problems and the corresponding embeddings are never trained). 


 Thank you very much for your review and comments. 

**Motivation of our research**

We believe our research is important, and can be useful, for four reasons.

1- Demonstrating that such tasks can be learned by transformers (or other sequence to sequence models) is important, if we want to use those models to solve problems of science. It has been observed by many authors that transformers struggle with arithmetic tasks. If such common computational building blocks as the linear algebra tasks we propose here could not be solved by transformers, the applicability of transformers to many problems of science would be severely limited.

2- Even though we already have algorithms for those problems, being able to solve them inside a neural network is important. Transformers (and other networks) are often used as end-to-end solutions, and “stepping out” of the network every time a computation has to be done would be very impractical. This line of reasoning led to the development of Neural Arithmetic Logic Units (NALU): neural networks components for addition and multiplication (which we obviously could perform with perfect accuracy long before neural networks). Our work shows that such components could be developed for more advanced (and very ubiquitous) computations, such as matrix computations.

3- Some of our results provide new insights on transformers. While mainstream research focuses on very deep models, with a very large number of parameters, our results demonstrate the efficiency of shallow models (and the new results on universal transformers suggest that this could be pushed even further). 

4- The encodings we propose can serve other purposes. Notably, they can be used in any setting that uses real numbers, such as symbolic regression.

**Asymptotic efficiency**

We will definitely experiment with linear attention models. However, we do not think asymptotic efficiency will translate into a large advantage for transformer-based implementations. There are several reasons for this, but the most important, in our opinion, is the memory constraint : as matrix dimension increases, so does GPU memory requirements. Transformers process full sequences that need to be held in GPU memory for the process to be efficient. We believe that GPU memory will put a hard limit to the size of the dense matrices that can be processed. 

**Small points**
- Page 2 and 7 were edited as you suggest
- Page 5 : The property is true for large n as well: sqrt(n) s is the standard deviation of the limit semicircle law, too. The correct statement of this property can be found in section 2.2, we rewrote the sentence accordingly.
We have verified it empirically for the small values of n we are considering in this paper (5-20), and added an illustration (and explanatory comments) describing the phenomenon in  appendix G.


This paper “Linear algebra with transformers” studies the application of seq2seq transformers to matrix operations. It studies their performance across different encodings of floating point numbers, different sizes of matrices, different operations, and different (synthetic) data distributions. The main findings are that transformers work surprisingly well on various matrix operations (addition, multiplication, eigenvalues, inversion, SVD, …) for small matrices (e.g. 5x5), and that generalization to OOD problems is not symmetric (I.e. generalization from one distribution to another does not imply the other way round). Strengths:
- The authors present a detailed study of many important questions. This is a fairly comprehensive work on the idea.
- Thought-provoking application of transformers.
- Very well written.

Weaknesses:
- I would have loved to see more than just L1 distances
- The paper studies the question only on random matrices. In other symbolic domains we have seen that insights gained from machine learning approaches trained on random data do not necessarily carry over to “real-world” distributions. I would have loved to see a study that includes a wider variety of training and evaluation data.
- The models used in this paper have sometimes rather odd (small) hyper parameters. E.g., for many experiments the models have only 2 layers. I’d love to see larger models and see if they improve the results (and the exact number of parameters). This work explores a wonderful idea: to solve linear algebra with transformer models. While these models use way more compute internally than the problem they are applied to requires to solve, it is an intriguing question whether these computations can be learned from scratch without further biases. The surprising answer is that this works relatively well for small matrices.


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. I found this to be an interesting paper overall.

- Framing

I found the following claim to be problematic: 'Our results on out-of-distribution generalization provide justification to the idea that models trained over random data can be used to solve "real world" problems'. First, the authors only evaluate on other synthetic distributions. Second, in most "real world" problems, matrices are gigantic relative to the tiny context windows of dense transformers. Third, since traditional methods are always perfectly accurate on all distributions, I think the claim carries with it some burden to elaborate on why such noisy (and much less scalable) methods might prove useful in practice. Note that even if the potential practicality cannot be argued for, I think the experiments are interesting enough to stand on their own. 

- Sparse data reporting

I would have liked to see much more data collected from the experiments, especially train-loss, validation-loss, and correctness-upto-tolerance curves over a range of architectures. The curves would also make it clear how many samples had been trained on for each measurement, which would be useful for understanding the relative performance of the different encodings. Note: it is not always clear from the current tables which encoding is even being used. It would also be interesting to see some analysis/visualization of the attention patterns, at least for tasks with relatively simple ground-truth algorithms like transposition and addition. Some experimental results also seem to be omitted; for example, S4.3 claims that "deeper decoders are needed" for matrix-matrix multiplication, but Table 5 does not include enough data to defend this claim.

- Out of distribution findings seem unsurprising

It is great that the authors assess out-of-distribution, and I appreciate the negative result of generalizing from Wigner to matrices with positive eigenvalues. However, although the details of the subsequent out-of-distribution experiments are interesting, I found it generally unsurprising that models trained on non-Wigner matrices would generalize better, and I thought that the authors tried to make too big a point of this finding.

- Co-training?

Although not essential, I would be interested to see how co-training on many of the tasks at once affects sample efficiency.
 Despite the concerns listed above, I think this paper does contribute to our emerging understanding of transformers, and that many people in the community will find it worth engaging with.

This paper describes several experiments where transformers are trained to perform real-valued linear algebra calculations (matrix transposition, addition, multiplication, eigenvalues & eigenvectors of symmetric matrices, SVD, inversion). In-distribution accuracy is generally very high, whereas care is needed in order to obtain out-of-domain generalization. The paper carries out a very thorough set of experiments on linear algebra calculations with transformers, using four different encodings of input matrices. In addition, the authors are aware of the importance of out-of-distribution generalization, varying both the matrix size and the distribution of input matrices of a given size. Results appear to be complete, and the conclusion drawn from them are generally sound.

However, the problem tackled in this paper does not appear to be particularly useful. In my opinion, the conclusions and findings of this paper are only interesting on a theoretical level (perhaps they can help understand what transformers can or cannot do), rather than being directly applicable in a meaningful way. After all, we do have algorithms for all linear algebra problems considered, and they work with 100% accuracy, perfect out-of-domain generalization, and faster run time.

As the authors note in the discussion, at the current stage transformers have quadratic complexity in the number of tokens, which translates into $O(n^4)$ complexity for $n \times n$ input matrices, and this is asymptotically slower than the exact algorithms we have. A potentially interesting future direction (which perhaps can be advertised more by the author, to strengthen the claim that this paper is useful) is to investigate linear-time transformers on tasks where the exact algorithm requires more than $O(n^2)$ time, so that perhaps transformers can be used to perform approximate computations with less time.

I also have the following minor comments.
- First line of page 2: $m \times n$ should be in a math formula.
- Section 5, fifth line: what does “for small values of n” mean here? The statement is very precise, so it is hard to believe that it is true up to a certain small number (say, 5) and false for a larger n.
- End of page 7, “This confirms that out-of-distribution generalization is possible when particular attention...”: actually, you showed that it is *necessary* to pay particular attention, not that it is *sufficient*. So I would write “out-of-distribution generalization requires particular attention...”. This paper provides a thorough and well-written investigation of the use of transformers to perform linear algebra computation. However, this does not appear to be particularly useful.

Meta-review: The paper demonstrates that transformer architectures can be trained to compute solutions of linear algebra problems with high accuracy. This is an interesting direction and, as the the reviews and the discussion show it is a "good data point and insightful", as one reviewer puts it. I fully agree with this but also agree with one other reviewer in that this is "yet another" application of a known transformer architecture. The author should place the model into context and provide some perspective. Without, the motivation behind solving the specific set of linear algebra problems considered is a bit unclear. For instance, could a transformer now learn to solve corresponding ML problems? Moreover, the dimensions of the considered matrices are rather small, and the generalization to larger dimension appear to be tricky.

Decision: Reject