Understanding and Relaxing the Limitations of Transformers for Linear Algebra
Authors: Andres Potapczynski
Alex Ali
Andrew Gordon Wilson
Venue: ICLR 2026 Poster (2026)
Keywords: Numerical Linear Algebra
Transformers
Out-of-distribution
Generalization
Primary area: interpretability and explainable AI
TL;DR: We show that current transformer methods for linear algebra are simply learning statistics about the data

Abstract: Matrix operations, such as linear solves, eigendecompositions, and log determinants, are foundational building blocks for any number of downstream applications. Therefore, any broadly capable learning system should be able to effectively approximate these operations in its internal representation. Accordingly, there is great motivation to study transformers for linear algebra --- for if transformers cannot even semi-competently perform matrix operations, then we cannot expect them to form a basis for a generally intelligent system. We demonstrate that current techniques developing transformers for linear algebra have striking failure modes, prohibitive scaling, and particularly poor out-of-distribution generalization to other matrix distributions, and matrices of different sizes. Investigating further, we find that current transformer approaches operate as statistical interpolators, rather than discovering algorithms that will generalize to matrices from other distributions. Based on our understanding of these limitations, we develop a sequence of interventions that substantially improve scaling and performance, including matrix embeddings through a learnable projection, linear attention, looping, and a data pre-training distribution of structured matrices. We term the resulting method the \emph{RangeFormer}, which we show has significantly improved scaling and performance on challenging OOD matrices from the \emph{matrix market}. Moreover, with RangeFormer we show for the first time that transformers can be successfully applied to downstream tasks that involve iterative matrix operations, including Gaussian process learning, and improving the sampling distribution of randomized methods.

Published as a conference paper at ICLR 2026
UNDERSTANDING AND RELAXING THE LIMITATIONS OF TRANSFORMERS FOR LINEAR ALGEBRA
Andres Potapczynski", Alex Ali & Andrew Gordon Wilson New York University
ap ap6604enyu edu, aa102538nyu edu, andrewgwecims nyu. edu
ABSTRACT
Matrix operations, such as linear solves, eigendecompositions, and log determi- nants, are foundational building blocks for any number of downstream applications. Therefore, any broadly capable learning system should be able to effectively ap- proximate these operations in its internal representation. Accordingly, there is great motivation to study how transformers can perform linear algebra- - for if trans- formers cannot even semi-competently perform matrix operations, then cannot expect them to form a basis for a generally intelligent system. We demonstrate that current techniques developing transformers for linear algebra have striking failure modes, prohibitive scaling, and particularly poor out-of -iistiiiiiiittoionn generalization to other matrix distributions, and matrices of different sizes. Investigating further, we find that current transformer approaches operate as statistical interpolators, rather than discovering algorithms that will generalize to matrices from other distri- butions. Based on our understanding of these limitations, we develop sequence of interventions that substantially improve scaling and performance, including matrix embeddings through learnable projection, linear attention, looping, and data pre- pre-training distribution of structured matrices. We term the resulting method the RangeFormin which show has significantly improved scaling and performance on challenging OOD matrices from the matrix market. Moreover, with Range- Former show for the first time that transformers can be successfully applied to downstream tasks that involve iterative matrix operations, including Gaussian process learning, and improving the sampling distribution of randomized methods.
INTRODUCTION
Linear algebra is at the heart of science and engineering. Linear solves, eigenvalue computations, log determinants, and other matrix operations, are the basis for Gaussian processes, normalizing
flows, principal component analysis, equivariant neural networks, differential equations, approximate
second-order optimizers, and any number of other applications (Rasmussen and Williams 2006 Anil
et al. Li et al. 2018 Martens
2020 Cuturi 2013 Dao et al. 2019 Finzi et al., 2023: Fu et al.] 20233 Kovachki et al.) 2021
Martens and Grosse [015)N et al. 2022 Perez et al. [0018] Potapezynski et al.
2023). It follows that broadly capable learning system must be able to perform matrix operations in its internal representation, as a fundamental building block.
Not only have transformers rapidly become the most popular architecture, they are being applied to distinctly wide array of settings. We have moved from hand-crafted feature engineering, to neural networks for particular modalities (CNNs for vision, RNNs for sequences, etc.), to transformers for everything (Goldblum et al. 2024). However, if transformers are to be general purpose, a pressing question arises: how good are transformers matrix operations, core primitive in many settings? Despite the importance of this question, there is so far relatively little work on transformers for linear algebra (e.g.. Charton 2022a Yang al. 2024).
In this paper, we shed light on the limitations and failure modes of existing transformer approaches for matrix operations. Standard approaches involve representing matrix sequence of real numbers or string tokens which is then mapped by transformer to the output of a particular operation,
Denotes equal contribution, names ordered by coin flip.

---

Published as a conference paper at ICLR 2026
such computing its eigenvalues (Charton, 2022a). These approaches have severe computational constraints, requiring a staggering O(N D²D compute and O + N2 memory, for an N x N matrix, and embedding dimension D. Moreover, when trained standard distributions, e.g.. random matrices of a fixed size with standard Gaussian entries, transformers generalize poorly to other matrix structures, even basic identity matrix. Probing further, we show current approaches perform statistical interpolation, rather than algorithm discovery, explaining their particularly poor OOD performance.
Based on our study of these limitations, we investigate a sequence of interventions for improved scalability and performance. We note that the goal of our work in understanding the potential capabilities and limitations of transformers for linear algebra, and methodological interventions are primarily in service of this understanding. First, rather than flatten the matrix, and represent each number as a string token, we embed the matrix through a learnable projection. We then pursue linear attention, which improves both scalability and the performance of matrix operations, as quadratic softmax attention is known to be poor for matrix multiplications (Arora et al. 2023 Liu et al. 2025). Together, these interventions significantly reduce compute to ((ND=) and memory to O
Next, we use looped transformer architecture (Giannou et al. 2023 Yang et al. 2024 Saunshi et al., 2023; Geiping et al.. (025). inspired by iterative methods in numerical linear algebra (such as linear conjugate gradients, and other Krylov subspace methods (Potapezynski et al. 2023)), ), that each matrix can be processed by the transformer variable number of times, depending on the difficulty of the operation and its conditioning. Finally, we construct and train on a novel data mixture comprised of a wide range of matrix structures, covering diverse cigenspectra decays, leading to significantly better out-of-distribution generalization than the standard Gaussian random matrix training: moving from statistical interpolation towards generalizable algorithms. We illustrate the cumulative effects of these interventions Figure I
We name our approach the RangeFormer (and release code here).
We show that the RangeFormer profoundly relaxes the limitations of prior transformer approaches, with substantial performance improvements on diverse OOD matrices from the matrix market (Boisvert et al., 1997). Finally, for the first time, we show how transformers can be successfully used in downstream applications requiring iterative matrix operations, including Gaussian process learning, and improving the sampling distribution of randomized methods. Indeed, the ultimate test of a numerical method is whether it can support a downstream application, as part of a pipeline, especially in setting that requires repeated matrix operations that build off of one another. In such settings, compounding errors, or systematic biases, could easily prevent convergence.
Accordingly, we summarize key contributions here.
RangeFormer enables scalable transformers for NLA. Past work (Charton 2022a; Garg et al., 2023] Liu et al. (025) has focused on matrices of size at most 50 x 50. Our key interventions of using a projection matrix for embedding, looping, and using linear attention enable us to train transformers for matrices up to size 1000 x 1000 (see Figure (Left)) Collectively, these interventions reduce the memory complexity of the attention layer from O(N*) to O(ND + D²).
we demonstrate that trans- formers when trained on Gaussian random matrices fail to generalize on even trivial OOD examples like the Identity matrix. This is because the eigenvalues of Gaussian matrices have a predicitable distribution and thus the transformer learns to predict statistics of the training data as opposed to learning a more robust algorithm.
Prior models only learnt statisties of the data. In Section 3
Diverse training data distribution improves OOD generalization. To remedy the above issue, we construct a more diverse, structured training data distribution whose eigenspectra can be observed in Figure] (a). Figure 5 (Left) illustrates the OOD performance gain that comes from training on this distribution.
RangeF RangeFormer ormer can be used in downstream linear algebra applications. For the first time, use transformers for downstream applications of linear algebra. Specifically, we warm-start Krylov subspace methods. improve the sampling distribution in randomized SVD, and perform marginal likelihood optimizations using the RangeFormer for dog-determinant calculations (Section 5)
2

---

Published as a conference paper at ICLR 2026
Cauadon Eigenspectm
Amar
Data
Flattening Embeldiig
Arch.
10
0OD
0G0 rel crr
Structonel Figmserrts
10a
102
10a
10
.TF
Runtime
20*
Learrshle Projectoon tlon
1
A
10'
Tr(A)
Ama(1)
logiA)
102
10a Matrix Size
10*
(a)
Figure I: How to profoundly improve the scaling and OOD performance of transformers for linear algehra. (a) We train on a wide variety of structured random matrices resulting diverse eigenspectra as opposed to only training with Gaussian random matrices. (b) In contrast to flattening the matrix and then projecting cach scalar into a high dimensional X E NR xDD with linear attention, achieve ((ND + D²) scaling instead of O(N*). (c) Performance improvement, over an ensemble of challenging OOD matrices (section B). applying set of interventions for our RangeF ormer method. over four canonical matrix operations. The baseline (Base) a vanilla transformer trained on Gaussian random matrices. We then train the same transformer but on our data mixture (Data) from section 4 Finally, we apply the architectural changes (Arch), such as looping, learnable projection (panel (b)) and linear attention. (d) OOD performance and runtime of RangeForme: (RF) against vanilla transformer baseline (TF) from section B TF cannot run matrices with N 100 on an 80 GB A100 and its performance worse than just predicting 0. In contrast, RangeFo r scales reasonably in compute and performance.
XD we embed the matrix via a learnable projection X E Nx Together
TRAINING TRANSFORMERS FOR LINEAR ALGEBRA TASKS
Assume that we have data in the form of f(A)) where A is N matrix and (AA is a linear algebra operation applied to A. say Tr( (t, log [A), Amax (A). or A and we would like to train a vanilla transformer model (TF) based on Radford et al. (2018) to approximate (A). For given input X E S where S represents the sequence length and 1 an embedding dimension, the main components of a TF are to apply a positional encoding X < PosEnc(X)) and then to iterate
X レ X Attn (1) (Norm (1)
A (x))
(1)
X < X MLP(() (Norm (1) M (X))
for I = 1,..., L number of blocks. Here Norm(.) is normalization function like layer norm, Attn(-) is the attention mechanism (Vaswani et al. (017)) which requires O(S + SD²) compute and O + SD) memory. MLP(.) is applied to the embedding dimension, thus requiring O(SD)) compute and O(SD) memory.
We consider two baseline approaches, inspired by prior work, to train TFs for linear algebra tasks. These approaches mostly differ by how to translate the input matrix A to X, and their loss functions. train TF with its inputs and
The first approach, following Charton (202a)) and Charton (022b), outputs represented as strings. Here, the elements in the vocabulary V would be strings that represent floating-poit numbers like "ds E el ce..... e..e"" where 8m. represents the matissa and e, the exponent. The size of [V] is approximately R 2 10 x2x 10 (counting redundant representations), which grows exponentially as we increase the precision of the floating ponnttt-onnttto representation by increasing M and K. In this setup, the input matrix A gets flattened and each A., gets converted to an element V and then embedded in a D-dimensional space to get E RNN xD, After X is passed into equation then, assuming
(AA) = A-1, the last head layer creates Y E 1R X|| which we use to compare it to the string representation of each A_ via the cross-entropy loss. In our experiments refer to this method as STRF ormer (string transformer).
The second approach, NumF o rmer (numerical transformer), essentially removes the use of strings. That is, we flatten A and then embed each A into D-dimensional space through a linear layer
3

---

