Provably Learning Multi-Head Attention with Queries

Insung Kim; Jian Guo; Sunyeop Kim

Provably Learning Multi-Head Attention with Queries

T0 review · 0 major / 4 minor · reviewed 2026-08-05 · deepseek-v4-flash

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

Pith's one-line read Multi-head attention parameters can be recovered exactly from black-box scalar outputs, without subspace assumptions.

desk verdict A genuinely new exact-recovery result for multi-head attention, with the caveats honestly placed where they belong: exact-real access and finite-precision fragility. read the letter →

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

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

classification cs.LGcs.CR MSC 68Q3268T07

keywords multi-headattentionvaluequeriesrationalinterpolationcanonicalrepresentationparameterrecoverysoftmaxTransformerquerylearning

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

The paper shows that the parameters of a multi-head softmax attention model can be extracted exactly from a black box that returns only the final-token scalar output, provided the learner may query arbitrary real-valued token sequences. It removes the orthogonal-subspace and known-basis assumptions of prior recovery algorithms, merging heads that share identical query-key matrices into a canonical form. The key move is to query sequences with many copies of a token, which makes the oracle response a rational function of the repetition count; interpolating this function separates the heads. The algorithm is nonadaptive, uses exactly $4Hd^2-2H+1$ value queries when the canonical head count $H$ is known, and succeeds with probability one. If only an upper bound on $H$ is known, the query count becomes $4H_0d^2-2H_0+1$.

What carries the argument

The repeated-token rational-function identity. For a fixed pair $(u,q)$, a sequence with one perturbed token $q+tu$ and $m$ copies of $q$ gives, after subtracting the one-token output, $R(m)=\sum_{h=1}^H \frac{c_h r_h}{m+r_h}$ with $c_h=tu^\top v_h$ and $r_h=\exp(tu^\top W_h q)$. Rational interpolation over the $2H$ values $R(1),\dots,R(2H)$ uniquely determines the numerator and denominator, and factoring the denominator recovers the unordered set $\{(s_h,c_h)\}$ for that pair. Bilinear identities such as $s_h(u_1+u_i,q_1)=s_h(u_1,q_1)+s_h(u_i,q_1)$ then assign a common head labeling across all queried pairs, after which linear systems recover $W_h$ and $v_h$.

What would settle it

Run Algorithm 1 with an oracle that rounds each scalar response to the nearest IEEE 754 binary64 value while keeping all offline arithmetic exact, for a canonical model with $(d,H)=(3,8)$. The paper's own experiment reports 0/100 runs recovering all heads, which would contradict the exact guarantee if it held for finite-precision outputs. Conversely, with 180-digit oracle responses the same configuration succeeds 100/100, isolating the oracle precision as the deciding factor.

Watch

Extended reading notes

Core claim

The central claim is Theorem 5.1: for a canonical multi-head attention model with pairwise distinct $W_h$ and nonzero $v_h$, an exact scalar-output oracle, and exact real arithmetic, Algorithm 1 recovers every pair $(W_h, v_h)$ up to permutation with probability one over its random query directions, using exactly $4Hd^2-2H+1$ value queries of maximum length $2H+1$. The recovery is built on the repeated-token identity $R(m)=\sum_{h=1}^H c_h r_h/(m+r_h)$, where $c_h=tu^\top v_h$ and $r_h=\exp(tu^\top W_h q)$. Rational interpolation over $m=1,\ldots,2H$ recovers the poles and residues, separating heads for a fixed pair $(u,q)$; bilinear identities then match heads across different query pairs,

Load-bearing premise

The learner has an oracle that returns exact real-valued outputs and can perform exact real arithmetic, including exponentials and polynomial factorization; if responses are rounded even to standard binary64 floating point, the theorem's probability-one guarantee collapses (the paper reports 0/100 success for $(d,H)=(3,8)$).

Editorial extensions

If this is right

Multi-head attention parameters can be extracted without knowing subspace bases or assuming orthogonal heads, eliminating the main structural restriction of prior query-based recovery.
The exact-query bound is $4Hd^2-2H+1$, matching the information-theoretic scale $H(d^2+d)$ up to a constant; with only an upper bound $H_0$ the cost becomes $4H_0d^2-2H_0+1$.
Equality of scalar final-token outputs on length-three inputs characterizes canonical equivalence, so the canonical representation is the right identifiability target for scalar-output attention.
For approximate outputs, parameter error is at most a model-dependent constant times the output perturbation, but no uniform constant exists over all identifiable models — the guarantee is inherently instance-dependent.
For a one-layer bias-free ReLU Transformer, the effective attention heads are recoverable with $2(4Hd^2-2H+1)$ Transformer queries, and under additional conditions a functionally equivalent Transformer can be constructed without a separate feed-forward-network learner.

Reading between the lines

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

