Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Provably Learning Attention with Queries

Satwik Bhattamishra, Kulin Shah, Michael Hahn, Varun Kanade

TL;DR

This work addresses parameter recovery for Transformer-like attention blocks under value-query access, establishing provable guarantees for single-head softmax attention. It provides an exact recovery algorithm with $O(d^2)$ queries for the single-head case and a faster $O(rd)$-query procedure in the low-rank regime via compressed sensing and rank-one projections, plus robustness to additive noise with AVQ. A fundamental identifiability barrier for multi-head attention is shown, indicating that additional structural assumptions are needed for analogous guarantees. Overall, the results delineate what value queries can reveal about attention blocks and offer practical implications for model extraction and security analyses.

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.

Provably Learning Attention with Queries

TL;DR

This work addresses parameter recovery for Transformer-like attention blocks under value-query access, establishing provable guarantees for single-head softmax attention. It provides an exact recovery algorithm with queries for the single-head case and a faster -query procedure in the low-rank regime via compressed sensing and rank-one projections, plus robustness to additive noise with AVQ. A fundamental identifiability barrier for multi-head attention is shown, indicating that additional structural assumptions are needed for analogous guarantees. Overall, the results delineate what value queries can reveal about attention blocks and offer practical implications for model extraction and security analyses.

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 , there is an elementary algorithm to learn the parameters of single-head attention exactly with 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 , we provide a randomised algorithm that learns single-head attention-based models with 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 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.
Paper Structure (10 sections, 6 theorems, 49 equations)

This paper contains 10 sections, 6 theorems, 49 equations.

Table of Contents

  1. Introduction
  2. Related Work
  3. Definitions
  4. Learning Attention with Queries
  5. Learning 1-Layer Transformers
  6. Low Rank Recovery
  7. Background: Matrix Sensing
  8. Result for Low Rank
  9. Robustness and Stability of Learning
  10. Learnability of Multi-head Attention with Queries

Key Result

Theorem 4.1

Assuming ${\bm{v}}^\star \neq {\bm{0}}$, given access to a value query oracle $\operatorname{VQ}(X) = f_{{\bm{W}}^{\star},{\bm{v}}^{\star}}(X)$, the parameters ${\bm{W}}^\star$ and ${\bm{v}}^\star$ are exactly learnable in polynomial-time with $O(d^2)$ value queries.

Theorems & Definitions (10)

  • Theorem 4.1
  • proof
  • Theorem 5.1: Recovery under ROP; zhang_rop
  • Theorem 5.2
  • Lemma 6.1
  • proof
  • Theorem 6.2
  • proof
  • Proposition 7.1
  • proof