[論文レビュー] In-Context Learning for Attention Scheme: from Single Softmax Regression to Multiple Softmax Regression via a Tensor Trick

Yeqi Gao, Zhao Song+1 more|arXiv (Cornell University)|Jul 5, 2023

Tensor decomposition and applicationsMathematics被引用数 3

ひとこと要約

本稿は、アテンションメカニズムにおけるコンテキスト内学習のためのテンソル・トリックに基づくフレームワークを提案する。行列ベースのソフトマックス回帰をベクトル化変換を用いて再定式化することで、安定性と一般化性能を向上させる。損失関数のリプシッツ連続性を確立し、モデル重みやパrameterの摂動が出力分布に有界なずれを引き起こすことを証明した。これは、正規化およびスケーリングされたソフトマックス定式化において、モデル容量および入力次元数に指数関数的依存する条件下でも、ロバストなコンテキスト内適応を可能にする。

ABSTRACT

Large language models (LLMs) have brought significant and transformative changes in human society. These models have demonstrated remarkable capabilities in natural language understanding and generation, leading to various advancements and impacts across several domains. We consider the in-context learning under two formulation for attention related regression in this work. Given matrices $A_1 \in \mathbb{R}^{n imes d}$, and $A_2 \in \mathbb{R}^{n imes d}$ and $B \in \mathbb{R}^{n imes n}$, the purpose is to solve some certain optimization problems: Normalized version $\min_{X} \| D(X)^{-1} \exp(A_1 X A_2^ op) - B \|_F^2$ and Rescaled version $\| \exp(A_1 X A_2^ op) - D(X) \cdot B \|_F^2$. Here $D(X) := \mathrm{diag}( \exp(A_1 X A_2^ op) {\bf 1}_n )$. Our regression problem shares similarities with previous studies on softmax-related regression. Prior research has extensively investigated regression techniques related to softmax regression: Normalized version $\| \langle \exp(Ax) , {\bf 1}_n angle^{-1} \exp(Ax) - b \|_2^2$ and Resscaled version $\| \exp(Ax) - \langle \exp(Ax), {\bf 1}_n angle b \|_2^2 $ In contrast to previous approaches, we adopt a vectorization technique to address the regression problem in matrix formulation. This approach expands the dimension from $d$ to $d^2$, resembling the formulation of the regression problem mentioned earlier. Upon completing the lipschitz analysis of our regression function, we have derived our main result concerning in-context learning.

研究の動機と目的

• ソフトマックス関数の行列形式を用いたアテンションベースの回帰におけるコンテキスト内学習を形式化すること。
• 重み空間および入力空間の両方におけるパrameter摂動下でのソフトマックス損失関数のリプシッツ連続性を分析すること。
• 正規化およびスケーリングされたソフトマックス回帰における、モデル重みや入力の変化が出力分布に与える影響の理論的境界を確立すること。
• 小さなパrameter更新がモデル出力に制御可能で有界な変化を引き起こすことを示し、安定したコンテキスト内適応を可能にすること。
• ベクトル化およびテンソル技術を用いて、行列形式のアテンションメカニズムに拡張されたソフトマックス回帰の先行研究を統合すること。

提案手法

• 本稿は、行列回帰問題を高次元のベクトル空間に変換するためのベクトル化技術を採用し、パrameter次元を $ d $ から $ d^2 $ に拡張する。
• 2つの定式化を導入する：正規化版 $ \min_X \|D(X)^{-1}\exp(A_1 X A_2^\top) - B\|_F^2 $ およびスケーリング版 $ \|\exp(A_1 X A_2^\top) - D(X)\cdot B\|_F^2 $、ここで $ D(X) $ は指数化された活性化の各行の和を対角成分とする対角行列である。
• 解析は、$ u(x), \alpha(x), f(x), h(x), c(x), q(x) $ およびそれらの逆関数の主要関数を有界化することに依存し、リプシッツ連続性を用いて感度を制御する。
• 損失関数の勾配に対するリプシッツ境界を導出し、出力の変化がパrameterの変化に比例することを示した。その割合は $ M = \exp(O(R^2 + \log n)) $ で表される。
• フレームワークは、$ x $ や $ \mathsf{A} $ における摂動が、出力分布に有界なずれを引き起こすことを証明した。$ \|\delta_c\|_2 \leq M \cdot \|x_{t+1} - x_t\|_2 $ が成り立つ。
• 関数ノルムの再帰的有界化と、行列指数関数および対角スケーリングの性質の活用により、理論的結果が導出された。

実験結果

リサーチクエスチョン

• RQ1 パラメータ行列を $ \mathbb{R}^{d^2} $ にベクトル化することで、ソフトマックスに基づくアテンション回帰の安定性および一般化性能にどのような影響を与えるか？
• RQ2 パrameter摂動下での正規化およびスケーリングされたソフトマックス損失関数の勾配のリプシッツ定数は何か？
• RQ3 モデル重みや入力の小さな変化が、アテンションメカニズムの出力分布に与える影響はどの程度有界になるか？
• RQ4 大規模言語モデルのコンテキスト内学習行動は、行列アテンション回帰のテンソル・トリック再定式化によって理論的に正当化可能か？
• RQ5 境界 $ M $ における $ R^2 $ および $ \log n $ への指数的依存性が、コンテキスト内適応のロバストネスにどのように影響を与えるか？

主な発見

• 正規化ソフトマックス損失関数は、パrameter ベクトル $ x $ および行列 $ \mathsf{A} $ においてリプシッツ連続性を示し、勾配の変化が $ M \cdot \|x_{t+1} - x_t\|_2 $ で有界である。ここで $ M = \exp(O(R^2 + \log n)) $ である。
• パrameter $ x $ を用いた更新では、出力のずれ $ \|\delta_c\|_2 $ が $ M \cdot \|x_{t+1} - x_t\|_2 $ で有界であるため、小さなパrameter変更下でも安定したコンテキスト内適応が保証される。
• パrameter $ \mathsf{A} $ を用いた更新では、出力のずれが同様に $ M \cdot \|A_{t+1} - A_t\|_2 $ で有界であり、重み摂動に対してロバストであることが確認された。
• 理論的解析により、$ x_t $ から $ x_{t+1} $ または $ A_t $ から $ A_{t+1} $ への移行が、摂動されたターゲット $ \widetilde{b} $ や $ \widehat{b} $ を持つ新たな安定な回帰問題を生成することを確認した。ここで $ \|\widetilde{b} - b\|_2 \leq M \cdot \|x_{t+1} - x_t\|_2 $ が成り立つ。
• 境界は、$ c(x) = \alpha(x)^{-1} u(x) $ のような中間関数へのリプシッツ性の再帰的適用により導出され、指数的および多項式的項を用いてノルムの成長がきめ細かく制御された。
• 本フレームワークは、モデル更新が出力にのみ制御可能で予測可能な変化を引き起こすことを示した。これにより、アテンションメカニズムにおけるコンテキスト内学習の理論的基盤が提供され、信頼性のある少数ショット一般化が可能になる。