The rational-interpolation mechanism suggests that query-length scaling with $2H+1$ is intrinsic to the separation strategy; a lower-bound argument on query complexity might be built from the need to resolve $H$ poles, though the paper does not prove one.
The binary64 failures indicate that applying the algorithm to real finite-precision APIs requires either higher-precision oracle responses, direct one-token queries to avoid error propagation, or a numerically stable reformulation — the paper notes the first two but does not provide an end-to-end finite-precision guarantee.
The head-matching step via random orthogonal matrices could be adapted to other bilinear sum-of-nonlinearities models, such as multi-head linear attention or kernel attention, wherever an analogous rational or polynomial structure arises.
The no-uniform-stability construction suggests that practical deployment should couple the algorithm with an a priori separation check or confidence scores, an open engineering question the paper leaves unresolved.

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

The central theorem checks out as far as I can tell. This paper really does remove the orthogonal-subspace and known-basis restrictions from multi-head attention recovery, and it gives an exact algorithm with a clean query count: 4Hd^2 - 2H + 1 value queries of length at most 2H + 1. That is a substantive step for query-based model extraction and attention identifiability.

What is actually new: the canonical representation that merges heads with the same W and drops zero sums, the repeated-token trick that turns oracle responses into samples of a rational function, and the cross-query additive matching that gives a consistent head labeling. The proofs are unusually complete. Rational interpolation uniqueness is proven directly, the probability-one nondegeneracy of the random directions is handled rigorously, including the dependent rows of the orthogonal draw, and the query count is exact. The stability section is also honest: Theorem 5.2 gives explicit constants, and Appendix D.2 shows by an information-theoretic construction that no uniform bound on C_stab follows from distinct W's and nonzero v's alone. That is the right way to state a conditional stability result.

The soft spots are real but mostly disclosed. The exact-real oracle and exact-arithmetic assumptions are standard for this line of work, but they mean the guarantee is about query complexity under ideal access rather than practical extraction. The paper's own experiments say this clearly: with IEEE 754 binary64 responses, success drops to 0/100 for (3,8) and (8,4), even with 180-digit offline arithmetic. That does not contradict Theorem 5.1, but it should temper any claim that the algorithm works on realistic APIs. The ReLU Transformer extension is conditional and has no uniform query bound, which the authors acknowledge. It is a useful extension, not the main selling point.

The citation pattern looks fine. The prior single-head result is cited properly, and the new material is genuinely different from it. I do not see a circular step: the heads are not assumed known, and the interpolation separates contributions from first principles.

Who is this for: researchers working on model extraction, attention identifiability, or query learning of neural networks. It deserves a serious referee. The core theorem looks sound, and the limitations are stated in the paper rather than hidden. I would accept it for peer review.

Referee Report

0 major / 4 minor

Summary. This paper considers the black-box problem of recovering the parameters of a multi-head softmax attention model from exact scalar final-token outputs. The model is parameterized by pairs (W_h, v_h); heads with identical W_h are merged by summing v_h, and zero sums are discarded, yielding a canonical representation. The main result (Theorem 5.1) is an algorithm that, with probability one over its random choice of U, Q, recovers all canonical pairs up to permutation using 4Hd^2 - 2H + 1 value queries of length at most 2H+1 when H is known. The algorithm uses repeated-token queries to sample a rational function, interpolates it to separate head contributions, and then uses additive identities in the bilinear scores to match heads across queries. The paper also provides an upper-bound version for unknown H, a conditional stability theorem for approximate outputs with model-dependent constants (and an information-theoretic lower bound showing no uniform constant), and an extension to a one-layer bias-free ReLU Transformer under additional conditions. The supplementary material contains complete proofs of the rational interpolation, probability-one nondegeneracy, stability, binary-membership conditional recovery, and the Transformer extension.

Significance. This is a strong theoretical contribution. It removes the orthogonal-subspace assumption and the need for known subspace bases that were required in prior work on multi-head attention recovery, and it establishes identifiability from final-token scalar outputs (Proposition 3.1). The exact recovery theorem is supported by complete proofs: the rational interpolation uniqueness (Lemma 4.2), the consistent head-labeling argument (Lemma 4.3), and the probability-one event in Appendix B are internally coherent. The exact query complexity is clean and matches the information-theoretic scale O(Hd^2). The paper also deserves credit for explicitly disclosing the limitations of its model: exact real arithmetic and exact oracle outputs are required for Theorem 5.1; the approximate-output guarantee is conditional and the stability constant is not uniform (Appendix D.2); and IEEE 754 binary64 outputs cause frequent failure (Appendix G.4). These disclosures are appropriately scoped and do not undermine the theoretical claims.

minor comments (4)

