Related papers: Provably Learning Attention with Queries

Provably Learning Attention with Queries

URL: http://arxiv.org/abs/2601.16873v1
Date: Fri, 23 Jan 2026 16:28:22 GMT
Title: Provably Learning Attention with Queries
Abstract summary: 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.
Score: 15.606567893781367
License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
Abstract: 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 corresponding real-valued output. We begin with the simplest case, a single-head softmax-attention regressor. We show that for a model with width $d$, there is an elementary algorithm to learn the parameters of single-head attention exactly with $O(d^2)$ 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, motivated by the regime where the head dimension $r \ll d$, we provide a randomised algorithm that learns single-head attention-based models 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 show that multi-head attention parameters are not identifiable from value queries in general -- distinct parameterisations can induce the same input-output map. Hence, guarantees analogous to the single-head setting are impossible without additional structural assumptions.

Related papers

Provably Learning Multi-Head Attention with Queries [5.8565739627322095]
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.
arXiv Detail & Related papers (2026-08-04T08:06:15Z)
Indexing: the Beginning and the End [45.94629852512322]
We study information bottlenecks in modern deep-learning architectures through the lens of the indexing primitive.<n>In this primitive, the input consists of $n$ bits and one integer $i$ from $1$ to $n$ called the index.<n>We show that architectures with low causal complexity cannot solve the indexing primitive in any constant number of layers when the index appears.
arXiv Detail & Related papers (2026-07-24T14:42:44Z)
Learning to Route Queries to Heads for Attention-based Re-ranking with Large Language Models [25.83277662723952]
We propose a query-dependent head selection method, RouteHead, for attention-based re-ranking with Large Language Models.<n>We learn a lightweight router that can map each query to an optimal head set, and relevance scores are computed by aggregating attention signals only from these heads.<n> Experiments on diverse benchmarks and multiple LLM backbones show that the proposed method consistently outperforms strong baselines.
arXiv Detail & Related papers (2026-04-27T15:36:54Z)
Storage capacity of perceptron with variable selection [10.64866985260943]
A central challenge in machine learning is to distinguish genuine structure from chance correlations in high-dimensional data.<n>We show that a simple perceptron can perfectly classify $P = N$ random patterns by optimally selecting $M = N$ variables out of $N$ variables.<n>This not only provides a quantitative criterion for distinguishing true structure in the data from spurious regularities, but also yields the storage capacity of associative memory models.
arXiv Detail & Related papers (2025-12-01T16:44:57Z)
Parameter-free Algorithms for the Stochastically Extended Adversarial Model [59.81852138768642]
Existing approaches for the Extended Adversarial (SEA) model require prior knowledge of problem-specific parameters, such as the diameter of the domain $D$ and the Lipschitz constant of the loss functions $G$.<n>We develop parameter-free methods by leveraging the Optimistic Online Newton Step (OONS) algorithm to eliminate the need for these parameters.
arXiv Detail & Related papers (2025-10-06T10:53:37Z)
DS-Det: Single-Query Paradigm and Attention Disentangled Learning for Flexible Object Detection [39.56089737473775]
We propose DS-Det, a more efficient transformer detector capable of detecting a flexible number of objects in images.<n>Specifically, we reformulate and introduce a new unified Single-Query paradigm for decoder modeling.<n>We also propose a simplified decoder framework through attention disentangled learning.
arXiv Detail & Related papers (2025-07-26T05:40:04Z)
Learning Compositional Functions with Transformers from Easy-to-Hard Data [63.96562216704653]
We study the learnability of the $k$-fold composition task, which requires computing an interleaved composition of $k$ input permutations and $k$ hidden permutations.<n>We show that this function class can be efficiently learned, with runtime and sample in $k$, by gradient descent on an $O(log k)$-depth transformer.
arXiv Detail & Related papers (2025-05-29T17:22:00Z)
Initialization is Critical to Whether Transformers Fit Composite Functions by Reasoning or Memorizing [10.206921909332006]
Transformers have shown impressive capabilities across various tasks, but their performance on compositional problems remains a topic of debate.<n>We discover that the parameter initialization scale plays a critical role in determining whether the model learns inferential (reasoning-based) solutions.<n>We further find that inferential (reasoning-based) solutions exhibit low complexity bias, which we hypothesize is a key factor enabling them to learn individual mappings for single anchors.
arXiv Detail & Related papers (2024-05-08T20:23:24Z)
Single-Stage Visual Relationship Learning using Conditional Queries [60.90880759475021]
TraCQ is a new formulation for scene graph generation that avoids the multi-task learning problem and the entity pair distribution. We employ a DETR-based encoder-decoder conditional queries to significantly reduce the entity label space as well. Experimental results show that TraCQ not only outperforms existing single-stage scene graph generation methods, it also beats many state-of-the-art two-stage methods on the Visual Genome dataset.
arXiv Detail & Related papers (2023-06-09T06:02:01Z)
Fast Interactive Search with a Scale-Free Comparison Oracle [17.38671584773247]
A comparison-based search algorithm lets a user find a target item $t$ in a database by answering queries of the form. We propose a scale-free probabilistic oracle model called $gamma$-CKL for such similarity triplets $(i,j;t)$. We develop a search algorithm with provably exponential rate of convergence under the $gamma$-CKL oracle, thanks to a backtracking strategy.
arXiv Detail & Related papers (2023-06-02T09:33:19Z)
Learning Hidden Markov Models Using Conditional Samples [72.20944611510198]
This paper is concerned with the computational complexity of learning the Hidden Markov Model (HMM) In this paper, we consider an interactive access model, in which the algorithm can query for samples from the conditional distributions of the HMMs. Specifically, we obtain efficient algorithms for learning HMMs in settings where we have query access to the exact conditional probabilities.
arXiv Detail & Related papers (2023-02-28T16:53:41Z)
Query Training: Learning a Worse Model to Infer Better Marginals in Undirected Graphical Models with Hidden Variables [11.985433487639403]
Probabilistic graphical models (PGMs) provide a compact representation of knowledge that can be queried in a flexible way. We introduce query training (QT), a mechanism to learn a PGM that is optimized for the approximate inference algorithm that will be paired with it. We demonstrate experimentally that QT can be used to learn a challenging 8-connected grid Markov random field with hidden variables.
arXiv Detail & Related papers (2020-06-11T20:34:32Z)
$O(n)$ Connections are Expressive Enough: Universal Approximability of Sparse Transformers [71.31712741938837]
We show that sparse Transformers with only $O(n)$ connections per attention layer can approximate the same function class as the dense model with $n2$ connections. We also present experiments comparing different patterns/levels of sparsity on standard NLP tasks.
arXiv Detail & Related papers (2020-06-08T18:30:12Z)
Non-Adaptive Adaptive Sampling on Turnstile Streams [57.619901304728366]
We give the first relative-error algorithms for column subset selection, subspace approximation, projective clustering, and volume on turnstile streams that use space sublinear in $n$. Our adaptive sampling procedure has a number of applications to various data summarization problems that either improve state-of-the-art or have only been previously studied in the more relaxed row-arrival model.
arXiv Detail & Related papers (2020-04-23T05:00:21Z)

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