Provably Learning Multi-Head Attention with Queries

Insung Kim; Jian Guo; Sunyeop Kim

Provably Learning Multi-Head Attention with Queries

T0 review · 0 major / 3 minor · reviewed 2026-08-10 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

Pith's one-line read Repeated-token trick recovers multi-head attention exactly.

desk verdict A genuinely nontrivial theoretical advance in query-based learning of multi-head attention, with an honest exact-arithmetic premise and a sound central proof. read the letter →

arxiv 2608.03294 v2 pith:C76NBRQ6 submitted 2026-08-04 cs.LG cs.CR

Sunyeop Kim , Insung Kim , Jian Guo This is my paper · ORCID

classification cs.LGcs.CR

keywords multi-headattentionsoftmaxparameterrecoveryvaluequeriesrationalinterpolationcanonicalrepresentationTransformerquerylearning

verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

This paper asks whether the parameters of a multi-head softmax attention layer can be recovered from black-box scalar outputs at the final token, without assuming the heads live in known orthogonal subspaces. It answers yes: by querying sequences that repeat the same token many times, the learner turns the hidden heads into the poles of a rational function, interpolates that function, and then uses additive queries to label the heads consistently. The recovery is exact up to permutation with probability one over randomly chosen query directions, using $4Hd^2-2H+1$ value queries of length at most $2H+1$ when $H$ is known. The paper also extends the construction to a one-layer ReLU Transformer and gives conditional linear stability under approximate outputs.

What carries the argument

The load-bearing object is the repeated-token identity $R(m)=\sum_{h=1}^H c_h r_h/(m+r_h)$, where $r_h=\exp(t u^\top W_h q)$ and $c_h=t u^\top v_h$. Varying the number $m$ of repeated tokens makes each hidden head appear as a simple pole of a rational function; rational interpolation and denominator factorization return the unordered set $\{(u^\top W_h q, u^\top v_h)\}$ for each direction pair $(u,q)$. Additive queries on sums of direction vectors then exploit bilinearity, for example $s_h(u_1+u_i,q_1)=s_h(u_1,q_1)+s_h(u_i,q_1)$, to match entries across local decoders, after which solving linear systems recovers $W_h$ and $v_h$.

What would settle it

Run Algorithm 1 on a fixed canonical model with rational parameters and an exact oracle, and test whether any draw of the random directions $(U,Q)$ with positive probability violates the distinctness or matching conditions of Lemma 4.3; finding such a draw would disprove the probability-one recovery claim, and exhibiting two distinct canonical models with identical scalar outputs on all length-three inputs would disprove Proposition 3.1.

Watch

Extended reading notes

Core claim

The central claim is that scalar final-token outputs uniquely determine, and an explicit query algorithm recovers, the canonical representation of a multi-head softmax attention model: heads sharing the same key-query matrix $W_h$ are merged by summing their value vectors $v_h$, and merged heads whose sum is zero are discarded. Equality on length-three inputs already pins down this representation. Algorithmically, a sequence with one perturbed token $q+tu$ followed by $m$ copies of $q$ makes the subtracted response $R(m)=F_M(X)-F_M([q])$ equal a rational function $\sum_{h=1}^H c_h r_h/(m+r_h)$; sampling $m=1,\dots,2H$, interpolating to recover the denominator, and factoring it separates the heads as an unordered set. Queries on sums of directions then give each head a consistent label, and linear systems recover every pair $(W_h,v_h)$. Under exact real arithmetic this succeeds with probability one over random orthogonal query directions, at the stated query cost.

Load-bearing premise

The whole guarantee rests on the oracle returning exact real values and the learner performing exact real arithmetic, including exact polynomial factorization; under ordinary IEEE 754 arithmetic, the authors' own experiments show recovery collapsing from near-perfect to zero for several configurations.

Editorial extensions

If this is right

Multi-head attention can be provably recovered from black-box scalar queries without orthogonal-subspace assumptions or known subspace bases.
The parameter count $H(d^2+d)$ and the query count $4Hd^2-2H+1$ are both order $H d^2$, so the recovery scales polynomially with the token dimension.
When only an upper bound $H_0$ on the raw head count is known, the same algorithm works with $H_0$ replacing $H$ and determines the true canonical head count from the smallest consistent denominator degree.
Under conditions guaranteeing separation and conditioning, an output error $\tau$ propagates to parameter error at most $C_{\mathrm{stab}}\tau$, so small oracle noise does not destroy recovery.
The same queries recover the effective attention heads of a one-layer ReLU Transformer via the odd-component reduction, and under extra conditions a functionally equivalent Transformer is constructed without a separate feed-forward learner.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

