Provably Learning Multi-Head Attention with Queries

• URL: http://arxiv.org/abs/2608.03294v2
• Date: Fri, 07 Aug 2026 05:20:15 GMT
• Title: Provably Learning Multi-Head Attention with Queries
• Abstract summary: We study the problem of learning multi-head softmax attention from black-box input-output access.<n>Recent work gives an algorithm using $O(d2)$ value queries to recover the single-head parameters $(W,v)$.<n>Applying the single-head recovery algorithm separately to the heads additionally requires bases for these subspaces to be known.
• Score: 5.8565739627322095
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• A Sharp Barrier for Consistent Submodular Maximization: Any Improvement over $2-\sqrt{2}$ Entails Exponential Queries or Linear Recourse [50.69285844345291]
  We study tradeoff between solution quality and stability when elements arrive over time.<n>Our algorithm has a future implementation for rational oracle$ with $O(varepsilon-1)$ recourse and separate the existence of universal-price certificates from their efficient algorithms.
  arXiv  Detail & Related papers  (2026-09-09T10:13:43Z)
• Small-Bias Quantum Approximate Counting via the Multiplicative Adversary Method [1.2799130177714182]
  We study the two-weight decision version of quantum approximate counting.<n>Given oracle access to $xin0,1N$, distinguish $|x|=M$ from $|x|=M+$ with success probability $1/2+$.<n>Using the multiplicative adversary method, we prove $left(maxleftsqrt(N-M)(M+)/,sqrtN/rightright)$.
  arXiv  Detail & Related papers  (2026-09-09T07:01:11Z)
• The Head Complexity of Boolean Functions in Single-Layer Attention [2.2727733134290813]
  We study the minimum number of attention heads required to compute a function in a one-layer attention-only model.<n>We establish an exact hierarchy under this measure: $k$ heads compute $k$-bit parity but cannot compute $(k+1)$-bit parity.<n>We also establish compactness bounds for embedding dimension and numerical precision.
  arXiv  Detail & Related papers  (2026-09-03T16:22:02Z)
• Near-Optimal Dimension Lower Bounds for Single-Vector Embeddings of Maximum Inner Product Similarity [75.14269295861845]
  For singleton queries, Chamfer becomes maximum inner product similarity (MAX-IP)<n>For every fixed $in(0,1)$, there are constants $A_,c_>0$.<n>Unit-sphere MAX-IP matrix is an exact two-valued affine image of the DNF pattern matrix with gap of at least $8$.
  arXiv  Detail & Related papers  (2026-07-22T17:27:20Z)
• Faster quantum linear system solver beyond the condition number [50.84794327094274]
  We present two quantum algorithms that produce the normalized solution $|xrangle$ of linear system $Ax=| b rangle$ to accuracy $$ with complexity independent of the condition number $=lVert A-1rVert$.<n>Our filtering-based solver is extremely simple with a favorable runtime prefactor.
  arXiv  Detail & Related papers  (2026-07-08T17:49:40Z)
• On the Computational Hardness of Transformers [14.73362105392153]
  We show that transformers can be computed more efficiently than $LH$ independent evaluations of attention.<n>In the small embedding regime, computing $LH$ attention heads separately takes $LHN2 + o(1)$ time.<n>In the large embedding regime, one can compute $LH$ attention heads separately using $LHN+ o(1)$ arithmetic operations.
  arXiv  Detail & Related papers  (2026-03-11T21:48:43Z)
• Provably Learning Attention with Queries [15.606567893781367]
  We study the problem of learning Transformer-based sequence models with black-box access to their outputs.<n>In this setting, a learner may adaptively query the oracle with any sequence of vectors and observe the corresponding real-valued output.
  arXiv  Detail & Related papers  (2026-01-23T16:28:22Z)
• Learning Multinomial Logits in $O(n \log n)$ time [56.23331174813387]
  A Multinomial Logit (MNL) model is composed of a finite universe of items $[n]=1,..., n$, each assigned a positive weight.<n>A query specifies an admissible subset -- called a slate -- and the model chooses one item from that slate with probability proportional to its weight.<n>This query model is also known as the Plackett-Luce model or conditional sampling oracle in the literature.
  arXiv  Detail & Related papers  (2026-01-07T22:07:44Z)
• Statistical and computational challenges in ranking [53.03724383992195]
  We consider the problem of ranking $n$ experts according to their abilities, based on the correctness of their answers to $d$ questions.<n>We investigate here the existence of statistically optimal and computationally efficient procedures for this problem.
  arXiv  Detail & Related papers  (2025-12-24T11:18:06Z)
• One-Query Quantum Algorithms for the Index-$q$ Hidden Subgroup Problem [1.4146420810689422]
  The Bernstein-Vazirani problem is an instance of the hidden subgroup problem (HSP)<n>We present the index-$q$ HSP: determine whether a hidden subgroup $H le G$ has index $1$ or $q$, and, when possible, identify $H$.
  arXiv  Detail & Related papers  (2025-10-12T10:42:50Z)
• Learning Partitions with Optimal Query and Round Complexities [16.815943270621638]
  We consider the basic problem of learning an unknown partition of $n$ elements into at most $k$ sets.<n>Non-adaptive algorithms require $Theta(n2)$ queries, while adaptive algorithms require $Theta(nk)$ queries.<n>Our algorithm only needs $O(log log n)$ rounds to attain the optimal $O(nk)$ query complexity.
  arXiv  Detail & Related papers  (2025-05-08T07:27:29Z)
• Active Sampling for Linear Regression Beyond the $\ell_2$ Norm [70.49273459706546]
  We study active sampling algorithms for linear regression, which aim to query only a small number of entries of a target vector. We show that this dependence on $d$ is optimal, up to logarithmic factors. We also provide the first total sensitivity upper bound $O(dmax1,p/2log2 n)$ for loss functions with at most degree $p$ growth.
  arXiv  Detail & Related papers  (2021-11-09T00:20:01Z)
• Learning a Latent Simplex in Input-Sparsity Time [58.30321592603066]
  We consider the problem of learning a latent $k$-vertex simplex $KsubsetmathbbRdtimes n$, given access to $AinmathbbRdtimes n$. We show that the dependence on $k$ in the running time is unnecessary given a natural assumption about the mass of the top $k$ singular values of $A$.
  arXiv  Detail & Related papers  (2021-05-17T16:40:48Z)
• Query complexity of heavy hitter estimation [6.373263986460191]
  We consider the problem of identifying the subset $mathcalSgamma_mathcalP$ of elements in the support of an underlying distribution $mathcalP$. We consider two query models: $(a)$ each query is an index $i$ and the oracle return the value $X_i$ and $(b)$ each query is a pair $(i,j)$. For each of these query models, we design sequential estimation algorithms which at each round, either decide what query to send to the oracle depending on the entire
  arXiv  Detail & Related papers  (2020-05-29T07:15:46Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.