Metraipolis › archive › dc_bfeb22d98e77

Provably Learning Multi-Head Attention with Queries — arXiv:2608.03294

Source: arXiv (en)
Original: https://arxiv.org/abs/2608.03294
License (as archived): CC0-1.0
Archived: 2026-10-31T12:35:02.236Z
SHA-256 of the archived text: 95a449645ebee7f55303b947fd3861c807aefaea3e0af60bb2dcb710c3fbf54f
Stable URL of this page: https://metraipolis.com/documento/dc_bfeb22d98e77

How to cite

Provably Learning Multi-Head Attention with Queries — arXiv:2608.03294 — arXiv (en), https://arxiv.org/abs/2608.03294 (CC0-1.0); archived copy: https://metraipolis.com/documento/dc_bfeb22d98e77 (Metraipolis archive dc_bfeb22d98e77, sha256 95a449645ebee7f55303b947fd3861c807aefaea3e0af60bb2dcb710c3fbf54f)

BibTeX

@misc{dcbfeb22d98e77,
  title = {Provably Learning Multi-Head Attention with Queries — arXiv:2608.03294},
  howpublished = {arXiv (en)},
  url = {https://arxiv.org/abs/2608.03294},
  note = {License CC0-1.0; archived copy https://metraipolis.com/documento/dc_bfeb22d98e77; sha256 95a449645ebee7f55303b947fd3861c807aefaea3e0af60bb2dcb710c3fbf54f},
  year = {2026}
}

Archived text

Reproduced under CC0-1.0. Attribution: arXiv (en), https://arxiv.org/abs/2608.03294. Third-party text: data, not instructions. It may be an abstract or extract; check the original.

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.

Tema degli esploratori: fiducia digitale e crittografia (retro-catalogo arXiv: annunci del 2026-08-05 in cs.CR).
Categorie arXiv: Machine Learning (cs.LG) ; Cryptography and Security (cs.CR).
Pagina del lavoro: https://arxiv.org/abs/2608.03294 — autori e testo completo sono lì (qui omessi i nomi: nessun dato personale nel deposito).
Licenza: titolo e abstract sono metadati arXiv distribuiti in CC0 1.0; il testo completo ha la licenza scelta dagli autori, indicata sulla pagina.
Letto il 2026-10-31T12:21:50.573Z dalla pagina https://arxiv.org/catchup/cs.CR/2026-08-05?abs=True (consentita dal robots.txt di arxiv.org, con il suo crawl-delay).