Training Dynamics of In-Context Learning in Linear Attention

要約

注意ベースのモデルはコンテキスト内学習の顕著な能力を実証していますが、これらのモデルが勾配降下トレーニングを通じてこの能力を獲得した方法の理論的理解は依然として予備的です。
この質問への回答に向けて、コンテキスト内線形回帰のために訓練されたマルチヘッド線形自己触媒の勾配降下ダイナミクスを研究します。
線形自己触媒の2つのパラメーター化を調べます。1つはキーとクエリの重みが単一のマトリックス（理論研究で一般的）と統合され、もう1つは個別のキーとクエリ行列（実際の設定に近い）を使用します。
マージされたパラメーター化の場合、トレーニングダイナミクスには2つの固定ポイントがあり、損失軌道が単一の急激な低下を示すことを示します。
特定のクラスのデータセットと初期化の分析時間コースソリューションを導き出します。
個別のパラメーター化のために、トレーニングのダイナミクスが指数関数的に多くの固定点を持っていることを示し、損失はサドルからサドルのダイナミクスを示し、それをスカラーの通常の微分方程式に減らします。
トレーニング中、モデルは、トレーニング時間の経過とともに増加する主成分の数とのコンテキストで主成分回帰を実装します。
全体として、線形注意の勾配降下トレーニング中にコンテキスト内の学習能力がどのように進化するかを特徴付け、異なるパラメーター化を伴うモデルの急激な獲得と進行性の改善のダイナミクスを明らかにします。

要約(オリジナル)

While attention-based models have demonstrated the remarkable ability of in-context learning, the theoretical understanding of how these models acquired this ability through gradient descent training is still preliminary. Towards answering this question, we study the gradient descent dynamics of multi-head linear self-attention trained for in-context linear regression. We examine two parametrizations of linear self-attention: one with the key and query weights merged as a single matrix (common in theoretical studies), and one with separate key and query matrices (closer to practical settings). For the merged parametrization, we show the training dynamics has two fixed points and the loss trajectory exhibits a single, abrupt drop. We derive an analytical time-course solution for a certain class of datasets and initialization. For the separate parametrization, we show the training dynamics has exponentially many fixed points and the loss exhibits saddle-to-saddle dynamics, which we reduce to scalar ordinary differential equations. During training, the model implements principal component regression in context with the number of principal components increasing over training time. Overall, we characterize how in-context learning abilities evolve during gradient descent training of linear attention, revealing dynamics of abrupt acquisition versus progressive improvements in models with different parametrizations.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google