Going beyond the paper, replacing the exact polynomial-factorization step with certified numerical root isolation would convert the probability-one guarantee into a finite-precision routine with explicit error bounds.
Going beyond the paper, the same separation-and-matching strategy should apply to any attention kernel whose repeated-token response is a rational function with simple poles; if such kernels satisfy the analogous matching identities, the query count carries over unchanged.
Going beyond the paper, the near-total failure under IEEE 754 binary64 outputs for larger head counts suggests the exact-arithmetic oracle, not the subspace assumption, is the operative boundary; a practical deployment would need either high-precision oracles or a conditioning-aware variant.

Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Desk Editor's Note

read the letter

Dear colleague,

Here is my read on arXiv:2608.03294. The headline: this is a genuinely nontrivial advance over the prior single-head recovery framework. It removes the orthogonal-subspace and known-basis assumptions, recovers a canonical merged representation, and gives a clean nonadaptive query bound: 4Hd^2-2H+1 value queries, length at most 2H+1, probability one for the exact-arithmetic oracle. The repeated-token rational interpolation plus additive-identity matching across random bases is the real work, and it holds together.

What I like best: the proof is self-contained, and the central chain is coherent. Lemma 4.1's rational-function identity is simple and correct. Lemma 4.2's interpolation uniqueness is exactly right. Lemma 4.3's head-matching argument is sound. Appendix B's probability-one nondegeneracy argument handles the finite family of bad events cleanly. The experiments go beyond a token check: 180-digit arithmetic recovers to error below 10^-100 on all 800 sampled targets, and the controlled-perturbation slopes confirm the linear-in-tau claim. That is real supporting evidence.

The soft spots, in proportion: the exact-arithmetic assumption is genuinely load-bearing. The paper says so, and its own binary64 experiments show recovery collapsing from 100/100 to 0/100 for (3,8) and (8,4) even with 180-digit offline arithmetic. That makes the title's 'provably' more fragile in practice than a finite-precision reader might hope. But it is not an internal flaw: Theorem 5.1 states its premise exactly, and the approximate-output Theorem 5.2 is an honest local stability result with no uniform constant, which the two-head impossibility construction backs up. The ReLU Transformer extension is more speculative; it requires additional conditions, and the algorithm's termination has no uniform bound. I read these as scope limitations, not contradictions.

The citation pattern looks appropriate: it builds directly on Bhattamishra et al. 2026 and Tran et al. 2025, and the comparison table is fair.

Who is this for? Researchers in query-based model extraction and attention identifiability. It deserves a serious referee and, in my view, acceptance with revision. The appendix is long, but the main text is readable.

Recommendation: send it to review.

Referee Report

0 major / 3 minor

Summary. The paper studies exact recovery of the parameters of a multi-head softmax attention model from black-box value queries that return only the scalar output at the final token. It defines a canonical representation obtained by merging heads with the same query-key matrix W_h, summing their value vectors, and discarding zero aggregates. It proves (Proposition 3.1) that this canonical representation is identifiable from outputs on length-three inputs, and presents a randomized, nonadaptive algorithm (Algorithm 1) that, under exact oracle responses and exact real arithmetic, recovers every canonical pair (W_h, v_h) up to permutation with probability one, using exactly 4H d^2 - 2H + 1 value queries of maximum length 2H+1 (Theorem 5.1). The algorithm uses repeated-token queries to observe samples of a rational function whose poles and residues encode the per-head contributions, then uses additive identities from summed query directions to assign a common labeling across local decoders, after which linear systems recover the parameter pairs. The paper also proves a conditional stability bound for approximate oracle outputs (Theorem 5.2), shows that no model-independent stability constant exists, extends the recovery to a one-layer ReLU Transformer (Theorem 5.3), and provides variants for a known upper bound on the head count, low-rank W_h, and binary membership queries.

Significance. The result is a substantial advance over the prior art of Bhattamishra et al. (2026), which required pairwise orthogonal head subspaces and known orthonormal bases for those subspaces. The rational-interpolation technique is elegant and the proof chain, including the probability-one nondegeneracy argument in Appendix B and the rational-interpolation uniqueness lemma, appears sound. The paper is unusually honest about its computational model: it explicitly assumes exact real arithmetic and reports its own negative binary64 results. The conditional stability theorem and the lower bound showing that no uniform stability constant follows from the qualitative identifiability assumptions are valuable. The extension to a one-layer ReLU Transformer, while more involved and conditional, broadens the applicability of the query-based recovery approach.

minor comments (3)

