Max-Margin Token Selection in Attention Mechanism

要約

アテンション メカニズムは、大規模言語モデルの驚異的な成功につながったトランスフォーマー アーキテクチャの中心的なコンポーネントです。
ただし、アテンション メカニズムの基礎となる理論原理、特にその非凸最適化ダイナミクスはほとんど理解されていません。
この研究では、独創的なソフトマックス注意モデル $f(\boldsymbol{X})=\langle \boldsymbol{Xv}, \texttt{softmax}(\boldsymbol{XWp})\rangle$ を探索します。ここで $\boldsymbol
{X}$ はトークン シーケンス、$(\boldsymbol{v},\boldsymbol{W},\boldsymbol{p})$ はトレーニング可能なパラメーターです。
$\boldsymbol{p}$、または同等の $\boldsymbol{W}$ で勾配降下法を実行すると、$\textit{局所的に最適な}$ トークンと非最適なトークンを分離する最大マージン解への方向に収束することを証明します。
もの。
これは明らかに、最適なトークン選択メカニズムとしての注意を形式化します。
注目すべきことに、私たちの結果は一般的なデータに適用でき、値埋め込み $\boldsymbol{Xv}$ と問題の幾何学に関してトークンの $\textit{optimality}$ を正確に特徴づけます。
また、非線形予測ヘッドに対しても注意のマージンを最大化する性質を確立する、より広範な正則化パス分析も提供します。
$\boldsymbol{v}$ と $\boldsymbol{p}$ をロジスティック損失と同時に最適化する場合、正則化パスがそれぞれのハードマージン SVM 解に方向的に収束する条件を特定します。
ラベルに基づいてフィーチャを入力します。
興味深いことに、$\boldsymbol{p}$ の SVM 定式化は、$\boldsymbol{v}$ のサポート ベクター ジオメトリの影響を受けます。
最後に、数値実験を通じて理論的発見を検証し、洞察を提供します。

要約(オリジナル)

Attention mechanism is a central component of the transformer architecture which led to the phenomenal success of large language models. However, the theoretical principles underlying the attention mechanism are poorly understood, especially its nonconvex optimization dynamics. In this work, we explore the seminal softmax-attention model $f(\boldsymbol{X})=\langle \boldsymbol{Xv}, \texttt{softmax}(\boldsymbol{XWp})\rangle$, where $\boldsymbol{X}$ is the token sequence and $(\boldsymbol{v},\boldsymbol{W},\boldsymbol{p})$ are trainable parameters. We prove that running gradient descent on $\boldsymbol{p}$, or equivalently $\boldsymbol{W}$, converges in direction to a max-margin solution that separates $\textit{locally-optimal}$ tokens from non-optimal ones. This clearly formalizes attention as an optimal token selection mechanism. Remarkably, our results are applicable to general data and precisely characterize $\textit{optimality}$ of tokens in terms of the value embeddings $\boldsymbol{Xv}$ and problem geometry. We also provide a broader regularization path analysis that establishes the margin maximizing nature of attention even for nonlinear prediction heads. When optimizing $\boldsymbol{v}$ and $\boldsymbol{p}$ simultaneously with logistic loss, we identify conditions under which the regularization paths directionally converge to their respective hard-margin SVM solutions where $\boldsymbol{v}$ separates the input features based on their labels. Interestingly, the SVM formulation of $\boldsymbol{p}$ is influenced by the support vector geometry of $\boldsymbol{v}$. Finally, we verify our theoretical findings via numerical experiments and provide insights.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google