[論文レビュー] Provably Learning Attention with Queries

Satwik Bhattamishra, Kulin Shah+2 more|arXiv (Cornell University)|Jan 23, 2026

Stochastic Gradient Optimization Techniques被引用数 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.

研究の動機と目的

• ブラックボックスの値クエリを用いた注意ベースのシーケンスモデルの学習を動機付け、形式化する。
• 多項式のクエリ複雑性で単頭注意の正確なパラメータ復元を示す。
• ReLU FFN 学習仮定の下で1層Transformerへ拡張する。
• 圧縮センサリングを用いてクエリ複雑性を減らす低ランク regime のアルゴリズムを開発する。
• 加法的 Oracle ノイズに対するロバスト性と、マルチヘッド注意の同一性可能性を分析する。

提案手法

• 単頭注意を f_{W,v}(X) = α(X,W)^{T}(Xv) として、 α は s_i = x_i^T W x_N のスコアをSoftmax で取る。
• 短いシーケンスでSoftmaxを分離し、線形方程式へ変換することで O(d^2) 回の値クエリで (W*,v*) の厳密な回復可能性を証明する。
• 2段階アプローチとFFN learner を組み合わせて、単頭注意を持つ1層Transformerを得る方法を示す。
• 低ランク領域 (rank(W*) ≤ r) では、ランク1の測定を設計し、圧縮センサリングを適用して O(rd) クエリで回復する。
• 近似的な値クエリを分析し、緩いノルム境界とマージン仮定の下で ε 正確性の回復を導出してロバスト性を扱う。
• 一般には値クエリからマルチヘッド注意パラメータの同一性は成立しないことを示す。

実験結果

リサーチクエスチョン

• RQ1 単頭ソフトマックス注意パラメータは値クエリから厳密に回復できるか。
• RQ2 埋め込み次元 d に対するクエリ複雑性は、単頭 vs 低ランク W* でどうスケールするか、圧縮で低減可能か。
• RQ3 FFN 学習をサブルーチンとして活用することで、値クエリアクセスを用いて1層Transformerを学習できるか。
• RQ4 ノイズのあるまたは近似的な出力をオラクルが返す場合、回復保証はどれだけロバストか。
• RQ5 値クエリからマルチヘッド注意パラメータは同一性を持つか。どの構造的仮定の下で同一性が達成可能か。

主な発見

• 単頭注意パラメータは O(d^2) の値クエリで厳密に回復可能（定理4.1）。
• 2段階アプローチにより、FFN値クエリ学習器が存在するという仮定のもと、単頭注意を持つ1層Transformerを得られる。
• 低ランク領域で rank(W*) ≤ r の場合、圧縮センサリングを介して O(rd) クエリで回復可能。
• 近似的な値クエリの場合でも、緩いノルム境界とマージン条件の下で ε 精度の回復が可能。
• 一般には値クエリからマルチヘッド注意パラメータの同一性は成立せず、追加の構造なしには単頭同様の保証を導くことはできない。