[Section 6 / G.4] The paper reports that rounding oracle responses to IEEE 754 binary64 collapses the success rate from 100/100 to 0/100 for (d,H)=(3,8) and (8,4) even when offline arithmetic is 180-digit. A brief statement in the main text, rather than only in the appendix, that Theorem 5.1 is an exact-arithmetic statement and that the binary64 results are outside its scope would help calibrate reader expectations.
[Appendix D] The symbol A is used both for the set of query pairs in (29) and for the interpolation matrix A(y) in (38). This dual use makes the stability proof harder to follow; a different symbol for one of the two objects is advisable.
[Section 5, Theorem 5.2] The statement assumes exact subsequent computations, but the experiments use 180-digit arithmetic. The paper could note explicitly that the stability bound is a property of the exact algebraic reconstruction and does not by itself account for offline rounding errors.

Circularity Check

0 steps flagged · score 0.0 of 10

No significant circularity: the recovery proof is self-contained and does not reduce to fitted inputs or self-citations.

full rationale

The paper's derivation chain is self-contained against an external oracle. The canonical identifiability result (Proposition 3.1) is proved in Appendix A using the linear-independence lemmas A.1 and A.2, not imported from prior work. The rational interpolation lemma (Lemma 4.2) is a direct polynomial-degree argument, and the head-matching construction (Lemma 4.3) uses bilinear identities (10)-(11) together with probability-one nondegeneracy conditions proved in Appendix B. Query directions U and Q are drawn independently of the target model, and no fitted parameter is inserted into the recovery equations. The one-token outputs computed from the recovered v_h are derived consequences of those recovered values, not assumed inputs; they are used only after v_h has been determined from independent local decoders. The stability theorem (Theorem 5.2) is explicitly conditional on quantitative separation and conditioning constants, and the paper proves that no uniform bound can follow from the qualitative assumptions alone. The Transformer extension re-derives the odd-component reduction rather than merely citing it. The exact-real-arithmetic assumption is a clearly stated scope limitation, and the binary64 experimental collapse is honestly reported as a finite-precision phenomenon, not hidden or repackaged as a success. No load-bearing step reduces to its own inputs, and no self-citation chain is used to force the conclusion.

Assumptions & free parameters 1 free parameters · 8 assumptions · 0 invented entities

The central exact-recovery claim rests on the idealizations listed as axioms: exact real oracle outputs and exact arithmetic (including polynomial factorization), knowledge of d and H or H0, continuous random query bases, and the model conditions (pairwise distinct W_h, nonzero v_h). None of these are fitted to data; the only hand-chosen scalar is t, which the theorem permits to be any nonzero value. The approximate and Transformer extensions add further instance-dependent conditioning and generic-position assumptions. No new particles, forces, dimensions, or auxiliary models are introduced; the canonical representation is a proved equivalence class of the existing model, not an added entity.

free parameters (1)

repeated-token scale t = 1 (fixed; any nonzero t valid)
Hand-chosen query scale in Equation (2). It is not fitted to data and cancels from the interpolation; in finite precision it affects conditioning, which is why the authors fix t=1 in experiments.

assumptions (8)

domain assumption Oracle returns exact real values and the learner can perform exact real arithmetic, evaluate exponentials/logarithms, solve linear systems, and factor polynomials.
Section 3 'Computational assumptions'; without it the probability-one exact recovery and query count do not hold as stated. This is an idealization standard in query-learning theory but not achievable on finite-precision hardware.
domain assumption The target model has H canonical heads with pairwise distinct W_h and nonzero v_h.
Theorem 5.1 and the canonical representation require these identifiability conditions; merged heads are only recoverable as sums, and zero-sum groups are invisible.
domain assumption The learner knows d and H (or a known upper bound H0).
Problem statement in Section 3; the query count and the rational interpolation use H or H0 explicitly.
domain assumption Random query matrices U = Lambda_U O_U and Q = O_Q Lambda_Q are drawn from continuous distributions with singular values in [1,2].
Equation (6) and Appendix B; probability-one nondegeneracy of the local decoder and of the matching equalities hinges on this randomness.
domain assumption Exact univariate polynomial factorization is available as a primitive operation.
Section 4 'Computational assumptions' and offline complexity; the exact algorithm factors degree-H denominators, and the numerical version replaces this with companion-matrix eigenvalues.
standard math Standard algebraic facts: rational interpolation uniqueness (Lemma 4.2), identity theorem for meromorphic functions (Appendix A), and Gaussian/orthogonal measure-zero arguments (Appendix B).
These are unproved background results used in the proofs.
domain assumption For approximate outputs, positive separation and conditioning margins exist (r_min, delta_r, c_min, gamma_int, gamma_c, Delta_match) and downstream computations are exact.
Theorem 5.2 and Appendix D; the stability bound is conditional on these instance-dependent constants and does not hold uniformly, as the two-head construction in D.2 shows.
domain assumption For the one-layer ReLU Transformer extension, assumptions (i)-(iv) of Theorem 5.3 hold, including nonzero weights and directions, no proportional b_j, and Z one-to-one; termination is finite but non-uniform.
Appendix F; without these the feed-forward stage cannot be recovered, and the query count is not bounded uniformly.

