{"id":"43bba076-bbc0-4608-8bce-70cf6ca04f10","arxiv_id":"2508.19557","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":4.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":4,"one_line_summary":"NLAFormer is a transformer construction that expresses core numerical linear algebra operations in one or two layers and can be trained to emulate conjugate gradient iterations for SPD systems.","lead":"This paper introduces NLAFormer, a transformer design that can express basic matrix and vector operations in one or two layers, and shows that a trained version can mimic and sometimes outpace the conjugate gradient solver on small synthetic systems. The reading point is whether large AI models can be repurposed as compact numerical solvers, though the proof and the experiments still have important gaps.","discovery_kind":"extension","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Theorem 3.1's CG construction does not compute the global scalars alpha_k and beta_{k+1}; the column-wise FFN cannot perform the required reductions, so the central CG expressiveness claim is unsupported.","rationale":"The reader's weakest assumption identifies the same load-bearing gap: the global scalars alpha_k and beta_{k+1} are never constructed, and the column-wise FFN cannot compute them from the stated prompt layout. My read agrees with this assessment, and I do not find an additional concern that would change the verdict. The paper's Theorem 2.1 constructions for basic operations are plausible extensions of known softmax-linearization techniques, and the empirical section still shows that a transformer can learn CG-like updates on small systems. However, the CG theorem is not proven as stated, so a conditional verdict is appropriate: the construction must be repaired, or the theorem must be weakened, before the central claim of assembling iterative solvers is accepted. I would not move to REJECT because the broader architectural and empirical contributions remain informative and potentially salvageable.","tokens_in":15382,"tokens_out":6443,"duration_ms":63553,"concrete_test":"Independently re-derive Lemma 3.2 by writing the prompt (3.3) with explicit n-dimensional vectors d_k, x_k, and r_k, and symbolically compose the two specified attention heads with the column-wise FFN. Check whether d_k^T A d_k, ||r_k||^2, and ||r_{k+1}||^2 appear anywhere in the pre-FFN features. If they do not, the universal-approximation argument cannot produce alpha_k or beta_{k+1}, settling that Theorem 3.1 fails as stated. A sharper check: show that the FFN output at column j can depend only on (d_k)_j, (x_k)_j, (r_k)_j, (A d_k)_j, and fixed markers; then the identical scalar alpha_k required in every column cannot be encoded unless it is already present in each column.","verdict_should_be":"UNCHANGED","load_bearing_attack":"In Lemma 3.2, after the two attention heads, the matrix U in (3.6) contains d_k, x_k, and r_k as final-column entries and c A d_k in row 6. The FFN then processes each column independently with the same weights, as the paper itself states. To execute the CG update, the network must produce alpha_k = ||r_k||^2 / (d_k^T A d_k) and beta_{k+1} = ||r_{k+1}||^2 / ||r_k||^2, which are reductions over all n coordinates of the state vectors. No attention head in TF_loop constructs d_k^T A d_k, ||r_k||^2, or ||r_{k+1}||^2; the proof only constructs A d_k. Universal approximation on a single column cannot recover these global inner products from the scalar entries that are actually present, because the FFN never sees the full vectors in any one column. The sentence in the proof 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 is therefore insufficient: those mappings are not functions of the per-column scalars alone, and the FFN cannot aggregate across columns. The same gap affects beta_{k+1}. Since Theorem 3.1 relies entirely on Lemma 3.2, the claim that a one-layer, two-head loop block executes CG for n-dimensional systems is not established. The experiments do not repair this gap, as they train only n=20 systems and release no code or data.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","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.","tokens_in":15750,"tokens_out":12254,"duration_ms":111733,"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":[{"comment":"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":"Section 3, Lemma 3.2, Eqs. (3.3) and (3.5)"},{"comment":"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":"Section 3, Lemma 3.2, Eq. (3.7)"},{"comment":"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":"Section 2, Theorem 2.1 proofs"},{"comment":"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.","section":"Section 3, proof of Theorem 3.1"}],"minor_comments":[{"comment":"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":"Section 4"},{"comment":"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":"Section 4.1"},{"comment":"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.","section":"Section 3"},{"comment":"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.","section":"Equations (2.16)-(2.23) and (3.4)-(3.8)"}],"recommendation":"reject","confidential_remarks":"The paper relies heavily on the softmax-linearization machinery of references [4,5], with [4] sharing an author; the current manuscript does not re-derive or verify those lemmas. Given the gap in the CG construction, the dimension inconsistencies in the weight matrices, and the absence of code or data, I do not see a repair that stays within a normal revision cycle."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Quick take: the paper gives explicit one- and two-layer transformer constructions for standard NLA primitives with O(n^2) input, which is a real improvement over the loop-transformer's O(n^3). That part is worth engaging with. The CG section, however, does not hold up as written: the loop block stores d_k, x_k, r_k as scalar entries in a single column, and the FFN processes columns independently, so it cannot compute the global inner products that define alpha_k and beta_{k+1}. The proof hand-waves this with a continuity claim that doesn't match the architecture. This is the central expressiveness result, so the paper's main theorem is unsupported without a substantial rewrite.\n\nWhat is new: the specific constructions for pointwise ops, shifts, transpose, inner product, matmul, matvec, and the input-size bounds. They are a direct application of the softmax-linearization technique from [4,5], but the authors do contribute the particular encodings and layer counts. The experiments on n=20 SPD systems show the model can imitate CG when trained with intermediate supervision, and the 'learned faster solver' result is suggestive, though the comparison is not apples-to-apples because the model sees the true solution during training.\n\nSoft spots: (1) The CG proof gap I mentioned. The prompt stores vectors as scalars; no attention head constructs d_k^T A d_k or the norms; a column-wise FFN cannot aggregate over n coordinates. This is load-bearing, not a cosmetic issue. (2) Several parts of Theorem 2.1 are left as 'can be proved similarly'; I'd want to see the remaining cases written out before trusting them. (3) Experiments: no code or data, one matrix distribution, n=20 only, and the faster-convergence claim overstates what the loss design shows. The comparison with CG/PCG is on relative error per iteration, not wall-clock time or FLOPs, and the solution-supervision term is an oracle the classical methods don't get.\n\nIf the authors fix the CG construction—perhaps by storing full vectors in the prompt and adding attention heads that compute the reductions, or by proving a different mechanism—the paper could be a solid contribution to the transformer-expressiveness literature. As it stands, the basic operations part is worth reading; the CG theorem needs major revision. I'd send it to review, but with a clear request for a corrected proof and a more rigorous empirical section.","headline":"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.","tokens_in":16267,"tokens_out":4497,"would_cite":false,"duration_ms":41076,"reading_group":"maybe","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":["65F10","68Q32"],"pacs":[],"model":"deepseek-v4-flash","headline":"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…","keywords":["transformers","numerical linear algebra","conjugate gradient","attention mechanism","universal approximation","iterative solvers","expressiveness","matrix multiplication"],"falsifier":"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$.","tokens_in":15153,"feed_emoji":"🧮","tokens_out":11027,"duration_ms":99429,"temperature":0.7,"pith_summary":"This paper claims that transformers are not merely pattern-matching tools for numerical tasks but can be explicitly wired, with hand-chosen weights, to execute the standard operations of numerical linear algebra. Theorems 2.1 and 3.1 give constructions: one-layer blocks for pointwise arithmetic, shifts, transposes, inner products, $a^\\top b$, and $A^\\top B$; two-layer blocks for $AB$ and $Ab$; and a pre-processing block plus a shared loop block that carries out the conjugate gradient iteration for symmetric positive definite systems. The payoff is a sharp reduction in cost versus the loop-transformer approach: $O(n^2)$ input and one or two layers instead of $O(n^3)$ input and many layers. The experiments then show that when NLAFormer is trained rather than hand-built, it reproduces CG's trajectory and, with an extra solution-matching loss, reaches a given accuracy in fewer iterations than CG, in some cases nearly half, and similarly outperforms PCG when trained against PCG trajectories.","feed_headline":"NLAFormer: one or two transformer layers run numerical linear algebra","feed_subtitle":"Compact O(n²) prompts carry matrix operations; the trained loop block uses fewer iterations than CG or PCG.","key_machinery":"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.","core_discovery":"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.","pith_inferences":["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."],"forward_implications":["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."],"supporting_citations":[{"why":"Supplies the universal approximation theorem for ReLU networks, used to justify the column-wise FFN computations in Theorems 2.1 and 3.1.","marker":"[9]"},{"why":"Defines the loop-transformer baseline and contains the softmax linearization estimates that the attention constructions rely on.","marker":"[5]"},{"why":"Provides the pre-processing, loop, and post-processing module design and the two-head constant-cancellation trick adopted for NLAFormer.","marker":"[4]"},{"why":"States the conjugate gradient iteration that Theorem 3.1 is designed to implement.","marker":"[10]"}],"fun_headline_variants":["NLAFormer: self-attention as an affine linear algebra operator","NLAFormer: compact transformer runs CG with fewer iterations","NLAFormer: a transformer that solves linear systems via CG","Transformers learn matrix operations and even the conjugate gradient method","NLAFormer shrinks input size and layers for linear algebra"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"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.","fun_headline_variants_meta":{"raw":{"variants":["NLAFormer: self-attention as an affine linear algebra operator","NLAFormer: compact transformer runs CG with fewer iterations","NLAFormer: a transformer that solves linear systems via CG","Transformers learn matrix operations and even the conjugate gradient method","NLAFormer shrinks input size and layers for linear algebra"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.001955,"raw_usage":{"total_tokens":7631,"prompt_tokens":918,"completion_tokens":6713,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":534,"completion_tokens_details":{"reasoning_tokens":6628}},"tokens_in":534,"tokens_out":6713,"duration_ms":40962,"temperature":1.0,"reasoning_tokens":6628,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-15T16:51:01.073601+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"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$.","supporting_citations":[],"review_version":2}