[論文レビュー] What can a Single Attention Layer Learn? A Study Through the Random Features Lens

Hengyu Fu, Tianyu Guo+2 more|arXiv (Cornell University)|Jul 21, 2023

Domain Adaptation and Few-Shot LearningComputer Science被引用数 3

ひとこと要約

本稿は、クエリおよびキー行列をランダムに初期化して凍結し、値行列を学習可能とするランダム特徴設定における、単一のマルチヘッドアテンション層の表現能力および一般化特性を研究する。このモデルが、特にクエリ・キー行列の平均がゼロでない場合に、相関に基づく関数の近似が向上し、標準的な2層ネットワークよりも優れたサンプル複雑度を達成できることが示されている。

ABSTRACT

Attention layers -- which map a sequence of inputs to a sequence of outputs -- are core building blocks of the Transformer architecture which has achieved significant breakthroughs in modern artificial intelligence. This paper presents a rigorous theoretical study on the learning and generalization of a single multi-head attention layer, with a sequence of key vectors and a separate query vector as input. We consider the random feature setting where the attention layer has a large number of heads, with randomly sampled frozen query and key matrices, and trainable value matrices. We show that such a random-feature attention layer can express a broad class of target functions that are permutation invariant to the key vectors. We further provide quantitative excess risk bounds for learning these target functions from finite samples, using random feature attention with finitely many heads. Our results feature several implications unique to the attention structure compared with existing random features theory for neural networks, such as (1) Advantages in the sample complexity over standard two-layer random-feature networks; (2) Concrete and natural classes of functions that can be learned efficiently by a random-feature attention layer; and (3) The effect of the sampling distribution of the query-key weight matrix (the product of the query and key matrix), where Gaussian random weights with a non-zero mean result in better sample complexities over the zero-mean counterpart for learning certain natural target functions. Experiments on simulated data corroborate our theoretical findings and further illustrate the interplay between the sample size and the complexity of the target function.

研究の動機と目的

• 単一のマルチヘッドアテンション層が、良好なサンプル複雑度で学習可能な関数クラスを理解すること。
• 有限のサンプルと有限だが大きなヘッド数を用いた場合のランダム特徴アテンション（RFA）モデルの一般化性能を分析すること。
• クエリ・キー行列の初期化、特に非ゼロ平均の分布が、相関関数の学習効率に与える影響を調査すること。
• RFAのサンプル複雑度を標準的な2層ランダム特徴ネットワーク（RFMLP）と比較し、アテンションが優位を示す状況を同定すること。
• 理論的知見を、シミュレートされたデータを用いた実験を通じて検証し、サンプルサイズとターゲット関数の複雑度の相互作用を検討すること。

提案手法

• クエリベクトルが1つで、N個のキーべクトルを持つ単一アテンション層のメカニズムをモデル化し、ランダムかつ凍結されたクエリ・キー行列と学習可能な値行列を用いる。
• アテンション機構がクエリ・キー内積に基づくアテンションスコアから値の重み付き和を計算するランダム特徴アテンション（RFA）モデルを導入する。
• ターゲット関数 $f_\star$ の複雑さを測る $B(f_\star)$ とサンプルサイズ $n$ を用いて、過剰リスクバウンド $\widetilde{\mathcal{O}}(\sqrt{B(f_\star)/n})$ を導出する。
• クエリ・キー行列積に非ゼロ平均（例：単位行列）を持つ、バイアス付きRFA（BRFA）を提案する。これは実用的なアテンション初期化パターンをモデル化する。
• 理論的分析により、クエリ・キーべクトル間の相関に依存する関数に対して、BRFAは標準RFAよりも優れたサンプル複雑度を達成することが示された。
• シミュレートされたデータを用いた実験により、RFA、BRFA、RFMLP を、さまざまなシーケンス長 $N$、入力次元 $d$、ターゲット関数の複雑度の下で比較し、理論的サンプル複雑度の傾向を検証した。

実験結果

リサーチクエスチョン

• RQ1 クエリ・キー行列がランダムかつ凍結されている場合に、単一のマルチヘッドアテンション層は、広範な順列不変ターゲット関数クラスを表現可能か？
• RQ2 RFAモデルを用いた場合、このようなターゲット関数の学習におけるサンプル複雑度は何か？ また、標準的な2層ランダム特徴ネットワークと比較するとどうなるか？
• RQ3 クエリ・キー行列積の初期化分布（特に非ゼロ平均対ゼロ平均）が、相関関数の学習効率に与える影響は何か？
• RQ4 理論的サンプル複雑度バウンドは、現実的なモデルサイズとターゲット関数を想定した有限サンプル設定でも検証可能か？
• RQ5 RFAおよびBRFAを用いたターゲット関数の学習において、シーケンス長、入力次元、モデル性能の実験的相互作用は何か？

主な発見

• RFAモデルは、2つのトークンの関数の平均として表現可能な、広範な順列不変ターゲット関数クラスを表現可能であり、多項式関数や相関に基づく関数を含む。
• RFAの過剰リスクバウンドは $\widetilde{\mathcal{O}}(\sqrt{B(f_\star)/n})$ であり、これは入力次元にのみ依存し、キー・トークン数 $N$ には依存しない。このため、RFMLPに比べてサンプル複雑度が向上する。
• クエリ・キーべクトル間の相関に依存する関数に対しては、非ゼロ平均を持つバイアス付きRFA（BRFA）モデルが、標準RFAよりも優れたサンプル複雑度を達成し、$d \gg p$ の場合に $\text{Poly}(d)$ 要因の節約が可能となる。
• 実験により、RFAは特にシーケンス長 $N$ が増加する際、キーの順列不変関数の学習においてRFMLPを上回ることが確認され、理論的サンプル複雑度の優位性が裏付けられた。
• BRFAは相関に基づく関数の学習においてRFAを著しく上回り、相関強度（$\gamma$ で制御）が高くなるほど性能差が拡大する。これは非ゼロ平均初期化の理論的利点を確認するものである。
• 分析により、BERTの事前学習済みクエリ・キー行列は対角成分が支配的であることが判明し、実世界のアテンション機構が既に非ゼロ平均初期化の誘導的バイアスを活用している可能性があることが示唆され、本稿の理論的知見と整合的である。