how reviews work

0 comments

Cite this review

Pith. "Pith review of Provably Learning Multi-Head Attention with Queries." pith.science (2026). https://pith.science/paper/C76NBRQ6

@misc{pith2026260803294,
  author       = {Pith},
  title        = {Pith review of: Provably Learning Multi-Head Attention with Queries},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/C76NBRQ6}},
  note         = {Machine review of arXiv:2608.03294}
}

abstract

We study the problem of learning multi-head softmax attention from black-box input-output access. The learner may query arbitrary real-valued token sequences and observe only the scalar output at the final token. Recent work gives an algorithm using $O(d^2)$ value queries to recover the single-head parameters $(W,v)$. For multiple heads, the same work establishes identifiability under the assumption that the heads occupy pairwise orthogonal subspaces. Applying the single-head recovery algorithm separately to the heads additionally requires bases for these subspaces to be known. We recover a canonical representation by merging heads with the same $W_h$, summing their corresponding $v_h$, and discarding a merged head when this sum is zero, without these subspace assumptions. By varying the number of copies of a token, our algorithm obtains samples of a rational function whose interpolation separates the canonical heads. Additional queries formed by adding selected token vectors then match the same head across different queries. When the oracle outputs and all subsequent computations are exact, the learner chooses its query vectors at random and recovers the canonical pairs $\{(W_h,v_h):h\in[H]\}$ up to permutation with probability one. When $H$ is known, it uses exactly $4Hd^2-2H+1$ value queries of maximum length $2H+1$. If only a known upper bound $H_0$ is available, the algorithm uses $4H_0d^2-2H_0+1$ value queries of maximum length $2H_0+1$. For approximate oracle outputs, we give conditions under which the parameter error is at most a model- and query-dependent constant multiple of the output error. Finally, we extend our result to a one-layer Transformer with multi-head attention followed by a bias-free ReLU feed-forward network. Under additional conditions, we recover a functionally equivalent Transformer without relying on a separate algorithm for learning the feed-forward network.

Review history (2 revisions) →

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

27 extracted references · 2 linked inside Pith

[1]

2026 , note =

Provably Learning Attention with Queries , author =. 2026 , note =. 2601.16873 , archivePrefix =

arXiv 2026
[2]

Advances in neural information processing systems , volume=

Attention is all you need , author=. Advances in neural information processing systems , volume=
[3]

Tony and Zhang, Anru , title =

Cai, T. Tony and Zhang, Anru , title =. The Annals of Statistics , volume =. 2015 , doi =

2015
[4]

International Conference on Learning Representations , year =

Equivariant Neural Functional Networks for Transformers , author =. International Conference on Learning Representations , year =
[5]

Forty-third International Conference on Machine Learning , year=

Functional Equivalence in Attention: A Comprehensive Study with Applications to Linear Mode Connectivity , author=. Forty-third International Conference on Machine Learning , year=
[6]

Proceedings of the 57th Annual ACM Symposium on Theory of Computing , year =

Chen, Sitan and Li, Yuanzhi , title =. Proceedings of the 57th Annual ACM Symposium on Theory of Computing , year =
[7]

2025 , eprint =

Finite Samples for Shallow Neural Networks , author =. 2025 , eprint =

2025
[8]

2025 , eprint =

Provably Extracting the Features from a General Superposition , author =. 2025 , eprint =

2025

Show all 27 references

[9]

Linear Algebra and its Applications , volume =

Matrix Methods for Pade Approximation: Numerical Calculation of Poles, Zeros and Residues , author =. Linear Algebra and its Applications , volume =. 2018 , doi =

2018
[10]

Applied and Computational Harmonic Analysis , volume =

On the Accuracy of Prony's Method for Recovery of Exponential Sums with Closely Spaced Exponents , author =. Applied and Computational Harmonic Analysis , volume =. 2024 , doi =

2024
[11]

Proceedings of the 42nd International Conference on Machine Learning , pages =

Interpreting the Repeated Token Phenomenon in Large Language Models , author =. Proceedings of the 42nd International Conference on Machine Learning , pages =. 2025 , volume =

