REVIEW 4 major objections 4 minor
NLAFormer: Transformers Learn Numerical Linear Algebra Operations
T0 review · 4 major / 4 minor · reviewed 2026-08-15 · deepseek-v4-flash
Pith's one-line read NLAFormer constructs explicit transformer weights showing that one or two layers can express the basic operations of numerical linear algebra, and that assembling these blocks yields a transformer which can execute the conjugate gradient…
desk verdict The basic-operation expressivity claims are plausible and worth a look; the CG theorem has a load-bearing gap that the authors would need to close. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The load-bearing mechanism is the prompt-matrix encoding: data rows hold $A$ and $b$, working rows hold the current CG scalars $d_k$, $x_k$, $r_k$, and the bottom row of unit vectors $e_1,\ldots,e_{n+1}$ acts as a positional key. Attention weights are chosen so that softmax scores become approximately $1+Z$ for large $C$ and small $c$, which makes multi-head attention behave like a linear operator that computes matrix-vector and matrix-matrix products and writes them into empty rows. A second ingredient is the column-wise feedforward network, which uses universal approximation to apply the same continuous function, whether pointwise arithmetic or the scalar CG updates, to every column. A third is the structured loop: a pre-processing block initializes $d_0=r_0=b-Ax_0$, a weight-shared loop block performs one conjugate gradient iteration, and a post-processing block emits the final iterate.
What would settle it
Run the Lemma 3.2 construction with $n=2$ on a non-trivial symmetric positive definite matrix, inspect the last column of the loop block's output, and compare with true CG's $\alpha_0$ and $\beta_1$: if the values differ by more than the claimed $\varepsilon$, the construction does not execute CG. More directly, check whether the FFN's last-column input, which contains $d_k$, $x_k$, $r_k$ but not $A$, determines $d_k^\top A d_k$; if it does not, no column-wise function can produce $\alpha_k$.
Extended reading notes
Core claim
The central discovery is a recipe for turning a transformer layer into a linear algebra kernel. By arranging the input as a prompt matrix whose bottom rows are unit vectors, and by choosing query, key, and value weights so that $(e^C \mathrm{softmax}(Z))_{i,j} \approx 1+Z_{i,j}$, self-attention becomes an affine operation that can write $Ad_k$, $A^\top B$, $a^\top b$, or a transposed matrix into a designated row. The feedforward network, acting column-wise and guaranteed by the universal approximation theorem to realize any continuous column function, supplies pointwise arithmetic and the per-iteration updates of the conjugate gradient method. The paper's claim is that these ingredients assemble into a complete solver: four attention heads in a pre-processing layer produce $d_0$ and $r_0$, and one two-head loop layer advances $d_k$, $x_k$, $r_k$ by one CG step, so $n$ loop applications yield the solution of $Ax=b$ to arbitrary approximation accuracy. Trained versions of the same architecture reproduce CG's iterates and, when supervised by the final solution, learn update rules that converge faster than the classical algorithm.
Load-bearing premise
The load-bearing premise is that the column-wise feedforward network can compute the conjugate gradient scalars $\alpha_k$ and $\beta_{k+1}$ from the single last-column entries $d_k$, $x_k$, $r_k$, even though those scalars require global reductions such as $d_k^\top A d_k$ and $\|r_k\|^2$, and the proof does not construct an attention head that supplies these reductions.
Editorial extensions
If this is right
- A transformer with $O(n^2)$ input and one or two layers can express every basic operation in the paper's list, replacing loop-transformer constructions that require $O(n^3)$ input and, for pointwise operations, up to $12n$ layers.
- Composing these operations yields a transformer that can execute the full conjugate gradient method for any symmetric positive definite system, with one loop block per iteration.
- Because the loop block has fixed weights across iterations, the construction shows that transformers can represent iterative algorithms without simulating a general computer's control flow, which is the overhead NLAFormer removes.
- In training, the architecture internalizes CG's recursion: intermediate variables align with CG's residuals and search directions even when only final solutions are supervised, and joint supervision improves that alignment.
- With solution-level supervision, the trained model reaches target accuracy in roughly half the iterations of CG on the test distributions, and training it on PCG trajectories yields a solver that finishes in fewer iterations than classical PCG.
Reading between the lines
- If the expressiveness theorems are correct, then any numerical algorithm built from these primitives, not just CG, has a constant-depth, weight-shared transformer implementation whose depth depends on iteration count rather than on $n$; GMRES, Richardson, or Chebyshev acceleration are natural next candidates.
- The CG theorem as written leaves a gap at exactly one point: the last-column scalars $d_k$, $x_k$, $r_k$ must feed the scalar updates $\alpha$ and $\beta$, which are global reductions over $n$-vectors. A direct extension would add explicit reduction heads or store full vectors in working rows, separating a construction issue from a fundamental limitation.
- The empirical speedup over CG suggests the trained model may have discovered a data-dependent preconditioner or a better step-size rule; probing the hidden state for an operator resembling $M^{-1}A$ would make that mechanism visible and could inform new hybrid solvers.
Signed reviews
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. This paper proposes NLAFormer, a transformer architecture claimed to execute numerical linear algebra primitives and, by assembly, the conjugate gradient method. Theorem 2.1 asserts that one- or two-layer transformers can perform pointwise operations, shifts, transposition, inner products, and matrix and matrix-vector products with O(n^2) inputs; Theorem 3.1 asserts that a pre-processing block plus a one-layer, two-head loop block can execute CG for symmetric positive definite systems. The proofs combine a softmax-linearization lemma imported from prior work with universal approximation for the column-wise FFN. Experiments train and test NLAFormer on n=20 systems and report that it matches or outpaces CG and PCG in relative error.
Significance. The high-level idea of replacing control-flow simulation with compact algebraic building blocks is attractive, and Table 1 gives a useful resource comparison with loop-transformers. If Theorem 2.1 held as stated, the layer counts would be a notable expressiveness improvement, and the structured preprocess-loop-postprocess design is worth studying. However, the proofs are too sketchy to verify, and the CG theorem contains a fundamental gap: the input prompt never stores the matrix A, and the global scalars alpha_k and beta_{k+1} cannot be computed by the column-wise FFN from the scalar state entries. The experiments on n=20 systems, without code or data release, do not repair this gap. The paper would need a substantially reworked construction for the iterative theorem; as written, the central claim is not established.
major comments (4)
- [Section 3, Lemma 3.2, Eqs. (3.3) and (3.5)] The prompt (3.3) has only n+1 columns and stores vectors a and b in rows 1 and 2 with scalars d_k, x_k, r_k in the last column; it nowhere stores the n by n matrix A. The attention construction uses the first row of P as an n-vector a (see W_K P = [0 a_1 ... a_n; C 0 ... 0]), yet Eq. (3.5) claims the attention output is c A d_k. Since A is not part of the prompt and cannot fit in the stated shape, the claimed computation of A d_k is not established. This is a load-bearing failure because every CG update in Eq. (3.1) depends on A d_k.
- [Section 3, Lemma 3.2, Eq. (3.7)] Even if A d_k were somehow available, the FFN is applied independently and with shared weights to each column, as stated in Section 2. The state vectors d_k, x_k, r_k appear only as scalar entries in the last column of U in (3.6), and no attention head is constructed to compute the global reductions ||r_k||^2, d_k^T A d_k, or ||r_{k+1}||^2. The coefficients alpha_k and beta_{k+1} in CG (3.1) require aggregation over all coordinates, so they are not functions of the scalar values in any single column. The sentence that x_{k+1} - x_k, r_{k+1} - r_k, and d_{k+1} - d_k are continuous mappings of x_k, r_k, d_k, and A d_k is therefore insufficient; universal approximation on a fixed column can neither see the full vectors nor perform the required reductions. Consequently Lemma 3.2 does not prove that a one-layer, two-head loop block executes one CG step.
- [Section 2, Theorem 2.1 proofs] The proofs of the inner product, outer product, matrix transpose, A^T B, AB, and Ab cases are not written out; the manuscript refers to similar arguments and to references [4,5]. More concretely, the displayed weight matrices in the column-shift and inner-product proofs have dimensions inconsistent with the prompt matrices: in part (ii) P is 3 by 3 but W_Q is displayed as a 3 by 4 matrix, and in part (iii) P is 4 by (n+1) but W_Q is displayed as a 2 by 3 matrix. Because Theorem 2.1 is one of the two central contributions, these unverified constructions and dimension mismatches prevent a rigorous check of the claimed expressiveness results.
- [Section 3, proof of Theorem 3.1] The proof of Theorem 3.1 iterates Lemma 3.2 and claims that after n applications the output has the form (3.14). Since Lemma 3.2 is not established, this does not prove the theorem. In addition, no post-processing block TF_post is constructed, so the full pipeline NLAF(P) in (3.2) is not shown to deliver the solution x_n in the claimed output format.
minor comments (4)
- [Section 4] No code or data are released, all experiments use n=20, and the plots do not include error bars or standard deviations; the claims that NLAFormer generalizes or converges faster than CG/PCG are therefore only weakly supported.
- [Section 4.1] The text mentions training over 700k epochs; this is likely a misstatement for optimization steps, especially since the learning-rate schedule in Figure 4(a) is plotted against steps.
- [Section 3] There are small presentation errors: 'deonated' should be 'denoted', the reference to 'Theorem 1' in the proof of Lemma 3.2 should presumably be to Theorem 2.1, and the notation A = [a1, ..., an]^T is not defined for a matrix A in n by n.
- [Equations (2.16)-(2.23) and (3.4)-(3.8)] The approximation relation is used without explicit error bounds; stating that C is sufficiently large and c is sufficiently small is not enough to turn the displayed approximations into a formal theorem with controllable error.
Circularity Check
No circular derivation: the CG issue is an omitted proof, not a circular reduction, and the only self-citation is a supporting lemma.
full rationale
The derivation chain for Theorem 2.1 rests on explicit constructions of W_Q, W_K, W_V and on the standard softmax linearization e^C softmax(Z) ≈ 1+Z, imported from [4,5]. Reference [4] shares an author, so this is a self-citation; however the linearization lemma is a parameter-free mathematical approximation whose assumptions (large C, small c) are stated in the paper, and the central expressiveness claims are not obtained by renaming a result from [4]. The same applies to the 'another head to cancel the constant' trick: it is a cited construction, not a conclusion equated to an input. I therefore find no self-definitional, fitted-input-called-prediction, or uniqueness-imported circularity. The main problem in the paper is a correctness gap rather than circularity: Lemma 3.2's proof states that x_{k+1}-x_k, r_{k+1}-r_k, and d_{k+1}-d_k are continuous mappings from x_k, r_k, d_k, and A d_k and then invokes universal approximation, but in prompt (3.3) d_k, x_k, and r_k appear as scalar entries in the last column and no attention head constructs the global reductions ||r_k||^2 or d_k^T A d_k needed for alpha_k and beta_{k+1}. A column-wise FFN cannot aggregate across the n coordinates from the data actually placed in its column, so Theorem 3.1 is not established by the written proof. This affects the claim's completeness, not its circularity. Similarly, Section 4.2's faster-convergence result is trained with the explicit solution-supervision term in (4.2), so the observed speed-up is an optimization objective rather than an unprompted prediction; the paper states that objective transparently. Overall, no load-bearing step is equivalent by construction to its input, so circularity is minimal; I score 2 to reflect the self-citation dependence of the attention lemma and the omitted CG proof.
Assumptions & free parameters
free parameters (4)
- C and c (attention scaling constants in Theorems 2.1 and 3.1)
- eta (intermediate supervision weight) =
0.0005
- lambda (solution supervision weight) =
10
- embedding dimension d =
256 (ablation: 128, 512)
assumptions (4)
- standard math Universal approximation theorem for multilayer feedforward networks with nonpolynomial activation (Leshno et al. [9])
- domain assumption Softmax linearization: (e^C softmax(Z))_{i,j} approximately 1 + Z_{i,j} for large C and small c, with a second head canceling the constant
- domain assumption All vectors and matrices lie in compact subsets K^n and K^{n x n}
- ad hoc to paper The FFN in Lemma 3.2 can compute alpha_k and beta_k from x_k, r_k, d_k, Ad_k despite acting column-wise
Cite this review
Pith. "Pith review of NLAFormer: Transformers Learn Numerical Linear Algebra Operations." pith.science (2026). https://pith.science/paper/W6XCZXPC
@misc{pith2026250819557,
author = {Pith},
title = {Pith review of: NLAFormer: Transformers Learn Numerical Linear Algebra Operations},
year = {2026},
howpublished = {\url{https://pith.science/paper/W6XCZXPC}},
note = {Machine review of arXiv:2508.19557}
}
read the original abstract
Transformers are effective and efficient at modeling complex relationships and learning patterns from structured data in many applications. The main aim of this paper is to propose and design NLAFormer, which is a transformer-based architecture for learning numerical linear algebra operations: pointwise computation, shifting, transposition, inner product, matrix multiplication, and matrix-vector multiplication. Using a linear algebra argument, we demonstrate that transformers can express such operations. Moreover, the proposed approach discards the simulation of computer control flow adopted by the loop-transformer, significantly reducing both the input matrix size and the number of required layers. By assembling linear algebra operations, NLAFormer can learn the conjugate gradient method to solve symmetric positive definite linear systems. Experiments are conducted to illustrate the numerical performance of NLAFormer.
Reviewed August 15, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.