Published as a conference paper at ICLR 2026
W R lx to obtain X
R vec(A))()) and, after passing X through the TF
(equation 1. we obtain Y = XW E & x1 through a final linear layer W E Dx1 In this case we use an approximation loss between Y and f(A) such as the nuclear norm |Y f(A)ll, if
f(A) = A (with Y reshaped as N x N) or YYx= f(A) if f(A) E R.
The use of flattening allows both $ STRE ormer and NumF o ormer to run for any matrix size. However, following the discussion on equation | each TF block requires a probibitive O(N++ + +N² memory, As seen Figure [ we cannot run matrices larger than N 50 as we run out of memory.
In terms of the training data, the standard procedure in the literature is to sample random Gaussian matrices, Ar, (0, 1). This sampling strategy is used in prior work (Charton| 2022a b) (among other sampling distributions). and in all the works that study solving regression tasks with transformers
(Garg et al. 2023 von Oswald et al., 2023; Vladymyrov et al. 2024 Fu et al. 2024 Liu et al. 2025
Yang et al. 2024). By sampling new random matrices during training we never repeat data and
therefore follow a single epoch training as is done for current LLMs.
DIAGNOSING THE SHORTCOMINGS OF USING TRANSFORMERS FOR NLA
We now seek to shed light on the limitations of using transformers for numerical linear algebra tasks. We then consider several interventions in section 4 that further advance our understanding.
In this section we only show results with NumF orrerr it consistently performs order of magni- tude better than STRF o rmer, as we will see in section 4 The main reason for the poor performance of STRRormms is that the string representation of the numbers sacrifices numerical precision. Thus, NumF CumFormee rmer trained with Gaussian random matrices Ai ~ (0, a) (with o = 1 unless specified otherwise) is our main baseline and starting point. See Appendix D for additional experimental details and Appendix for for further comparisons between STRFo rmer and NumF ormer. While simple but substantial improvement on STRF o rmer, and our baseline, we note NumF ormer itself is novel and has not been proposed in prior work to the best of our knowledge.
Are transformers learning an algorithm? Prior work argues how transformers appear to learn in- context regression (Garg et al. 2023) through iterative algorithms like gradient descent (von Oswald al. 2023) or Newton's 's method (Fu et al.] 2024). The argument von Oswald al. (2023) that simplified transformer architecture has set of parameter combinations that represent the gradient descent algorithm and, empirically, as we traverse the blocks equation
the approximations
become better similar rate Newton's method (Fu et al. 2024
In our linear algebra context, our model has learned an algorithm to solve the task, could then pass any matrix A that is not a random Gaussian and expect a reasonable solution. Throughout our paper, we will evaluate trained models two sets of such OOD matrices S and M. The first set S contains canonical structured matrices such identity I, diagonal D with Da.; ~ (0, 1), and Toeplitz with the bands sampled from N (0, 1). Even though the entries of the structured matrices in S are also sampled from a N/ (0, 1). their spectra are quite different, as we show in Figure [ The second set M contains more than 100 non-random matrices from the matrix marker (Boisvert et al. 1997). These matrices are gathered from challenging real-world applications like finite element approximations, structural engineering, fluid flow, power system networks and many more (see appendix B.2).
10
Test Rel Etr
Id
10a
MM
Ga
10-1
2
4
layers
For Figure 2 we train 6 layered NumF (umForrme ormer model on 20 x 20 Gaussian symmetric matrices Ar,; ~ (0 1) to calculate the inverse NN (A) R 1 We compute the relative error using the nuclear norm rerr( (A) = A NN, A) IlI. /| after each block of the transformer model (equation T]. We test on three different matrices: (1) the distribution of matrices that the model was trained on, Gaus- sian symmetric matrices, (2) the identity A = and (3) on all the matrices M E M, computing 1
The error does not decrease as we traverse through the transformer blocks. We plot the relative error the in- verse A -1 task of NunF ormer model for Gaussian symmetric matrices (Ga), for the Identity (Id) and the mean of the matrix market ensemble (MM).
Figure 2:
[M EMEM rer. In contrast to
4

---

Published as a conference paper at ICLR 2026
5
10
N
6
10 1
-10
MSE
10-3
5
2
ElAmar]]
TF
10-5
min
0
1.0 -10
0.0
0.5 masking
HP SF
NF
Figure 3: Unexpected failure cases for transformer methods for NLA. Left: We compute the Amax for distinct symmetric Gaussian matrices A.j ~ (0, 1) as we mask their inputs with different proportions of non-zero values (green). compared to the output of NumE ormer model trained the same symmetric Gaussian matrices (orange). As we incrcase the proportion of zeros the input, the NumFo o rmer models goes from predicting E Amax] without masking to min Amax- Middle: In the top panel, we compare the full eigenspectrum of a symmetric Gaussian matrix Ar., ~ NN 1) (grey) to the cigenspectrum prediction for the same matrix for NumF oreer model trained on the same distribution of symmetric Gaussian matrices (green). The NumF ormer learns linear decay in the spectrum that closely matches the training data. In the bottom panel, we compare the eigenspectrum prediction a NumFormer but now for a symmetric Toepliz matrix with Gaussian bands A.,, ~ (O, 1) (orange). The NumF o ormer erroneously predicts the same linear decay for the Toepliz input. Right: We display the MSE loss of 3 different methods (Liu et al., [2025 Charton [222a by for training transformers to solve least-squares problems when pass an in in-distribution Gaussian matrix (green) versus passing the identity matrix (orange). There significant performance degradation when evaluating on canonical matrices like the identity and STRFormer failed decode for See section 3 for full details.
von Oswald al. 2023) and Fu et al. (2024), we do not observe monotonically decreasing error as we traverse the different layers in the transformer. This finding suggests that the NumF o ormer model at least not learning an iterative algorithm solve the task. Moreover, we see the NumF o rmer has high error for the OOD matrix market, and even worse for the trivial identity matrix! We revisit this observation shortly, in validating with canonical matrices (below).
Transformers as Statistical Machines From random matrix theory, we know that a symmetric Gaussian matrix Ar, No (0, o²) with Ai, = Az, has the bulk of its eigenvalues A (A) in the interval (N 20 / N), and E [Ammxx (A) //VVVV + 2a as N + 0o Tao 2012). In Figure (left) we train NumF o rm r model on symmetric 20 20 Ar, (O, 1) matrices to solve >max- We then probe its predictions on set of masked inputs M C A where M.. 1} and vary the overall proportion of nonzero entries. That is, we go from the zero matrix 0
A (masking proportion 1) to A = 1 @ A (masking proportion 0). Note how the NumF o omer model assigns to the zero matrix an estimate of the min Amax (A) most likely obtained through training. In other words, NN₂ (0) x min(A (A (grey dot and line), which is much larger than Amax (0) 0. We compute our estimate with B = 10K.
In Figure 3 (middle) we train a NomForme r model on symmetric 20 x 20 Ar, ~ N(00 1) matrices to predict the whole eigenspectrum A(A) E 22 In the top panel we have the prediction of the model NN,(Ax) and the actual spectrum X (Ax) for AN (O, I) and A = Av- On the bottom panel we have the prediction of the model and the actual spectrum but for Toeplitz structured matrix As E S. The NumForme r learned through training that linear decay a close approximation to the X Ax) and it erroneously predicts linear decay also for As.
Overall, we found these results alarming as the transformer models appear to be learning statistical properties of the random matrices that they were trained and not an algorithm to solve the linear algebra operations. Similar results can be obtained when Arj is sampled from a Laplace or uniform distribution (see Appendix D.I).
Validating with Canonical Matrices It is common to sanity check an NLA algorithm on ma- trices like the identity A I, zero A O, diagonal A = diag(d) d), as these are considered trivial matrices with closed-form solutions. In Figure 3 (left), we show the discrepancy between the test loss achieved in-distribution versus the loss obtained by passing A I as the input
5

---

Published as a conference paper at ICLR 2026
for distinet transformer models for least-squares problems || bl1-- Namely, NumE ormer
STRFormer (Charton 2022a b), and the high-precision solver (HPS) in Liu et al. (2025).
We train our models following the setup in Liu et al. (2025) which is also almost identical to Charton (2022a) and Vladymyrov et al. (2024). That is, A E 20x6 A.j N(, 1). b E IK",, and by
N where we cap the condition number x(A) 5. In Figure 3 (right) we see how all the methods struggle to produce reasonable outputs for
Test Rel Error
10a
On the one hand, is OOD since matrices close to I have low probability density under A., ~ (0, 1). Yet. it is surprising we believe that the transformer models are learning a reasonable algorithm, as I should be trivial example.
10-1
On Different Matrix Sizes clear benefit of transformers is that they can handle different sequence lengths, here different matrix sizes, without any additional intervention. When training a LLM such as GPT-2 (Radford et al. 2018), we can create input batches of the same context length by collecting sequences of different lengths and denoting the length difference using the <EOS> token. When training transformers for linear algebra problems, it is not clear how accommodate for different matrix sizes in the same batch, the transformers models in the literature are trained with a fixed size. In Figure 4 we train NumF ormer model on 50 50 Gaussian symmetric matrices Aij (O, 1) on the Tr(A) task. We measure the relative error when passing as input Gaussian symmetric matrices but of sizes 5 x 5. 10 x 10, 20 20, 30 x 30, 45 45 and 50 50. Even reduction from N 50 N 45 worsens the performance almost 10 times.
Matrix Size
Figure 4: The approximation error exponentially increases for smaller matrices.
4 THE RANGEFORMER METHOD
RangeFormer
NumFomer
StFoomer
Amss!A
Tr(A)
10
1.0
16
Relative (MM)
10
0.9
Relative Eror (MM)
10"
0.8
logl4]
A-1
30'
0.8
Base
Loop
Rango
Altn
Dist
16'
0.6
15
Matrix
Figure 5: Our data and architectural interventions significantly improve performance. (Left): Intervention breakdown and its effect on improving the test relative error on OOD matrix market data for the A task for 20 x 20 matrices. Our baseline is the NunFormer model trained Gaussian random matrices. and the interventions (Range), (3) incorporating subquadratic attention mechanisms (Art)) (4) using our data distribution (Dist). (Right): Performance of kangel rerr against STииогме- (Charton 2022a and NumFo rmer on Amax (A), Tr(A), log AA| and A- across mairix market data for sizes 5 x 5, 10 10 and 20 x 20. We expand upon the different model configurations and loss functions in appendix D.7 The missing numbers for STRFormer decoding emrrs.
(1) adding iterative bias (Loop), (2) embedding through the range of the operator
We now construct, step-by-step our RangeF rmer architecture, and introduce our synthetic data procedure. We will study the effects of each intervention classical linear algebra problems such as max eigenvalue Amax (A), trace Tr(A), ) log determinant log([A) and inverse - (additional details in Appendix D.2

---

Published as a conference paper at ICLR 2026
Figure 5 (left) shows the cumulative and substantial performance increase from all the interventions, previously summarized in Figure II that transform NumFormer to Range or In Figure 5 (right) we compare the performance of RangeFormez across different matrix sizes and linear algebra operations. For each linear algebra operation, we train each model using our batching technique (discussed below) that enables the models to perform for different matrix sizes and avoids the lack of size generalization discussed in Figure Overall, Range orrer is an order of magnitude better when testing our OOD matrix market ensemble. For STRFormer, we use the code and checkpoints provided by Charton (2022a b). Interestingly, the checkpoints perform well when ran on in ni-distribution random matrices but they usual fail to decode correct string when ran on the OOD mulrix market data. We now explain each intervention in detail.
Positional encodings Though not particular intervention in Figure introducing positional encodings is vital for our application as the results of linear operations like trace, solves, or max eigenvalue are not permutation invariant. For example, permuting the identity changes the value of the trace. In appendix D.3 we compare the performance of the different encoding methods.
Looping To introduce an algorithmic bias into our transformer model, and avoid the situation depicted in Figure 2 we tie the transformer weights in equation to obtain looped transformer (Saunshi et al. 2025] Giannou et al., 2023; Yang et al., 2024, Geiping et al. 2025). Via looping, now the model is forced to use the recurrences to improve the solutions, similar to how most iterative linear algebra routines work, such as linear conjugate gradients, or stochastic Lanczos quadrature (Saad 2003). In Figure 6 (left), we compare the quality of the approximation of the transformer model as we traverse the layers in equation | on a OOD Toepliz matrix for the A- task. The NumF orer model does not improve its approximation as we traverse the layers. In contrast, our RangeF o rmer model monotonically decreases the error as we traverse the layers, more sharply in the beginning and flattening at the end. Additional experimental details are in Appendix [A.2]
Range embedding This intervention circumvents the ONN + 22D) scaling of the model by not flattening the matrix. Rather, we learn projection matrix I E N that allows us work directly with the range of the operator as X AT E NxD by now considering the rows of X as the sequence dimension and the columns as the embedding, scaling as O(N2 + ND). To underscore the relevance of this modification, we name our transformer models as RangeF (left) shows a minor performance improvement in substituting to this approach despite a reduction in total parameter count. See Appendix D for additional details on the architecture.
Figure S
Attention alternatives Giannou et al. (223) and Liu et al. (2025) argue that the scale normalization and the nonlinear functional form of the standard softmax attention introduces approximation distortion that is not favorable to linear algebra primitives like matrix-multiplies. Thus, we use linear attention alternatives like Taylor attention (Arora al. 2023) BaseConv (Liu et al. 2025) that allow us to improve upon standard attention and, most notably, achieve additional scaling improvements of O(ND + D²) by avoiding the construction of the full attention matrix. Figure | (c) shows the performance improvement captured by this architectural intervention (Arch) (among others). Moreover, Figure (d) shows the staggering runtime and memory benefits achieved through this intervention. We can now train transformers on matrices of size N = 1K in less time than prior methods on N = 50 as seen in Figure (d). Additional large scale experiments in Appendix D.I
Training data distribution As discussed in section 3 a notable limitation of previous methods is to solely train on Gaussian random matrices (Liu et al. 2025 Charton 2022a b) Dutta and Sra 2025; Vladymyrov et al.] (024)) as transformer models would then only learn statistical properties about the data. As seen in Appendix B.3 Figure g different Gaussian matrices almost have an identical eigenspectrum, making them essentially the same linear operator. We take different approach and construct a data mixture consisting of two main elements: structured matrices and matrices with different eigenspectrum decays. To generate structured matrices we use the continuous Einsum parameterization from Potapezynski et al. (2024) (discussed in Appendix B) which allows us randomly sample structures like Kronecker, Low-rank, Tensor Train, Block Tensor Train and Monarch, and many more. Interestingly, our way of sampling from the Einsum parameterization prevents us from sampling matrices like I, O, diagonal, or Toeplitz (Appendix (B. Furthermore, in Appendix B.3] Figure g we see how the different structures lead to diverse eigenspectra, even when the elements of these matrices are still sampled from /O. 1). Finally, as discussed in Appendix we use a series of functional forms to create a wide variety of eigenspectra A as seen in Figure Once we have A, then we sample a random basis Q and create A = Q AQ. Figure (Right)
B.3
10
7

