# Provably Learning Multi-Head Attention with Queries

arXiv: https://arxiv.org/abs/2608.03294 · HTML: https://mlsanity.com/paper/2608.03294/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 13 of its 20 references that are in the map.
- First posted on arXiv 4 Aug 2026.

## Builds on (most cited first, 13 of 13)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Stealing Machine Learning Models via Prediction APIs](https://mlsanity.com/paper/1609.02943/) (2016)
- [High Accuracy and High Fidelity Extraction of Neural Networks](https://mlsanity.com/paper/1909.01838/) (2019)
- [Stealing Part of a Production Language Model](https://mlsanity.com/paper/2403.06634/) (2024)
- [On Identifiability in Transformers](https://mlsanity.com/paper/1908.04211/) (2019)
- [Model Reconstruction from Model Explanations](https://mlsanity.com/paper/1807.05185/) (2018)
- [Equivariant Neural Functional Networks for Transformers](https://mlsanity.com/paper/2410.04209/) (2024)
- [Provably learning a multi-head attention layer](https://mlsanity.com/paper/2402.04084/) (2024)
- [ROP: Matrix recovery via rank-one projections](https://mlsanity.com/paper/1310.5791/) (2013)
- [An Exact Poly-Time Membership-Queries Algorithm for Extraction a three-Layer ReLU Network](https://mlsanity.com/paper/2105.09673/) (2021)
- [Polynomial Time Cryptanalytic Extraction of Deep Neural Networks in the Hard-Label Setting](https://mlsanity.com/paper/2410.05750/) (2024)
- [Learning Linear Attention in Polynomial Time](https://mlsanity.com/paper/2410.10101/) (2024)
- [Functional Equivalence in Attention: A Comprehensive Study with Applications to Linear Mode Connectivity](https://mlsanity.com/paper/2606.17830/) (2026)

## Built on by (most cited first, 2 of 2)

- [Three Tokens Force Exponential Feature Rank in Nonnegative Kernel Attention](https://mlsanity.com/paper/2608.11427/) (2026)
- [The Identifiability and Observability of Deep Normalized Attention](https://mlsanity.com/paper/2610.09620/) (2026)

## Abstract (arXiv)

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(d2) 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 Wh, summing their corresponding vh, 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 (Wh,vh):h∈[H] up to permutation with probability one. When H is known, it uses exactly 4Hd2-2H+1 value queries of maximum length 2H+1. If only a known upper bound H0 is available, the algorithm uses 4H0d2-2H0+1 value queries of maximum length 2H0+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.