[Section 5, Theorem 5.1] The theorem states 'for any fixed t != 0', while the algorithm and implementation fix t=1. It would be clearer to state that t=1 is one valid choice of the fixed t and that the proof goes through for every nonzero t.
[Theorem 5.2 and Appendix D] The constants tau0 and Cstab are explicit functions of model- and query-dependent quantities (delta_r, gamma_int, gamma_c, Delta_match, etc.). Since these quantities are not known to the learner a priori, the main text should more prominently describe the theorem as a conditional, instance-wise stability result. The abstract does this, but the theorem statement in Section 5 could be clearer.
[Appendix F.1, Lemma F.1] The proof of critical-point refinement is correct but terse. A sentence explaining that the subdivision points are grid endpoints or midpoints at finer levels, and hence are avoided by the random shift with probability one, would help the reader follow the argument.
[Section 4] The notation D_t(u,q) is introduced as an unordered set, but the same symbol is later used for the local decoder procedure. Consider distinguishing the decoder procedure from its output to avoid ambiguity.

Circularity Check

0 steps flagged · score 0.0 of 10

No significant circularity; central derivation is self-contained.

full rationale

The paper's central claim (Theorem 5.1) is derived from first principles rather than from its own target parameters. The repeated-token identity (Lemma 4.1) expresses the observable difference R(m) as a rational function with poles -r_h and residues c_h r_h; the rational interpolation lemma (Lemma 4.2) establishes uniqueness of P and Q from 2H oracle values, and the poles are then recovered exactly. No fitted constant is renamed as a prediction, and no parameter of the target model is used in the reconstruction formulas (12), which require only the learner's chosen invertible matrices U,Q and the decoded scalar sets. The head-matching step (Lemma 4.3) is an algebraic consequence of bilinearity, with the required distinctness and nonzero conditions proved to hold with probability one in Appendix B; it does not assume the heads are known. The canonical representation is an identifiability target, not an input: the paper proves in Appendix A that scalar outputs on length-three inputs determine it. The approximate-output result (Theorem 5.2) is explicitly conditional on separation and conditioning constants, and Appendix D.2 proves that no uniform constant follows, so the stability bound is not a back-fitted claim. The one-layer ReLU Transformer extension proves the odd-component identities in Appendix F rather than importing them as an unverified ansatz. The only author self-citation (Chen et al. 2024, co-authored by Jian Guo, in the related-work list) is contextual and not load-bearing. The appendix also contains explicit limitations (e.g., no uniform query bound in Remark F.3, conditional nature of binary-membership recovery, and empirical failure under binary64 rounding), which further confirms that the claims are not being made circular by construction. Overall, the derivation chain is self-contained and the central recovery theorem does not reduce to its inputs.

Assumptions & free parameters 0 free parameters · 6 assumptions · 0 invented entities

The central recovery relies on the real-arithmetic oracle model and on standard analytic/algebraic facts (identity theorem, Rouché, Cauchy determinant, random-direction nondegeneracy). These are not introduced ad hoc; they are stated or proven. The canonical distinctness conditions (1) are part of the target definition, not hidden assumptions.

assumptions (6)

domain assumption Oracle returns exact scalar outputs; learner performs exact real arithmetic and exact polynomial factorization (Section 3, Computational assumptions).
The exact recovery theorem and query count assume this; finite-precision responses break the guarantee as shown in Appendix G.4.
domain assumption Canonical heads satisfy W_h pairwise distinct and v_h nonzero (Eq. (1)).
This is the definition of the target canonical representation; if heads coincide, only merged sums are identifiable.
standard math Random query directions from (6) avoid algebraic degeneracies with probability one (Appendix B).
Uses Gaussian/orthogonal invariance; nonzero polynomial events have measure zero.
standard math The identity theorem, Rouché's theorem, and the Cauchy determinant nonsingularity are used in Appendices A and D.
Standard analytic results invoked for linear independence and stability.
standard math For the low-rank variant, the rank-one projection operator satisfies the Cai-Zhang RIP (Appendix C.2).
Imported from the matrix recovery literature; not proven in this paper.
domain assumption Z is one-to-one and the b_j normals are pairwise nonproportional in the ReLU transformer theorem (Theorem 5.3(iv), Appendix F).
These are explicit additional assumptions for the Transformer extension.

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.

Discussion (0). Sign in to comment.

Reference graph

Works this paper leans on

27 extracted references · 24 canonical work pages

[1]

Provably Learning Attention with Queries

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

work page Pith review 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 =

work page 2015
[4]

International Conference on Learning Representations , year =

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

work page
[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=

work page
[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 =

work page
[7]

2025 , eprint =

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

work page 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 5, 2026 · model on record in the stance chip above.

[1] [1]

Provably Learning Attention with Queries

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

work page Pith review 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 =

work page 2015

[4] [4]

International Conference on Learning Representations , year =

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

work page

[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=

work page

[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 =

work page

[7] [7]

2025 , eprint =

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

work page 2025

[8] [8]

2025 , eprint =

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

2025