Table of Contents
Fetching ...

Linear algebra with transformers

François Charton

TL;DR

The paper investigates whether transformers can learn numerical linear-algebra computations from data. It develops four real-number encodings and trains encoder-decoder transformers on nine problems from basic matrix operations to decompositions, achieving high accuracy (>90%) and demonstrating robustness to noise and out-of-domain distributions. It provides a detailed analysis of when and why failures occur (notably in eigenvectors and inverses) and shows that with appropriate training data, transformers can generalize beyond their training distribution. The findings suggest transformers can complement traditional numerical methods and offer insights into AI for science, while outlining scalability limits and directions for future work with larger or sparse matrices.

Abstract

Transformers can learn to perform numerical computations from examples only. I study nine problems of linear algebra, from basic matrix operations to eigenvalue decomposition and inversion, and introduce and discuss four encoding schemes to represent real numbers. On all problems, transformers trained on sets of random matrices achieve high accuracies (over 90%). The models are robust to noise, and can generalize out of their training distribution. In particular, models trained to predict Laplace-distributed eigenvalues generalize to different classes of matrices: Wigner matrices or matrices with positive eigenvalues. The reverse is not true.

Linear algebra with transformers

TL;DR

The paper investigates whether transformers can learn numerical linear-algebra computations from data. It develops four real-number encodings and trains encoder-decoder transformers on nine problems from basic matrix operations to decompositions, achieving high accuracy (>90%) and demonstrating robustness to noise and out-of-domain distributions. It provides a detailed analysis of when and why failures occur (notably in eigenvectors and inverses) and shows that with appropriate training data, transformers can generalize beyond their training distribution. The findings suggest transformers can complement traditional numerical methods and offer insights into AI for science, while outlining scalability limits and directions for future work with larger or sparse matrices.

Abstract

Transformers can learn to perform numerical computations from examples only. I study nine problems of linear algebra, from basic matrix operations to eigenvalue decomposition and inversion, and introduce and discuss four encoding schemes to represent real numbers. On all problems, transformers trained on sets of random matrices achieve high accuracies (over 90%). The models are robust to noise, and can generalize out of their training distribution. In particular, models trained to predict Laplace-distributed eigenvalues generalize to different classes of matrices: Wigner matrices or matrices with positive eigenvalues. The reverse is not true.
Paper Structure (38 sections, 3 equations, 3 figures, 29 tables)

This paper contains 38 sections, 3 equations, 3 figures, 29 tables.

Figures (3)

  • Figure 1: Learning accuracies for different problems measured with norms $L^1$, $L^2$and $L^\infty$ (Max).
  • Figure 2: Learning curves for different problems. All problems except addition use $5 \times 5$ matrices. All models have $512$ dimensions and $8/8$ heads (except when mentioned in the legend). Inversion models have $6/1$ layers. Epochs correspond to 300,000 training examples. Test loss is cross-entropy.
  • Figure 3: Empirical distributions of eigenvalues for Wigner matrices, dimension 5x5 (left) to 20x20 (right), with uniform (top), gaussian (middle) and Laplace (bottom) coefficients. All distributions computed from 100 000 random matrices.