# Provably Learning Attention with Queries

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 22 of its 31 references that are in the map.
- First posted on arXiv 23 Jan 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Stealing Machine Learning Models via Prediction APIs](https://mlsanity.com/paper/1609.02943/) (2016)
- [Knockoff Nets: Stealing Functionality of Black-Box Models](https://mlsanity.com/paper/1812.02766/) (2018)
- [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)
- [In-Context Convergence of Transformers](https://mlsanity.com/paper/2310.05249/) (2023)
- [Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer](https://mlsanity.com/paper/2305.16380/) (2023)
- [Model Reconstruction from Model Explanations](https://mlsanity.com/paper/1807.05185/) (2018)
- [Transformers as Support Vector Machines](https://mlsanity.com/paper/2308.16898/) (2023)
- [Max-Margin Token Selection in Attention Mechanism](https://mlsanity.com/paper/2306.13596/) (2023)
- [Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot](https://mlsanity.com/paper/2406.06893/) (2024)
- [Provably learning a multi-head attention layer](https://mlsanity.com/paper/2402.04084/) (2024)
- [Implicit Regularization of Gradient Flow on One-Layer Softmax Attention](https://mlsanity.com/paper/2403.08699/) (2024)
- [ROP: Matrix recovery via rank-one projections](https://mlsanity.com/paper/1310.5791/) (2013)
- [Attention layers provably solve single-location regression](https://mlsanity.com/paper/2410.01537/) (2024)
- [Benign Overfitting in Single-Head Attention](https://mlsanity.com/paper/2410.07746/) (2024)
- [An Exact Poly-Time Membership-Queries Algorithm for Extraction a three-Layer ReLU Network](https://mlsanity.com/paper/2105.09673/) (2021)
- [Asymptotics of SGD in Sequence-Single Index Models and Single-Layer Attention Networks](https://mlsanity.com/paper/2506.02651/) (2025)
- [High-Dimensional Analysis of Single-Layer Attention for Sparse-Token Classification](https://mlsanity.com/paper/2509.25153/) (2025)
- [Learning Linear Attention in Polynomial Time](https://mlsanity.com/paper/2410.10101/) (2024)

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

- [Finite Probes Suffice: Identifiability and Universality for Weight-Space Learning](https://mlsanity.com/paper/2609.33901/) (2026)

## Abstract (arXiv)

We study the problem of learning Transformer-based sequence models with black-box access to their outputs. In this setting, a learner may adaptively query the oracle with any sequence of vectors and observe the output of the target function. We begin with studying the learnability of the simplest formulation, that is, learning a single-head attention-based regressor with queries. We show that for a model with width d, there is an elementary algorithm to learn the parameters of single-head attention with O(d2) queries. Further, we show that if there exists an algorithm to learn ReLU feedforward networks (FFNs), then the single-head algorithm can be easily adapted to learn one-layer Transformers with single-head attention. Next, we show that, in the common regime where the head dimension r ll d, single-head attention-based models can be learned with O(rd) queries via compressed sensing arguments. We also study robustness to noisy oracle access, proving that under mild norm and margin conditions, the parameters can be estimated to varepsilon accuracy with a polynomial number of queries even when outputs are only provided up to additive tolerance. Finally, we consider the learnability of multi-head attention and show that they are not identifiable from queries, and hence, learnability in the same sense is not feasible without additional assumptions. We discuss potential approaches to learn multi-head attention-based models under certain structural assumptions.