2025
[12]

Advances in Neural Information Processing Systems , volume=

Learning linear attention in polynomial time , author=. Advances in Neural Information Processing Systems , volume=
[13]

Machine Learning , volume =

Queries and Concept Learning , author =. Machine Learning , volume =. 1988 , doi =

1988
[14]

Advances in Neural Information Processing Systems , editor =

Recovering a Feed-Forward Net From Its Output , author =. Advances in Neural Information Processing Systems , editor =. 1993 , url =

1993
[15]

and Meka, Raghu , booktitle =

Chen, Sitan and Klivans, Adam R. and Meka, Raghu , booktitle =. Efficiently Learning One Hidden Layer
[16]

An Exact Poly-Time Membership-Queries Algorithm for Extracting a Three-Layer

Daniely, Amit and Granot, Elad , booktitle =. An Exact Poly-Time Membership-Queries Algorithm for Extracting a Three-Layer
[17]

Stealing Machine Learning Models via Prediction

Tram. Stealing Machine Learning Models via Prediction. 25th USENIX Security Symposium (USENIX Security 16) , pages =. 2016 , isbn =

2016
[18]

29th USENIX Security Symposium (USENIX Security 20) , pages =

High Accuracy and High Fidelity Extraction of Neural Networks , author =. 29th USENIX Security Symposium (USENIX Security 20) , pages =. 2020 , isbn =

2020
[19]

Proceedings of the 41st International Conference on Machine Learning , series =

Stealing Part of a Production Language Model , author =. Proceedings of the 41st International Conference on Machine Learning , series =. 2024 , url =

2024
[20]

Advances in Cryptology -- EUROCRYPT 2025 , editor =

Polynomial Time Cryptanalytic Extraction of Deep Neural Networks in the Hard-Label Setting , author =. Advances in Cryptology -- EUROCRYPT 2025 , editor =. 2025 , doi =

2025
[21]

2026 , howpublished =

Cryptanalytic Extraction of Deep Neural Networks with Non-Linear Activations , author =. 2026 , howpublished =

2026
[22]

2026 , howpublished =

Cryptanalytic Extraction of Recurrent Neural Network Models , author =. 2026 , howpublished =

2026
[23]

IEEE Std 754-2019 (Revision of IEEE 754-2008) , year =

2019
[24]

Dragan and Moritz Hardt , title =

Smitha Milli and Ludwig Schmidt and Anca D. Dragan and Moritz Hardt , title =. Proceedings of the Conference on Fairness, Accountability, and Transparency , pages =. 2019 , publisher =

2019
[25]

International Conference on the Theory and Application of Cryptology and Information Security , pages=

Hard-label cryptanalytic extraction of neural network models , author=. International Conference on the Theory and Application of Cryptology and Information Security , pages=. 2024 , organization=

2024
[26]

arXiv preprint arXiv:1908.04211 , year=

On identifiability in transformers , author=. arXiv preprint arXiv:1908.04211 , year=

1908 arXiv
[27]

Annual international cryptology conference , pages=

Cryptanalytic extraction of neural network models , author=. Annual international cryptology conference , pages=. 2020 , organization=

2020

Pith tools

Reviewed August 10, 2026 · model on record in the stance chip above.

[1] [1]

2026 , note =

Provably Learning Attention with Queries , author =. 2026 , note =. 2601.16873 , archivePrefix =

arXiv 2026

[2] [2]

Advances in neural information processing systems , volume=

Attention is all you need , author=. Advances in neural information processing systems , volume=

[3] [3]

Tony and Zhang, Anru , title =

Cai, T. Tony and Zhang, Anru , title =. The Annals of Statistics , volume =. 2015 , doi =

2015

[4] [4]

International Conference on Learning Representations , year =

Equivariant Neural Functional Networks for Transformers , author =. International Conference on Learning Representations , year =

[5] [5]

Forty-third International Conference on Machine Learning , year=

Functional Equivalence in Attention: A Comprehensive Study with Applications to Linear Mode Connectivity , author=. Forty-third International Conference on Machine Learning , year=

[6] [6]

Proceedings of the 57th Annual ACM Symposium on Theory of Computing , year =

Chen, Sitan and Li, Yuanzhi , title =. Proceedings of the 57th Annual ACM Symposium on Theory of Computing , year =

[7] [7]

2025 , eprint =

Finite Samples for Shallow Neural Networks , author =. 2025 , eprint =

2025

[8] [8]

2025 , eprint =

Provably Extracting the Features from a General Superposition , author =. 2025 , eprint =

2025