---

Published as a conference paper at ICLR 2026
10a
MM
No
mix
10
WC
ld
10a
10
Test Rel Error
OOD Rel Err
Loss
10-1
10-3
100
10"4
10-5
10-2
10
2000 Steps
4000
40 Matrix Size
20
ld
Ze
To
MM
Figure Different benefits of training with our methodology. (Left): Curriculum learning enables our models to learn problems at larger scales. Train loss of RangeFormee for the linear solves task in (Liu et al. 2025). (Middle): We compare the performance on the same two OOD cases, of our RangeF oeeeFormee orer model across different matrix sizes. The RangeF Former model was trained using batching method on 50 50, 30 and 10 10 matrices but can perform relatively
20 and 45 45. (Right): Incorporating a diversity of structures, as well as matrices with a wide variety of eigenspectra decay, in the training distribution highly improves OOD performance. We show the relative error on different OOD matrices for Rangel model trained Gaussian random matrices and one trained (mix) for Amas. The matrices are the Identity (Id), zeros (Ze), Toeplitz (To) and the matrix market ensemble (averaged).
other cases
data mixture
shows the benefits of transitioning from training RangeFormes model on Gaussian random data, to one trained on our data mixture for the inverse A task. Also Figure (c) displays the benefits of training with this data mixture on other linear algebra operations.
Curriculum learning and varying sequence length At the N = 1K scale, particularly for solving linear systems, we could not get our RangeFo ormer models to converge unless we first got a model checkpoint that solves N 100 and then that checkpoint as starting point to solve N 1K. In Figure [6] (Left) we see how only when doing the curriculum would the loss start to decrease. To avoid the situation depicted in Figure we train models with diverse set of matrix sizes (N1,.., NR). Figure 6 (Middle): shows how our model can stabilize its OOD performance across different sizes, even ones that
seen during training.
DOWNSTREAM APPLICATIONS
Ultimately, matrix operations rarely end in and of themselves but instead are used part of downstream application. For the first time, we replace linear algebra routines with transformers to compute the matrix operations in downstream applications.
5.1 ACCELERATING THE CONVERGENCE KRYLOV SUBSPACE METHODS
Krylov subspace methods are classical iterative methods for computing matrix operations, such as conjugate gradients (CG) for solving linear systems. As discussed in appendix [. any Krylov subspace algorithm requires an initial vector to start forming the basis (x() A). We now use our trained RangeF r models NNe(-) to provide 2o NN, (A). The intuition is that since NN, provides good approximation to the solution, then the overall effect is that the Krylov subspace algorithm will converge the solution faster when compared to simply picking a standard random or zero initialization for Zo- In Figure 7 (Right) we show how CG converges much rapidly on the bessık02 matrix from the matrix market. We provide additional results for other algorithms and other matrices in appendix D.4, Since matrix operations are fundamental primitive for learning tasks, our goal in this paper is to expose limitations of transformers for linear algebra, and show how they can be improved, rather than to argue that they preferable to classical solvers. However, this experiment demonstrates that there is indeed promise for transformers to combined with classical solvers for better performance than these solvers could achieve alone.
8

---

Published as a conference paper at ICLR 2026
10
0.35
true
2.5
10-²
0.30
rsvd
rfsvd
2.0
||xx -bll2
10
0.25
1.5
10-6
0.20
1.0
10-8
10-10
cold
0.15
0.5
true
warm
0.0
0
5
10
50
CG iterations
index
iters
Figure 7: Performance RangeFormer on downstream applications. (Left):: CG convergence with and without RangeF ormer warm starts on besskk02 from the matrix market. (Middle): Spectrum approximation for RSVD sampling Q ~ N (0, I) and sampling f ~ (A) (Right): The RangeF ormer linear algebra primitives allow recover the length scale parameter of RBF kernel for data generated from Gaussian process with an RBF kernel.
5.2 IMPROVING THE SAMPLING DISTRIBUTION FOR RNLA
All randomized linear algebra methods rely sampling some random embedding Z usually Gaussian 52 ~ / (O,I). In this application, we train RangeF ormer models that define p(S2) = (N)(A) such that our samples Ω p(22) improve the performance of the randomized algorithm. In other words, NN, will take into consideration the characteristics and structure A and bias p(2) to provide better sampling distribution. Since RangeFo rmer uses looping, the process of sampling Ω is similar to diffusion, where we start with Zo N (0, I) and progressively pass this noise through the layers of NN₂ to get Ω NN₂(A) (A). Figure (Middle), we compare the performance of randomized SVD using RangeF o rmer, namely RF SVD, and the usual algorithm RSVD that uses Ω V (0, I). We use matrix E R 1100x stemming from a Gaussian Process RBF kernel on a uni-variate grid that has fast decaying spectrum. As seen in the figure, RFSVD (with D = 16) is able to more closely resemble the spectrum of A when compared to RSVR (rank 16) approximation (appendix D.5 has additional details). Moreover, Figure 15| (Right) also shows contour plots of the learned distribution p(52) which are quite distinct from the contours of (0, I) Figure 15 (Left).
5.3 GAUSSIAN PROCESSES
In our last application, we RangeFormer train Gaussian process (GPs) kernel hyperparame- ters. This is a particularly challenging stress-test, as the kernel learning pipeline requires a sequence of matrix operations that build off one another, where compounding errors and biases could preclude convergence.
GPs are flexible distributions over functions, with properties controlled by kernel (Rasmussen and Williams, 2006). Given some data (2,, D:) / and kernel function (z(xx²)) parametrized by Ø. training GP requires constructing K2 [ = h (x,, Z.y) and then minimizing the negative log-likelihood L (+) log K + T K d y which involves and solves (Rasmussen and Williams 2006). our case, we use RangeFo o rmer r for these operations in order to learn ø. In Figure 7 (Right) we perform a sanity check showing that our RangeF ooeer pipeline recovers the true length scale when trained on data generated from a GP. Finally, we fit RBF kernel on data generated y, sin(2Tx;)) + 6,, for and 2; ~ [4(0, 1), using our RangeF o rmer pipeline and using Cholesky. Both strategies give almost identical results the hyperparameter values learned, their trajectories, and the test RMSEs: 0.87996 Choleskyy and 0.87978 (RangeFormer) We provide additional experimental details in appendix D.6
computation of log-determinants
1,... ..... 50 with
(00 0.04),

---

Published as a conference paper at ICLR 2026
6 DISCUSSION
We have found that current transformer approaches for matrix operations profoundly fail on even trivial out-of -utto--distribution matrices, such the identity matrix. We shed light on these failure modes, showing that current approaches perform statistical interpolation on in-distribution matrices, rather discovering more generalizable algorithms. We also find that current approaches are computationally intractable, beyond even tiny 30 x
30 matrices.
Through sequence of interventions, we found that introducing a learnable projection improves scaling by several orders of magnitude, and new richer training distributions improve performance by several orders of magnitude, especially in OOD settings. These interventions are not merely making transformers better at numerical linear algebra, but moving from system that was essentially broken, one that can start perform competently on challenging real-world matrices, and shifting closer to algorithm discovery.
able profoundly relax these limitations. In partieular,
This research area is in its early stages. Our goal is not compete with purpose-buil classical algorithms, just like transformer could not be expected to compete with a calculator for arithmetic. At the same time, we also provide the first results using transformers successfully for matrix operations in downstream applications, which suggests that transformers could become practical in this setting. For example, GP marginal likelihood optimization is particularly challenging, since is an iterative pipeline that would quickly become unstable under compounding errors. Moreover, our randomized linear algebra application also shows how transformers can play complementary role to classical methods, learning how to improve randomized SVD. And our warm-start CG application additionally shows how transformers can naturally be combined with classical methods.
We hope that our results will inspire research into transformers for matrix operations. Presently, this research area is in its infancy, relative to decades of work on purpose-built classical solvers, representing millions of human hours of effort. But with sustained program, over years, will see both more general purpose architectures, and transformers that become practically compelling for matrix operations. Indeed, after years of work on neural graph methods, which initially were outperformed by classical approaches, there were eventually significant breakthroughs.
REFERENCES
Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable Second Order Optimization for Deep Learning. Preprint arXiv 2002.09018v2, 2020.
Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré. Zoology: Mcasuring and Improving Recall in Efficient Language Models. arXiv:2312.04927, 2023.
Ronald Boisvert, Roldan Pozo, Karin Remington, Richard F. Barrett, and Jack J. Dongarra. Matrix market: a web resource for test matrix collections. In Proceedings of the IFIP TC2/WG2.5 Working Conference on Quality of Numerical Software: Assessment and Enhancement, page 125-137, GBR, 1997. Chapman & Hall, Ltd. ISBN 041 2805308. URL https: patpss:/mmattin htt///athhh. n st gov/
Matrix MatrixMarket//
François Charton. Linear algebra with transformers. Transactions in Machine Learning Research (TMLR), 2022a.
François Charton. What is my math transformer doing? Three results on interpretability and generalization. arXiv:2211.00170 2022b.
Marco Cuturi. Sinkhorn Distances: Lightspeed Computation of Optimal Transport Advances in Neural Information Processing Systems (NeurIPS), 2013.
Tri Dao, Albert Gu, Matthew Eichhorn, Atri Rudra, and Christopher Ré. Learning Fast Algorithms for Linear Transforms Using Butterfly Factorizations. International Conference on Machine Learning (ICML), 2019.
Tri Dao, Beidi Chen, Nimit Sohoni, Arjun Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher Ré. Monarch: Expressive Structured Matrices for Efficient and Accurate Training. International Conference Machine Learning (ICMI.), 2022.
10

---

Published as a conference paper at ICLR 2026
Sanchayan Dutta and Suvrit Sra. Toward generalizable learning of all (linear) first-order methods via memory augmented Transformers. arXiv:2410.07263, 2025.
Marc Finzi, Andres Potapczynski, Matthew Choptuik, and Andrew Gordon Wilson. Stable and Scalable Method for Solving Initial Value PDEs with Neural Networks. International Conference on Learning Representations (ICLR), 2023.
Daniel Y. Fu, Tri Dao, Khaled K. Saab, Armin W. Thomas, Atri Rudra, and Christopher Ré. Hun- gry Hungry Hippos: Towards Language Modeling with State Space Models. Preprint arXiv 2212. 14052v3, 2023.
Deqing Fu, Tian-Qi Chen, Robin Jia, and Vatsal Sharan. Transformers Learn to Achieve Second- Order Convergence Rates for In-Context Linear Regression. Advances in Neural Information Processing Systems (NeurIPS), 2024.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What Can Transformers Learn In-Context? A Case Study of Simple Function Classes. arXiv:2208.0 2023.
Jonas Geiping, Sean McLeish, Neel Jain, John Kirchenbauer, Siddharth Singh, Brian R. Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, and Tom Goldstein. Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach. arXiv:2502.05177 2025.
Angeliki Giannou, Shashank Rajput, Jy yong Sohn, Kangwook Lee, Jason D. Lee, and Dimitris Papailiopoulos. Looped Transformers as Programmable Computers. arXiv 2301. 13196, 2023.
Micah Goldblum, Mare Finzi, Keefer Rowan, and Andrew Gordon Wilson. The no free lunch theorem, kolmogorov complexity, and the role of inductive biases in machine learning. In International Conference on Machine Learning, 2024.
Gene H Golub and Charles F Van Loan. Matrix Computations. The Johns Hopkins University Press, 2018. Fourth Edition.
Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Arizzadenesheli, Kaushik Bhattacharya, Andrew Stwart, and Anima Anandkumar. Neural Operator: Learning Maps Between Function Spaces. Preprint arXiv 2108.0804400 08481v3, 2021.
Chunyuan Li, Heerad Farkhoor, Rosanne Liu, and Jason Yosinski. Measuring the Intrinsic Dimension of Objective Landscapes. International Conference Learning Representations (ICLR), 2018.
Jerry Liu, Jessica Grogan, Owen Dugan, Ashish Rao, Simran Arora, Atri Rudra, and Christopher Ré. Towards Learning High-Precision Least Squares Algorithms with Sequence Models. International Conference on Learning Representations (ICLR), 2025.
James Martens and Roger Grosse. Optimizing neural networks with kronecker factoord approximate curvature. In International conference on machine learning, pages 2408-2417. PMLR, 2015.
Per-Gunnar Martinsson and Joel Tropp. Randomized Numerical Linear Algebra: Foundations & Algorithms. arXiv 2002.01387v3 2020.
Eric Nguyen, Karan Goel, Albert Gu, Gordon W. Downs, Preey Shah, Tri Dao, Stephen A. Baccus, and Christopher Ré. S4ND: Modeling Images and Videos Multidimensional Signals Using State Spaces. Preprint arXiv 2210.06583v,, 2022.
I. V. Oscledets. Tensor- Train Decomposition. SIAM Journal on Scientific Computing, 2011.
Ethan Perez, Florian Strub, Harm de Vries, Vincent Dumoulin, and Aaron C. Courville. Film: Visual reasoning with a general conditioning layer. In Proceedings of the Thirty-Second AAAI Conference on Artificial Intelligence, (AAAI-18), the 30th innovative Applications of Artificial Intelligence (1AAI-18), and the 8th AAAI Symposium on Educational Advances in Artificial Intelligence (EAAI- 18), New Orleans, Louisiana, USA, February 2-7, 2018, pages 3942-3951. AAAI Press, 2018.
Andres Potapczynski, Mare Finzi, Geoff Pleiss, and Andrew G Wilson. Cola: Exploiting com- positional structure for automatic and efficient numerical linear algebra. Advances in Neural Information Processing Systemx, 36:43894 43917, 2023.

---

Published as a conference paper at ICLR 2026
Andres Potapczynski, Shikai Qiu, Mare Finzi, Christopher Ferri, Zixi Chen, Micah Goldblum, Bayan Bruss, Christopher De Sa, and Andrew Gordon Wilson. Searching for Efficient Linear Layers over Continuous Space of Structured Matrices. Advances in Neural Information Processing Systems (NeurlPS), 2024.
Shikai Qiu, Andres Potapczynski, Mare Finzi, Micah Goldblum, and Andrew Gordon Wilson. Com- pute Better Spent: Replacing Dense Layers with Structured Matrices. International Conference on Machine Learning (ICML), 2024.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving Language Under- standing by Generative Fre-Training OpenAl, 2018.
Carl Edward Rasmussen and C. K. Williams, Gaussian Processes for Machine Learning. MIT Press, 2006.
Yousef Saad. Iterative Methods for Sparse Linear Systems, SIAM, 2003.
Yousef Saad. Numerical methods for large eigenvalue problems. SIAM, 2011.
Nikunj Saunshi, Nishanth Dikkala, Zhiyuan Li, Sanjiv Kumar, and Sashank Reddi. Reasoning with Latent Thoughts: On the Power of Looped Transformers. International Conference on Learning Representations (ICLR), 2025.
Terence Tao. Topics in Random Matrix Theory. American Mathematical Society, 2012.
Lloyd N. Trefethen and David Bau. Numerical Linear Algebra. SIAM, 1997.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention All You Need. Advances in Neural Information Processing Systems (NeurIPS), 2017.
Max Vladymyrov, Johannes von Oswald, Nolan Miller, and Mark Sandler. Efficient Linear System Solver with Transformers. Al MATH Workshop International Conference Machine Learning (ICML), 2024.
Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. International Conference Machine Learning (ICML), 2023.
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self- Self-AAteenion with Linear Complexity. arXiv 2006.04768, 2020.
Liu Yang, Kangwook Lee, Robert Nowak, and Dimitris Papailiopoulos. Looped Transformers are Better Learning Learning Algorithms. International Conference on Learning Representations (ICLR), 2024.
APPENDIX OUTLINE
Appendix provides additional background on Linear Algebra problems and methods.
2. Appendix describes our data construction procedure, evaluation metrics, and provides an analysis of the the eigenspectra of different data distributions.
3. Appendix describes the proposed RangeFo rmer architecture in detail.
4. Appendix provides additional experimental details.
5. Appendix formally analyzes the runtime for NumF ormer and RangeForme::
A
BACKGROUND
In this section we briefly introduce and describe the techniques, methods and problems that will be used throughout the paper.
12

---

Published as a conference paper at ICLR 2026
A.1 LINEAR ALGEBRA
Canonical Linear Algebra problems Assume that we have A E RNN, Arguably the most important linear algebra operations, besides matrix-veetor-multippe
(MVM) Ax. are: (1) linear solves x: b. (2) eigen decompositions A = QAQT (assuming A is symmetric), (3) applying a function f to A, as f(A) = ( or (4) applying a matrix to scalar function like trace Tr(A), determinant |A], or condition number <(A). The previous list is far from exhaustive but encompasses the majority of the operations that are common place in ML. For example, linear solves and determmannn computations are part of Gaussian Process (GP) training (Rasmussen and Williams, 2006) while eigen- aggn-ddccompositions and matrix functions (AA) are part of approximate second-order optimizers like Shampoo (Anil et al. 2020). In section 4 we will train transformers to tackle all of the aforementioned operations.
Linear Algebra methods There is vast set of algorithms to solve all the linear algebra problems from above (Trefethen and Bau 1997, Golub and Loan, 2018). In this paper, we will focus on two broad families of methods that are the most commonly used in practice: (1) Krylov subspace methods (Saad) 2003 201T) and (2) randomized methods (Martinsson and Tropp 2020).
The main idea of Krylov subspace methods is to iteratively apply MVMs to construct a subspace that can be used to solve a target problem. Namely, given an initial vector Vo. and a number of iterations f = 1,..., N. we will implicitly construct and utilize the Krylov subspace K(() (æo. A) = span Aæo.... ..... A'-1 .. Based on the previous equation, the cost of these methods can thus be upper-bounded by O (TN') by assuming the worst case MVM cost for A and where T is the number of iterations. To give some examples, some well-known Krylov subspace methods for linear solves are CG, MINRES, and GMRES (Saad) 2003) and for eigenvalue problems we have the Power Tteratiii i on, Lançzos, LOBPCG, and A Arnoldi (Saad 2011). In section we will train transformers to suggest an initial 2o that will accelerate the convergence of these methods.
The key insight of randomized methods is to reduce the problem size by multiplying A with a random matrix Ω E R x where é < N and then apply distinct deterministic methods to Y = An R to solve target problem. A famous example is RSVD (Martinsson and Tropp 2020). where we extract
an orthonormal basis Q = QR(Y) (Y), then construct C' C= Q A E R'X get U Σ, V] SVD(C) and thus provide an approximation of the form A z (QU)VVT where U E R EXE V E RR are orthonormal matrices, and Σ e cx is a diagonal matrix. With RSVD, we avoid the O(N3) cost of applying SVD to A but rather incur a ( N) cost of the SVD to C. Most commonly, Ω is sampled s N (0, I) but there are also other more efficient alternatives such as sparse sign matrices or subsampled trigonometric transforms (Martinsson and Tropp 2020). However, the overall decision of the sampling distribution is either speed or convenience and, in section, 5.2 we explore other learnt distributions that our transformer models uncovered to improve performance.
Matrix structures By matrix structure we refer to various patterns and relations between the entries of A that ultimately lead to more efficient MVMs and more efficient algorithms to solve distinct linear algebra problems. A trivial example would be when A is diagonal, that is Ai, = 0 if i + j. In this case, a MVM simply amounts to Aæ = diag O œ, a (NN) operation not (N Additionally, a linear solve would also be O(N), since diag -1 c b. A popular example is when A is Toeplitz. In this case, the PFT is used to efficiently compute a MVM in O (N log N). In section 4 we will make extensive use of different matrix structures pre-train the transformers. Specifically, we will use the Einsum parameterization of matrix structures from Potapczynski et al. (2024). The Finsum parameterization is able capture many well-known structures such as low rank, Kronecker, Tensor Train (Oseledets) 2011), Monarch (Dao et al. (2022). BTT (Qiu et al. (024) and more by representing the structures as tensor contractions between two factors. All that required in the Einsum parameterization is to specify a real valued vector & E [0, 1]" and the values on the entries & dictate the structure of the matrix.
A.2 TRANSFORMERS
Architecture Our architecture is inspired and resembles Radford et al. (2018). That is, for a given input X E RSD, where S represent the sequence length and 72 an embedding dimension, we will
13

---

Published as a conference paper at ICLR 2026
first apply a positional encoding X < PosEnc(X) and then iterate
X <- X + Attn (NormA(X)
X + X + MMLP (Normy/ (N (X))
for 1 = 1, 1,..., I number of blocks. Here Norm(.) is some sort of normalization function like layer norm, Atin(-) refers to the attention mechanism that is detailed in the next subsection and MI.(-)) an MLP applied to the embedding dimension, thus requiring OSSD²) operations.
We depart from Radford et al. (2018) in several ways, First, our transformer not a text to text model. At a high level, representing numbers as strings as Charton (2022a), has two clear drawbacks (1) the vocabulary grows exponentially in the precision and (2) it adds redundant, lossy transformations like representing numbers as strings to then embed them as numbers again. In terms of the vocabulary, if we have a string: "#$1 s - 61 Cg" then the size of the vocabulary would be approximately [V] = 2 x 10 2 10 (counting redundant representations). Second, our transformer model does not apply a causal mask as the sequence length is related to the input matrix where we have complete access. Third, our training loss not next token prediction but rather approximation error. Arguably the deviations from Radford et al. (0118) make our models completely different in nature. However, we are not trying to train a LLMs to do linear algebra, rather we are training neural networks to solve linear algebra tasks and transformer architecture is a natural choice as it can gracefully handle different sequence lengths (matrix sizes).
Attention mechanisms Following the classical attention mechanism from Vaswani et al. (017)) we construct (Ak) = -WWQ(h) (() X XWK((t) and V(h) = -WWV with WQU) (K(h), WV
Dxd where d represents the head dimension with Hd = D and h ..... H indexes the head number. For each head, we then output y softmax (M (Q(N) ( and finally concatenate the output Y concat( (1))
()) E RSDD This attention mecha- nism requires O(S D) operations for computing Q ((h) (R)T for each h.
There have been several efforts to reduce the quadratic requirement of the standard attention mech- anism, most notably the use of linear attention (Wang et al.) (020) where the goal is circum- vent the construction of softtax(( K (+)T) by instead applying feature map (-)) such that (1) (Q (1b( K (1))TV)) and therefore only require O(SD²) operations. Two popular feature
((b(K maps are (x)) 1+ + clu(z) and (xx) = + x + x (Arora et al. 2023).
Looped transformers The idea of looped transformers GGannoou et al. 2023; Yang al. 2024 Saunshi et al. 2025: Geiping et al.) 2025) is to incorporate an algorithmic inductive bias into the transformer architecture by sharing weights across blocks. Namely, Attn (1nn)) (.) = Attn(.) and (L) MLP(-) for all ( = ......L. Interestingly, this inductive bias allow us to increase the amount of computation (by extending the number of loops) and hopefully improve predictions.
In our models we add the regularization techniques from Geiping et al. (0255 of injecting input noise at each iteration as well as sampling randomly the number of blocks to run each iteration during training. We used looped transformers as they are the model class that gave us better empirical results but also because we want our transformer models to learn iterative algorithms.
B DATASETS
B. TRAINING DATA
Previous results transformers for NLA have used random matrix families to train their models. However, random matrices exhibit certain statistical properties which may make tasks like Amax (A) predictible from those statistics rather than required an algorithm to be learned. To avoid this pitfall, we train transformers on a distribution of structured matrices with a rich variety of eigenspectra, making this data less predictible from statistics only.
Base. Our baseline training dataset simply consists of square Gaussian matrices, Aij ~ . (0, o²). Usually, we would symmetrize the input A., = Ais to have well- well-defined data for the problem.
14

---

Published as a conference paper at ICLR 2026
Einsums. The Einsum parameterization of matrices described in Potapczynski al. (2024) captures structured matrices such as Kronecker, Low-rank, Tensor-train and many more. To sample a matrix of size N x N from the distribution of Einsums, give the following sampling algorithm.
Enumerate triplet factorizations, T {(N,, N₂, E N,N,, N). Sample (Na, Ne, N₂) and (N&, Ne, No) from T. Sampling can be done uniformly or biased to factorizations with more 1s as that induces more varied structures.
2. Uniformly sample the rank Np from the discrete interval [1, min min (Na, N.) + +1].
Construct tensors E NR,, XNNNN,N, E ReMxN, from standard normal distribution (0, 1). Together, these two matrices define linear operator A.
4. Let I be the identity where Finally, compute
E IR NxN, Reshape it such that I E RNNNN
A' = Σ
Brecpp X Xarōóp
ag7p
Now, A' is of shape (N, N8, , Ne, No). Reshape it to size N x N to recover A. In Figure 9 we observe the diverse eigenspectra obtain through the einsum sampling.
Diagonal Decays. To equip the training dataset with wide variety of eigenspectra, we directly sample various diagonals and then construct matrices from those diagonals. The experiments done in Figure 5 use uniform mixture over four different diagonal decay distributions. For all of the decays, we evenly space our input values over the [0, 1] interval,
i.c. {z, J 24 = (i 1)/(N 1) and sample scale parameter ~ [11,, 3]. Each of the following distributions also includes a rate parameter a which controls the rate of spectral decline.
Polynomial. We sample
5). Then, we compute our eigenvalues as 10-4 A 1 xu and add jitter of 10 4 to A to avoid a zero eigenvalue.
Cosine. We sample a from mixture of uniform distributions, a 0.2 2.2-2/(0.05. 0.1] + 0.2 [[1,4 + 0.6 2/00,3, 0.9]. We then compute, X; = 0.5 + cos(T))))).
Inverse exponential. We sample 30). Let f(x) Ite Then, we compute A f(a(z,
Log. We sample ~U[1, 1000) and compute X; log(1+aurs) log(1+a)
Given eigenvalues X, can compute D diag(A). Then, we randomly sample X ~ (0. 1) and compute the QR-ccoomponition QR X. Finally, we compute QDQT sample random bases. As Q is an orthonormal matrix, it does not change the eigenvalues present in D. In Figure TO we see the diverse eigenspectra obtained by through our diverse decay functions.
For the Tr(A) and the Amm (A) tasks, given matrix A taken from of the above distributions, we symmetrize in order to have real eigenvalues and a well-defined notion of Amax (A).
For the log [A] and A- tasks,
positive definite matrices. Specifically, compute A <- AAT + eI. Practically, we use e = 10 44 to ensure that the log |A| task is well-defined while not being trivial for the transformer model.
15

---

Published as a conference paper at ICLR 2026
Gaussian
Laplace
10
-5
-10
A
Uniform
Marchenko-Paste
10
-5
-10
5
10
15
5
10
15
Figure Eigenspectra of random matrices. We randomly sample 20 x 20 matrices from the selected random matrix families (Gaussian, Laplace, Uniform, and Marchenko-Pastuu and symmetrize them with the exception Marchenko-Parn -aattrr which is symmetric by construction. We plot their eigenvalues and observe that there is very little variance in the eigenspectra within any given random matrix family. We can see how the eigenspectra of each of this families is highly regular and predictible.
16

---

Published as a conference paper at ICLR 2026
B.2 MATRIX MARKET
To accurately evaluate the performance of our models on OOD matrices. we sample several "real- world", symmetric matrices from the matrix marker as described in Boisvert al. (1997). Here are the collections that we used in our experiments:
BCSSTRUCI. These matrices represent dynamic analyses in structural engineering, con- taining stiffness and mass matrices.
BCSSTRUC2, These matrices are linear equation problems "arising from applications of the GT-STRUDL structural engineering code".
BCSSTRUC3, These matrices were collected from structural engineering packages and serve as generalized symmetric eigenproblems.
BCSSTRUC4. Symmetric eigenproblems and linear equations.
CYLSHELL. Finite element discretized octant of cylindrical shell.
LANPRO. Lanczos with partial reorthogonalization.
LAPLACE. Finite different Laplacians.
PSADMIT. Four symmetric matrices used in modeling power system networks.
All the previous matrices come in wide variety of shapes. In our experiments, if we need a size of N x N, we simply take the first N columns and the first N rows. That is, we slice the matrix as Al: N,: N] following the notation in Python. We also considered projecting them via random matrices. Assuming that No is the original size of the matrix, we then have A' ΩAN where P a No x N matrix.
Also, normalize the matrices by its largest entry A |A where |Ar (ix( |A,yl- We follow this approach for two reasons. First, we avoid having to compute the max eigenvalue or Frobenious norm for each matrix. Second, many of this matrices come in the sparse format, where we only have access to the nonzero entries, and so grabbing the largest one is convenient and fast. This normalization would not make the max eigenvalues approximately 1 but with some divergences as seen in Figure
B.3 SPECTRAL ANALYSIS
One glaring issue in the usage of random matrices training data is the homogeneity of their eigenspectra. While previous results such as Charton (2022a) train on up to 360 million random matrices for eigenvalue tasks, there may be limited utility in doing as almost all of those matrices share very similar spectral properties. In this section, we visually examine the eigenspectra of different distributions of matrices.
Figure & visualizes the eigenspectra of different random matrix families. We observe that within any given random matrix family, the eigenspectra follow the same general pattern with very little variation. In fact, the Gaussian, Laplace, and Uniform eigenspectra are all quite similar as well which could explain results like those Figure 12 where training those different random matrix families gives very similar results. The fact that the eigenspectra are homogeneous, implies that the neural network can learn that type of pattern and not an algorithm solve the task at hand.
Figure 9 visualizes the eigenspectra of various Einsums. We observe much greater variation in both the overall shape of the eigenspectra and the values of the largest and smaller eigenvalues, even though the factors of the einsum are also sampled from A00 1).
While Einsums may offer diverse eigenspectra by virtue of representing structured matrices, we can directly control the properties of the eigenspectra by specifying various functional forms. Figure TO shows the eigenspectra of various diagonal decays which are more varied than both the Einsums and random matrices. By varying the parameter, we can further control the rate of the decay and the
'The matrix market is provided by the National Institute of Standards and Technology (NIST) as service, intended for public access and use. See https://wwwnnss www.aii.www/ppee/oooggggggegggge faiiopprrigg------e---eeeeeeeeee har-iggt--------a----iiiee anndlliceniing-stttementt wontwwaa---d----eeiiilllllael publiccctiissslisioionn for further information about licensing.
17

---

Published as a conference paper at ICLR 2026
-5
10
Figure 9: Eigenspectra of Einsum matrices. We display the eigenspectra of 16 symmetrized Einsum matrices and observe that the eigenspectra of structured matrices parameterized by Einsums are more varied than the eigenspectra admitted by random matrices.
way in which the eigenspectra decays. Empirically, observe that this richer distribution results in strong performance improvements (see Figure (Left))
Finally, visualize the eigenspectra of the matrix market matrices in 11 and that they
qualitatively different from the training distribution making them a good metric for OOD performance. Since these matrices are collected from real-world engineering problems, we observe behavior which is hard to capture through either structured matrices or functional forms such as the step decay in Besstm26 and Besstm02. In essence, the matrix market spectrums are unpredictiable and erratic, making them a great test case for our models.
C ARCHITECTURES
In this section we expand on the architectures NumFormer and RangeFormer as well as compare them. For given input X RSxD, where represents the sequence length and D an embedding dimension, the main components of our core transformer backbone are to apply positional encoding X PosEne X) and then to iterate times the following:
X X Attn(Norm (X))
X + X + MLP(Normy/ (X))
Here Norm(-) is LayerNorm layer, Attn(-) refers attention mechanism, and MLP(-) is MLP applied to the embedding dimension. This last MLP consists two layers with an expansion factor of 4 similar to (Radford et al.) 2018). We discuss below the types of attention mechanisms MLP activations used. Note that in contrast to equation]] [[ equation eliminates the dependency on essentially then forcing the model to share weights across the layers.
18

---

Published as a conference paper at ICLR 2026
Poly
Cos
values
0.05
0.10
0.30
1.00
3.00
10.00
0.2
0.0
A
Inv
Log
1.0
0.8
0.6
0.4
0.2
0.0
1.0
X
Figure 10: Eigenspectra of diagonal decays. The eigenspectra of diagonal decay matrices vary wildly by construction. The parameter a controls the rate and type of decay. For example, with polynomial decay, a large a corresponds to a gentle decline that accelerates later on. With cosine decay, a controls the smoothness of the decline.
19

---

Published as a conference paper at ICLR 2026
Bcsstm26
Bcsstk16
494_bus
1.5
1.0
0.5
0.0
Bcsstk06
662 bus
Bcsstk15
1.5
1.0
A
0.5
0.0
Bcsstm02
Nos3
Bcsstm 13
1.5
1.0
0.5
0.0
15
Figure 11: Eigenspectra of some matrix market examples. The eigenspectra of randomly selected symmetric matrices drawn from structural engineering problems, finite discretizations and other real-world applications are qualitatively different from synthetic random matrix eigenspectra. We take 20 x slices from the top left corner these matrices for this figure and normalize them by their largest element.
20

---

Published as a conference paper at ICLR 2026
D/N
20
50
100
6.6 x 1.2 10-1
6.5 x
10-1
8.7 10-10011 10
32
1.7 x 10-1 10-2
5.0 x 10-1
10-1
4.7 x
10
128
1.9 x 10-1 x 10-2
2.3 X 10 1 10
2.6 x 1.. 10
512
1,8 x 10-1
10-2
2.6 x 3.2 10-2
10
2.0 2 10
1024
1.2 x 10-1 X
10
24 1
10 x 10 2
Table 1: We observe how performance changes as embedding size D and matrix dimension N scale. Small sizes like 8 and 32 limit performance as expected but with D 128, performance plateaus. If we suppose that the matrices we work with have an exponentially decaying eigenspectrum or that the majority of eigenvalues are close to O, then it may that D does not need to scale with N. Fast decaying cigenspectra or small eigenvalues precisely why randomized linear algebra can achieve a good performance with a few vectors even for very large matrices.
common characteristics in matrices of interest and is
Op
20
100
1K
A
1,8 x 10
2.3 X 10
3.1 x 10
Table 2: Error for solves as we increase matrix size. We follow the least square settings from Liu al. 2025) and track the MSE error as we increase the problem size.
NumF o omer Given the input A, we embed each Ar, into a D-dimensional space through a linear layer ( E 1 to obtain X E IR xD vec( W and, after passing X through equation obtain Y (O E - through final linear layer W E Dx1
3
For this architecture, Att((-) is the standard quadratic attention mechanism and M-P(() uses the standard GeLU activations.
(Vaswani et al. 2017
RangeFormer Given the input A, we embed each A., into D-dimensional space through a linear layer r E NxD to obtain X E R NDD which ultimately capture the range of the operation (where all the action happens). Then we pass X through equation 3 to obtain Y (W R Nx1 through a final linear layer (O) TR which mixes the column dimension.
For this architecture, Attn(-) is either the linear Taylor attention mechanism (Arora et al. the linear polynomial attention mechanism from (Liu et al. 2025). Here for MLP(-) we use ReLU activations as we found that to be marginally better than Ge
2023) or
However, now the layer I depends N, which would allow us run matrices of different sizes. We have two approaches to circumvent this problem. One approach to define a Nmax (similar to the positional encodings I to the input size N < N Nmax. The other approach is to not make F learnable but just random Gaussian embedding. This an approached inspired by how the randomized linear algebra methods operate. We found that the first approach yields the best results and ultimately used that one in our experiments. Note that this is not problem as we already have to define max sequence length for the positional encodings. However, if the user wants train model to work for all matrix sizes then then can make I not learnable.
transformer that ultimately require max sequence) and then slice
D EXPERIMENTAL DETAILS
D.1 ADDITIONAL. EXPERIMENTS
In this section we expand our statistical analysis to show that the finding from section fection[]applic appplie also when the matrices are sampled from Lapacian distribution, that is A,,; ~ L(0, 1) or from a uniform that is Aij ~ 1111, 1]. In Figure 12] we see how, regardless of the sampling distribution, the NumFo rmer r model simply learns the predictible behavior of the matrices and fails to generalize. We also describe experiments the scaling of the embedding size D with matrix size N as that relationship affects the scalability of our method. Empirically, we do not observe that large N drastically affects performance long D is moderately sized.
21

---

Published as a conference paper at ICLR 2026
10
10
10
-5
-10
-10
10
NF
10 NF 10
NF
5
S
0
<
< 0
-5
-10
-10
-10
10
Figure 12: NumF ormer learns particular eigenspectrum decays for different sampling distributions. In all the top panels, we compare the full eigenspectrum of random matrix A sampled from different distributions (grey) to the cigenspeetrum prediction for the same matrix for a NumFormer model trained on the same distribution of random matrices (green). We see how the NumF ConFormme omee Icarns a linear decay in the spectrum that closely matches the training data. In all the bottom panels, we compare the cigenspectrum prediction of a NumĒ ormer but now for symmetric Toepliz matrix with Gaussian bands Ai, ~ (N, 1) (orange). We how the NumF ormer fails to predict different eigenspectrum decay. Left: Here Aig (o 1) with As,y = A,,,, that is A is symmetric Gaussian matrix. Middle: Here A.,, ~ [(,, 1) with A., = A,,s, that A is a symmetric Laplacian matrix, Right: Here Au, ~ 21[11,11 1] with AL,5 Aç.s. that is A is a symmetric Uniform matrix.
D.2 FIGURE DETAILS
Here, we provide experimental details for various figures referenced in the paper.
Figure Data We utilize two distinct data distributions random symmetric 20 20 Gaus- sian matrices which represent an easy way to obtain synthetic training data for transformers and diagonal decays which contain more widely varying eigenspectra.
Base Data. For our baseline Base on the Tr(A) and Ammax (A) tasks, we train using symmetric Gaussian matrices where each element Ai.j NO, 1). For the log |Al and
A tasks we use positive definite Gaussian matrices. We sample A (O, 1) and then compute A + AAT +I.
Decay Data. For the Data and Arch. models, we use a uniform mixture over the polynomial, cosine, inverse, and log decay distributions explained in B.I These matrices are symmetric by construction in the Tr(A) and Amax (A) tasks. For the log A and A tasks, we make these matrices dositive-definite in similar way as to with the bascline Gaussian data.
Figure 1 Hyperparameters. Different tasks perform better with a different selection of hyperparameters for the architecture. For Figure 1 (c), we swept over set hypers and ultimately used the ones that we detail below. For Figure (d), we did not sweep over hypers, but grab the ones from Figure (c) and use them train models for different sizes. Arguably, larger matrix sizes might require larger models. We now explain: Tr(A), Amax (A), log [A], and A-1 from Figure (c).
(r). The trace is control task as we do not usually employ an iterative algorithm to compute the trace. With the NumF ormer r models we use 2 layers, an embedding size of 256, attention heads, and learned positional encodings. We train for 5, 000 iterations using the AdamW optimizer and a batch size of 100 with step-wise decay scheduler and an initial learning rate of 10-3 With the RangeFormes model, we use 8 layers to achieve best performance. Empirically however, we find that we can still out- -erfforro the NumFormer models in Base or Data with only 2 layers.
Amax (A). For the the NumF o rmer model, we use 4 layers, an embedding size of 256 8 attention heads, and learned positional encodings. For the Rangel o omerr model, we use the same configuration but with 6 layers instead. Empirically. the RangeF o ormer performs better with few more layers while NumFormer performance begins to
22

---

Published as a conference paper at ICLR 2026
deteriorate with more layers. We train using the AdamW optimizer for 10, 000 iterations with a fixed learning rate of 10-3 and L2 as a loss function. The RangeF SereeFormn omerr omee can afford much larger batch size of 512 due to its reduced memory footprint while we used a batch size of 128 for NumF o ormer, else we would run out of memory in our hardware.
log |A]. For both the NumF and RangeFormer models, we use 4 layers, an embedding size of 256, 4 attention heads, and learned positional encodings. We use a batch size of 128 and train for 5, 000 iterations using the Adam W optimizer with a learning rate of 10-3 and step-wise decay scheduler. For the log-determinant task specifically, we empirically found that the L.1 loss function or MAF out-performed standard MSE.
A- For the NumF omer model, we use 4 transformer blocks, an embedding size of 256, attention heads and learned positional encodings. We train all the models using a learning rate of 10 and the Adam W optimizer with step-wise decay scheduler for 10, 000 iterations using a batch size of 128. For the RangeFore ormer model, we use the same hyperparameters except we a larger batch size of 256 as looping reduces the memory footprint of the model. We find this task to be difficult as practically, both the train and test set matrices can be relatively poorly conditioned, when using e = 10-1. Interestingly, we also find that RangeF Former performance does not drop even when we use up to 14 layers.
the order of 10
Figure 5 Data (Left). For the Base, Loop, Range, and Attn interventions, we train with random positive-definit Gaussian matrices. The Dist. intervention involves training on a uniform mixture over positive definite decay matrices, explained in Appendix B.I
Figure 5 Hyperparameters (Left). We fix the hyperparameters over all models to isolate the impact of each individual intervention when learning A-1, Specifically, we use 4 layers, an embedding size of 256, and 8 attention heads (although this last change does not apply to the Attn and Dist interventions which use the Ba BaseCovv alternative to attention). We train for 5, 000 iterations using the nuclear loss ||NN, (A) A I||= with a batch size of 12 128, with the Adam W optimizer, an initial learning rate of 10 and a step-wise wise decay learning rate scheduler. In the interest of controlling compute, we enforce a limit of GPU time. To get the error bars we loop over 3 random seeds. In contrast to Figure (c) we did not run a sweep over hyperparameters
each intervention. We simply used hypers detailed above and introduce changes one step at a time.
Figure Data (Right). We use a different training distribution than Figure 5 (Left) in order to analyze the impact of training on structured matrices. For the NumFo rmer r models, we simply use a uniform mixture of 5, 10 x 10 and 20 x 20 Gaussian random matrices. For the RangeF ormer models, we similarly use uniform mixture of 5 x 5, 10 x 10 and 20 x 20 random Einsums. For the \max (A) and Tr (A) tasks, we symmetrize these matrices to ensure that we have real eigenvalues. For the log A and A tasks, we ensure that these matrices are positive definite. We also observe that for the log |A| task, additionally training on polynomial decays can help improve performance for the RangeFo rme method.
Figure 5 Hyperparameters (Right). We train one model per task that is able to interpolate between observed matrix sizes during training (see appendix [c) The RangeF o ozmer models all use the BaseConv attention alternative, looping, and a learned projection. Empirically, we also find that
loss function may result in marginal performance gains. The NumF ormer models use standard quadratic attention with 4 attention heads.
Ar). For the trace task, the RangeF ormer uses 4 layers, an embedding size of 128 and RoPE. We train for 2000 iterations using the AdamW optimizer with a step- wise decay scheduler, an initial learning rate of 10-3, and a batch size of 100. The NumFormex possesses similar hyperparameters but uses only 2 layers.
Amax ()). For the maximum eigenvalue task, the Range7 ormer model uses 4 layers, an embedding size of 256 and RoPE. We train for only 1000 iterations using batch size of 100 with the Adam optimizer, an initial learning rate of 10-3 and a step-wise decay scheduler. The NumForme: possesses similar hyperparameters but uses only 2 layers.
- log [A]. For the log-determinant task, the RangeForme model uses 4 layers, an embedding size of 256, and RoPE, We train for 5000 iterations with a batch size of
23

---

Published as a conference paper at ICLR 2026
100, initial learning rate of 10-3, and step-wise decay scheduler. We observe that for the log-determinant task specifically, the GrokF ast optimizer seems to offer performance improvements albeit inconsistently. The NumFormer uses the exact same hyperparameters.
A For the RangeFo RangeFormer, ormer, we find that we can attain strong performance with 8 layers, an embedding size of 256, and RoPE. We train for 10, 000 iterations, an initial learning rate of 10-3, S, and a batch size of 256. We use the nuclear norm of the difference between NN, A) and I as a loss function. For the NumF ormer, we use the same hyperparameters but with 4 layers and 8 attention heads. We also find that batch size of 256 practically may substantially slow down training so we recommend a smaller batch size on the order of 128 or 64.
Figure 6 (Left). Here we a RangeFo rmer model fit the A-1 task on 50 x 50 matrices. For the RangeFo oreer, we use with 8 layers, an embedding size of 256, 8 heads, and learned embeddings. We train for 10, 000 iterations, initial learning rate of 10-3 and a batch size of 256. We use the nuclear norm of the difference between NN ( loss function.
A and I as
Figure 6 (Middle). Here we use RangeFormes model to fit the Tr(A) task on varying sizes of matrices: 10 x 10, 30 x 30 and 50 x 50 matrices. For the RangeF oreer we use with layers, embedding size of 128, 4 heads, and learned embeddings. We train for 10, 000 iterations, an initial learning rate of 10-3, and a batch size of 256. We use the L.1 loss.
Figure 6 (Right). Here we use RangeFormes model to fit the A task on 50 x 50 matrices. For the RangeFornee ormer, we use with 8 layers, an embedding size of 256, hnd learned embeddings. We train for 10, 000 iterations, an initial learning rate of 10 -3 and a batch size of 256. We use the nuclear norm of the difference between NNa(A). and loss function.
See Appendix D.3 for ablations the effect of different positional encodings models.
D.3 POSITIONAL ENCODINGS
we show the impact of different positional encodings on the performance of RangeForme: for the Tr(A) task. Different tasks exhibit different behaviors with the positional en- codings. For example, in the inverse task A the performance across different positional encodings does not significantly change.
In Figure 13
D.4 WARM-STARTING KRYLOV SOLVERS
In this section explain how trained our RangeForme models for this downstream application as well as present some additional results. For both CG and GMRES, we created a 20 x 20 dataset consisting of our spectrum decays. The RangeFormer model, for this case, uses Taylor attention, 128, looping of R = 8 and a step-wise exponentially decreasing scheduler starting the with learning rate of x 10-3
embedding
In Figure 14 we provide an initial solution to CG or GMRES based on our trained Range rmer model. We cap CG at 12 iterations and display the convergence behavior. The goal of the experiment is to show how neural network and classical method can interact. As seen in the plots, neural network trained many problem could provide some useful information classical methods. In principle, improvements on the performance of RangeF orerr would result in faster convergence for the cases shown. Finally, since the matrix market contains wide variety of matrices from diverse applications, it is understandable that our RangeFo Former models would perform better in some cases and not others.
D.5 RNLA
We now explain the details of our RFSVD method. For this task, train RangeFormes model consisting of embedding dimension of D 16, with H = 1 number of heads and L = layers. However, the input for the RangeFormer model is Gaussian noise On - (0, I16) and the
24

---

Published as a conference paper at ICLR 2026
None
Learned
Sine
Rope
0.5
OOD Rel Enror
0.4
0.2
10
20
50
Size
Figure 13: Performance for different positional encodings. We observe that the presence of learned, sinusoidal and rotary positional encodings each improve performance on the Tr(A) task for matrices taken from the matrix market. Here, we use RangeF o rmer models with 4 layers and an embedding size of 128. We use the BaseConv attention alternative and train 1000 iterations using the AdamW optimizer with an initial learning rate of 10-$, step-wise decay scheduler, batch size of 128, and MSE as the loss function. Error bars indicate + standard deviation around mean from running 3 different seeds.
|AA bilb
110
blb
-cold
10
10
10
cold warm
.cold -wam
10
ccol wwrr
10-2
10
10
10
10 GMRES iterations
CG Inerations
CG Iterations
GMRES berations
(a)
(b)
(c)
(d)
[AAX- olla
10a
Dila
10
[|xr bila
10
|Ax blb
10
10-3
.col -wom
10-11
.col
10-2
ccll -warm
wwam
warm
10
CG iterations
CG iterations
GMRES iteratioons
GMRES Rerations
(e)
(f)
(
(h)
Figure 14: Warm-start performance for CG and GMRES on matrix market. For some cases
RangeFormer approximation improves converge and, for others, least does not harm it. We now detail the matrix used per plot. For (a) used nos3, for (b) we used besstk09, (c) used gr3030, for (d) we used besstk02, for (e) we used besstm/3, for (f) we used besstm06, for (g) we used beasmm24, and, finally, for (h) we used bcsstk15.
25

---

Published as a conference paper at ICLR 2026
output is 52 (N) RN Then I2 is fed into the RSVD algorithm and then the approximate spectrum is compared against that the training matrices (and we backpropagate through the L1 loss). We train the models 50 x 50 matrices using all of our spectral decays.
In Figure [5 we sec how we achieve better rank 16 approximations. Additionally, we observe that the learned distribution Ω NNo (Figure 15 (e)) is distinct from Gaussian noise (Figure [3 (d)).
1.2
true . rsvd
1.4
1.8
M(A)
rfsvd
1.0
A(A)
1.6
1.2
true
1.4
true
0.8
rsvd
rsvd
1.0
rfsvd
1.2
rfsvd
5
10
index
index
index
(a)
(b)
(c)
-3.44
14.77
-3.64
5.0
-16.5
2
3388
-4.04
2.5
-20.1
-4.24 m
4.44 5
X
0.0
-21.
-23.7
-464
-2
-488
2.5
-25.5
5.04
5.0
227
5.24
-9.11
X1
X1
(d)
(e)
Figure 15: Impact of learning the sampling distribution. In (a), (b) and (c) RESVD, RangeF crmer model trained leam sampling distribution outpeforms RSVD using the same rank D 16. The matrices per plot are: (a) 1138 bus, (b) bessrk04, and, finally, (c) nos6. In (d), we show the contour the first two PCA directions of the sampling distribution for RSVD. As expected, each direction is decoupled. In (e), we show the contour of the first two PCA directions of the learned sampling distribution. Now the directions coupled.
D.6 GPs
For this application, we train a RangeFormer model with 8 heads, R = 8 loops, embedding size of D 512, using the Adamli optimizer with learning rate of 10-3 In contrast other applications, here we created specialized training dataset for this application. Let N denote the number of data points used per kernel and the total number of kernel matrices present in our dataset. For every b = 1,.... B, we sample Fn,6 ~ (0 10). for all n 1,...,N, N, & ~ 41(0, 10), and of ~ ((001,11 1]. Then, we construct Ki = exp [=.6- + of where d.,; =
i and 0 otherwise. The set of kernels (K,) where each K, E R N constitutes our training dataset. Indeed, the motivation for the training dataset is to ensure that the RangeForme model performs on wide range of RBF kernels, where the range is motivated by the bounds in the sampling uniform distributions above.
For the experiment discussed in Figure 7 (Right), we follow similar procedure as above but we only create one kernel Ki E XN (B 1) and then we sample Yn ~ A Ki) for = I,...,A N. After that, then provide the data set (z', yn) ) to then find Φ (E, o²) that minimized the loss C (Ø) = log Kg + 7 K" discussed in section 5.3 We then compared the train hyperparameters (* to then ones sampled & (Figure 7 (Right)). This procedure is common check done in probabilistic modeling, as we want to ensure that our training procedure can recover the parameters of the data generating process.
26

---

Published as a conference paper at ICLR 2026
follow the same procedure from the previous paragraph but now with the data (In, yn)a=1 where En ~ (O, 1], en ~ N (0, 0.04) and Yn = (2T2, cn for = ,,..., N.. In contrast to the previous experiments, there is no ground truth & that we should recover. However, we note that the learned hypers & are approximately the same those that would be learned when using choleske fChe1> that is l Ri fooo..
For the experiment discussed at the end of section 5.3
Overall, the main goal of these experiments is to test how would a RangeForme perform when substituting linear algebra routines, not advocate that GP kernels should now be trained with transformers. It is challenging task, at every iteration t. the hypers change d(t) and so now the RangeFor o model has to perform linear algebra operations against K aete and, worse, the trajectory of 6t is determined by how accurate the linear algebra operations compute the loss.
model would
D.7 Loss DETAILS
The different tasks illustrated in require variety of loss functions and model configurations. Basically, we have two sets of losses. The losses that apply to vectors of size N or scalars and the losses that apply to matrices of size N x N. For the first category we have the L1 loss | - yll
E, le; yi and the L2 loss || 3/12 = (z; - y.)², which if required we could report as | - -uli-- We mostly use the 1.1 as a default as we find that to be better performing, though marginally, than the L2 loss. We use the L1 loss for the following tasks: A- E R, log |A), X(A) E RN,, Amax (A) E IR, and, Tr)
For the matrices, we use the nuclear norm ||||, = E; o; where are the singular values of A. In our experiments, we only use matrix losses for the inverse task, and actually, our loss for this task would take the form of (A) Σ||.. Intuitively, we are asking that the neural network approximation NN, (A) behaves like the identity but not that it matches its entries ||NN((A) - - Our choice of loss also has the added benefit of not requiring to compute the inverse A- for the training task.
E RUNTIME ANALYSIS
We provide formal analysis of the time complexity of a forward pass through a single attention head in vanilla transformer (NumFormez) against the proposed RangeFornes
For the following analysis, we denote (I) as the input projection layer, and WQ(), ( WV() the query, key, and value projections. Here, the 1 indicates that we do this analysis for a single attention head. Additionally, let D be the embedding dimension of the model.
Suppose we have some input matrix A E NNN The vector of length N2, Each element of the sequence A is embedded into D-dimensional space resulting in X E N = vec( (T),
r then flattens A into a sequence
A single attention head then constructs Q() X XWQI), k = X XWK(K ((1) and V(1) V(I where W Q1))), WK(1), (11 e DDD, Thus, the construction of Q k(), v E R runs in D()
The attention matrix is then constructed as (Q(")Tr)V())) where is the softmax function. The matrix multiplication (L() runs in D( and constructs matrix of size N2 N incurring a space complexity of ( The final matrix-multiply with (()) incurs the same time complexity so the total time complexity of a single forward pass of an attention head with batch size 1 is O(N D + D2
The major change proposed by the RangeF oomer architecture is to avoid flattening the matrix and instead directly embed the range of the operator through learnable projection. Given A E N RangeFormm computes X E R (I)).
27

---

Published as a conference paper at ICLR 2026
The construction of the key, query, and value matrices is functionally the same but the sequence length is now N instead of N Thus, Q(1,, K and V E NR incur time complexity of O(ND²).
If a practitioner chooses to use RangeF r with the standard attention mechanism, then the total time complexity of a forward pass for one attention head would be O + N DD) with space complexity of (N*)). However, replacing the standard attention mechanism with a subquadratic alternative can offer an even more scalable architecture. We review two alternatives: linear attention and BaseConv.
Linear attention replaces the softmax function a with two feature maps separately applied to the key and query matrices. Concretely,
LinAttn(( K(),, v() o' )o (K
The first matrix multiplication o' ((((T))(()) has a time complexity of O(ND2) resulting in matrix of size D D. The additional matrix multiply against Q incurs the same time complexity resulting in total time complexity of (NN
Another subquadratic alternative to the attention mechanism is BaseConv introduced in Arora et al. (2023), meant to simulate gating and convolutions. Where (1) RNDD is the output of the l'th layer of the transformer, we compute linear projections.
U XW(1)
(5)
U (B
(6)
Where K represents the kernel of a I-dimensional convolution, we compute the depth-wise convolu- tion:
U
K (A)
Finally, compute:
X ((+1) oonn (B)
(8)
Equations
5 and 6 run in O(ND²) resulting in the same time complexity linear attention.
28

---

Published as a conference paper at ICLR 2026
F CHARACTERIZING THE EFFECTIVENESS OF RANGEFORMER
The core idea behind the RangeFormes ormer is to capture most of the action of a matrix through learnable low-rank projection matrix. We expect the low-rank projection work as there exist theoretical guarantees for using random Gaussian matrices as projection matrices. Below, we reproduce the problem statement and proof from ?, Theorem 11.5.
Let A E R= be an imput matrix with I < be the subspace matrix. The rangefinder problem is to find an orthonormal matrix Q gaxt whose range aligns with the largest left singular vectors of A. We define the spectral norm error
- - A = |(I AQ1)A|I
(9)
The rank-l matrix Á QQ* A serves
approximation for A.
Fix A E R"x" with singular values 01 2 d2 2 Draw a standard normal test matrix S2 R
and construct Y AΩt. Choose & < [ - and introduce the random variable Z |[t] where
E &x is also standard normal. Define Py as the orthogonal projector onto the range of Y. Then,
1/2
k
B|| - Pr)
1
0k+1 (EZ)
(10)
1
The randomized rangefinder finds an 1-dimensional subspace that captures as much of the action of A the best k-dimensional subspace.
Proof. Since S2; (0, 1), the matrices Ω == Vi Ω2 == V are also independent standard
normal since V1, V/2 orthonormal and mutually orthogonal. Then by Chevet's Theorem referenced in ?, Prop 10.1,
B | = Ea, En,
E
(11)
< E
Σg | |2
+
Σ> [
[|]
(12)
k
|Σ
(EZ) Ez|le
(13)
k
The last inequality involves a well-known estimate of the trace of an inverted Wishart matrix described in Prop 10.2.
We can insert estimates of the expectation for Z. When 2 1, EZ evi
I-E When &s <
EZ ~ VIvK These estimates lead to accurate performance bounds across many matrices and parameters.
For RangeFo Forerr specifically, would expect method perform better for matrices with rapidly declining cigenspectra as that way, the action of the matrix can be better captured in a low-rank projection. For example, in Figure I0 using a polynomial diagonal decay, would not expect our method to work so well for a = 10 as it has many large eigenvalues without decay but would expect RangeFormer to perform better for E [0.05, 0.3). Similarly, we would expect performance improvement matrices in the matrix market shown in Figure I1 due to their declining eigenspectra.
29


Reviews:
[{"review_id": "VYYQxIarBN", "reviewer": "Reviewer_LWQ3", "summary": "Starting with the argument that a \"generally intelligent system\" should be perfectly competent at performing matrix operations, this work studies Transformers' ability to perform primitives relevant to linear algebra (e.g., eigenvalue computations, trace computation, etc.). Results show the Transformers architecture explored in this work is unable to perform well, but interventions on the architecture based on seen failures leads to a new architecture---called rangeformer---which can fix these issues. In essence, rangeformer is a series of interventions involving changes from positional encodings to weight tying (aka looping) to data level interventions.", "questions": "See weaknesses.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 1, "strengths": "The paper's experiments are well written and the sequential structure made it fun to follow.", "weaknesses": "- My core apprehension is that I do not currently see sufficient novelty or utility to the paper. The paper is not necessarily a science of deep / machine learning project that tries to elicit new phenomena or explain known ones (i.e., does not ask question of the sorts \"why does a Transformer fail to perform a task\"), nor is it an explicit methodological improvement (i.e., it does not try to introduce a new architecture for practical tasks and verify the efficacy of proposal). This, personally, makes it hard for me to contextualize this paper with respect to existing literature.\n\n- Related to above, if I focus on the precise narrative the authors took for motivating their study, i.e., generally intelligent systems should be able to perform linear algebra, then I find myself not convinced at all by this narrative. For several problems that Transformers were benchmarked on in this paper, I argue I would personally fail to solve several of these tasks with a sufficiently high performance. That said, I can easily write a program or do a library call to perform these deterministic calculations, and so can Transformers. In this sense, I don't really see the motivation to wanting to replace a deterministic calculator with a fuzzy one.\n\n  - To be clear, I'd generally deem a point like this to be minor and not flag it in the review. However, given that this statement was made multiple times by the authors and because I do not see a sufficient motivation for the paper, I'm leaning into authors' proposed motivation.\n\n- Prior work: Authors state (L46--48) that there is little prior work exploring how Transformers can perform linear algebra tasks. I strongly disagree with this. There's several papers exploring, e.g., the ability of a model to perform tasks like matrix completion [1], broader math operations [2], identifying the benefits of looping for this [3], effects of data properties [4], and so on. \n\n[1] https://arxiv.org/abs/2410.22244\n\n[2] https://arxiv.org/abs/2506.13688\n\n[3] https://arxiv.org/abs/2502.17416\n\n[4] https://arxiv.org/abs/2307.03381", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NQ51cR40lc", "reviewer": "Reviewer_Qhzi", "summary": "This paper begins by identifying the inherent limitations and failure modes of Transformer-based approaches to matrix computation. The authors’ analysis reveals that standard Transformers primarily act as statistical interpolators rather than genuine algorithm discoverers, which leads to particularly poor performance on out-of-distribution (OOD) data. Building upon these insights, the paper introduces a series of architectural improvements aimed at enhancing both scalability and accuracy, culminating in a new model termed RangeFormer. The proposed architecture demonstrates substantially improved scalability and performance on OOD matrices from the Matrix Market collection. Moreover, the paper provides, to the best of the authors’ knowledge, the first empirical evidence that Transformers can be effectively applied to downstream tasks involving iterative matrix operations, such as Gaussian process learning and improving sampling distributions in randomized numerical methods.", "questions": "* The notation in the figures (e.g., “HP”) is insufficiently explained, making them difficult to interpret.\n* It is recommended that the main text include additional references to the literature on RandNLA and RSVD.\n* In the comparison between RSVD and RFSVD, the paper only reports the error metrics but does not provide a comparison of the running time. Although RFSVD improves accuracy, it remains unclear whether this improvement comes at the cost of increased computational time.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper introduces a learnable projection mechanism that embeds the input matrix into a higher-dimensional latent space. This design substantially improves computational efficiency and scalability, while enhancing the Transformer’s performance on matrix-multiplication–type tasks.\n\n* By adopting a looped Transformer structure inspired by iterative methods in numerical linear algebra, the model can repeatedly process the same matrix multiple times. This enables adaptive refinement and better handling of tasks with varying difficulty and condition numbers.\n\n* The authors construct a novel structured data mixture for training, consisting of matrices with diverse structures and eigenspectrum decay patterns. Compared to conventional Gaussian-random training data, this approach significantly improves out-of-distribution (OOD) generalization.\n\n* The paper provides the first empirical demonstration that Transformers can be successfully applied to downstream tasks involving iterative matrix operations, such as Gaussian process learning and improving sampling distributions in randomized numerical methods.", "weaknesses": "I will admit I am not very familiar with this paper, so I am basing my review more on the clarity of the writing. \n\n\n1. The theoretical explanation of the RangeFormer design remains largely intuitive. The paper proposes four key modifications, range embedding, looping, linear attention, and structured data mixture, but their effectiveness is demonstrated only empirically, without rigorous theoretical justification or detailed analysis. In particular:\n* Why can the projection operator $X = A\\Gamma$ preserve the matrix spectrum or operator properties?\n* Must the projection be Gaussian? In practice, Gaussian projections are dense and computationally expensive, could sparse alternatives such as CountSketch or other structured embeddings achieve similar effects?\n* Why does *linear attention* improve the approximation of matrix multiplications?\n* What are the convergence and stability guarantees of the *looping* mechanism?\n* Does the structured data mixture genuinely improve algorithmic generalization, or does it merely enlarge the training support domain?\n\n2. RangeFormer is derived from the baseline NumFormer through several architectural modifications. While the paper describes these individual changes, it does not provide a complete framework diagram of RangeFormer that would help readers clearly understand the overall workflow and methodological pipeline of the proposed model.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4MbdRJDt4H", "reviewer": "Reviewer_XjAK", "summary": "This paper presents a comprehensive empirical investigation into the capability of transformer architectures to execute fundamental linear algebraic operations. Through their simulations, the authors assess whether transformers can accurately perform tasks such as vector addition, matrix multiplication, and eigenvalue estimation. The results provide empirical evidence that, despite their expressive power, transformers—particularly those constructed following existing theoretical proofs—struggle to reliably perform even simple linear algebra operations. \nFinally, the paper proposes a new version of transformer that is capable of performing those tasks.", "questions": "1. Can the authors explain why would adding an extra matrix embedding layer would help the performance? At least in the linear regression task [Garg et al., 2023, Von Oswald et a. 2023], multiplying the data matrix with a random matrix results in another linear regression problem. Therefore, intuitively I do not understand how would it work.\n\n2. Can the authors address my two concerns? It is possible that I misunderstand some of the goals of this paper.\n\n3. One observation from the paper is that trained transformers do not have good OOD generalization on linear algebra tasks, which seems to be an universal issue other than for this specific task. Would training transformers on a more diverse dataset gain performance?", "limitations": "", "rating": 2, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* The simulation section is comprehensive.\n* The research question this paper tries to address is indeed new to the ML community.", "weaknesses": "My main concern of this paper is that I do not see how the results from this paper would provide insights to either empirical or theoretical studies of transformer.\nI believe a main statement this paper tries to make is that the transformer constructed in [Garg et al., 2023, Von Oswald et a. 2023], does not appears to be the same as the learned transformer.\nHowever, this does not seem to be a significant insight to the theoretical community.\nIt is obvious that in a series of works on \"transformers simulate algorithms\", their proofs are mainly constructive, meaning they only show the existence of a solution instead of the uniqueness of it.\nTherefore, while analyzing transformers block-by-block seems to be a new approach to my knowledge, I think it is obvious that trained transformers will not necessarily converge the the exact construction from previous studies.\n\n\nMy second concern is that even from the empirical perspective, I am unsure what does the insight or the proposed model from this paper contribute to empirical ML.\nI understand that this paper is not trying to push SOTA performance. However, if ones goal is to utilize transformers for linear algebra operations, why not simply follow the constructions in previous studies?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fH6gYhPlQt", "reviewer": "Reviewer_w4zs", "summary": "This paper performs an empirical analysis of transformers trained in a supervised manner to perform linear algebra tasks. The authors design detailed experiments to identify the bottlenecks of transformers in performing such operations. Specifically, they find that transformers tend to learn properties of the input data distribution rather than the underlying matrix operation algorithms.\nTo address this, the authors propose a new training recipe and a novel combination of architectural components to mitigate the identified limitations. Finally, they evaluate the improved architecture on three downstream applications.", "questions": "*  I appreciate the detailed and carefully designed experiments examining the transformer’s limitations in solving linear algebra tasks. However, as mentioned in the weakness section, I would appreciate a more in-depth explanation of the importance and broader motivation behind this problem. Specifically, what kind of future research directions or practical applications does understanding transformers’ ability to perform linear algebra enable?\n\n*  Could the authors clarify why they selected the three specific downstream tasks in Section 5? Are these tasks directly related to the failure cases identified in Section 3, or do they represent scenarios where classical solvers perform poorly or are computationally expensive, thus providing opportunities where RangeFormer offers advantages?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "* The paper conducts a detailed and comprehensive analysis of the limitations of transformer-based architectures in performing matrix operations. It further proposes a new model architecture and training recipe, and evaluates them on a wide range of tasks to demonstrate their effectiveness.\n\n* The paper provides thorough details of all experimental settings, which is good for future reproducibility.\n\n* The work uses simple linear algebra tasks as a testbed to reveal that transformers may primarily learn the statistical properties of the data distribution rather than the underlying algorithmic procedures.", "weaknesses": "* Overall, the paper is a solid empirical study. However, my main concern is in the motivation for investigating the transformer’s ability to perform such simple linear algebra tasks.\nAs the authors repeatedly emphasize, (e.g., line 86: “We note that the goal of our work is in understanding the potential capabilities and limitations of transformers for linear algebra…”; line 472: “This research area is in its early stages… our goal is not to compete with purpose-built classical algorithms…”),  the stated goal is primarily exploratory. However, it remains unclear why understanding transformers’ capabilities on these tasks is important, and what future directions this understanding enables. Does this line of research aim to eventually develop a single transformer that outperforms classical linear algebra solvers? Why is a new transformer architecture needed to solve problems that classical solvers already handle perfectly? It is also not evident whether the insights obtained from these experiments will translate to a better understanding of transformer usefulness in real-world analytical settings.\n\n* Many experiments in the main text lack multi-run results with different random seeds. While the appendix provides details about whether hyperparameter optimization was performed, the plots in the main text appear to come from single runs. Including multiple seeds would strengthen the statistical reliability of the results. Since the models are relatively small and trained for fewer than 10,000 iterations, adding multi-run averages and std should be feasible.\n\n\n* Some aspects of the writing logic could be improved. For example, in line 297 -line 300, the authors state \"We now construct, step-by-step, our RangeFormer architecture, and introduce our synthetic data procedure....\", but the following paragraph immediately shifts to discussing experimental results instead of the mentioned architecture.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WI6ku24bqP", "reviewer": "Reviewer_gjAT", "summary": "The goal of the paper is to study the capability of the transformer architecture to perform numerical linear algebra tasks, such as linear system solving, computing SVD and EVD and so on. The first contribution of the paper is that the existing architecture approaches is unable to learn generalizable algorithms and show OOD generalization. Their main contribution is RangeFormer, which has various features such as learnable matrix embeddings, linear attention, and a diversifying the training distribution. These improvements first  reduce the make the transformer efficient for linear algebra tasks and also improves OOD generalization.\nThey show their new architecture can scale well compared to the classical transformer in performing various tasks such as Gaussian process learning and randomized numerical linear algebra. Furthermore the paper also explores using transformers to improve the performance of classical methods such as Krylov subspace methods and randomized SVD by providing better initialization and sampling distributions.", "questions": "I found the paper quite interesting and I have a few questions regarding it. Firstly there has been quite a lot of work on linear algebraic tasks such as linear regression and the ability of transformers to solve them in context ( for eg. refer to this work https://arxiv.org/abs/2211.15661). As far as I understand your current work pretrains the architecture for linear algebraic tasks, however it would be interesting to know if claims along the lines that the transformer has enough expressivity to perform tasks such as SVD or linear system solving in context would be quite interesting.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The main strength of the paper is in using linear algebraic tasks as a benchmark for understanding the limitations of transformer architecture. Various works have shown such limitations for different algorithmic tasks, and this paper delves deeper into tasks pertaining to implementing algorithms for numerical linear algebra.", "weaknesses": "Perhaps a minor weakness of the paper is in terms of the writing, and I would suggest the authors to include a main summary of their key findings in the beginning of the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)