One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention

要約

最近の研究では、コンテキスト内学習を実証的に分析し、合成線形回帰タスクで訓練されたトランスフォーマーが、十分な容量が与えられれば、ベイズ最適予測子であるリッジ回帰の実装を学習できることが示されました [Aky\’urek et al., 2023]。
線形セルフアテンションを備え、MLP 層を持たない 1 層トランスフォーマーは、最小二乗線形回帰目標で勾配降下 (GD) の 1 ステップを実装する方法を学習します [von Oswald et al., 2022]。
しかし、これらの観察の背後にある理論は依然としてよく理解されていません。
私たちは、合成ノイズを含む線形回帰データでトレーニングされた、線形自己注意の単一層を備えた変圧器を理論的に研究します。
まず、共変量が標準ガウス分布から抽出される場合、事前トレーニング損失を最小限に抑える 1 層変換器が最小二乗線形回帰目標で GD の単一ステップを実装することを数学的に示します。
次に、共変量と重みベクトルの分布を非等方性ガウス分布に変更すると、学習されたアルゴリズムに強い影響があることがわかります。事前トレーニング損失のグローバル ミニマイザーは、$\textit{pre の単一ステップを実装するようになりました。
-条件付き}$ GD。
ただし、応答の分布のみが変更される場合、これは学習されたアルゴリズムに大きな影響を与えません。応答が $\textit{nonlinear}$ 関数のより一般的なファミリーから来ている場合でも、
トレーニング前の損失でも、最小二乗線形回帰目標に対して GD の単一ステップが実装されます。

要約(オリジナル)

Recent works have empirically analyzed in-context learning and shown that transformers trained on synthetic linear regression tasks can learn to implement ridge regression, which is the Bayes-optimal predictor, given sufficient capacity [Aky\’urek et al., 2023], while one-layer transformers with linear self-attention and no MLP layer will learn to implement one step of gradient descent (GD) on a least-squares linear regression objective [von Oswald et al., 2022]. However, the theory behind these observations remains poorly understood. We theoretically study transformers with a single layer of linear self-attention, trained on synthetic noisy linear regression data. First, we mathematically show that when the covariates are drawn from a standard Gaussian distribution, the one-layer transformer which minimizes the pre-training loss will implement a single step of GD on the least-squares linear regression objective. Then, we find that changing the distribution of the covariates and weight vector to a non-isotropic Gaussian distribution has a strong impact on the learned algorithm: the global minimizer of the pre-training loss now implements a single step of $\textit{pre-conditioned}$ GD. However, if only the distribution of the responses is changed, then this does not have a large effect on the learned algorithm: even when the response comes from a more general family of $\textit{nonlinear}$ functions, the global minimizer of the pre-training loss still implements a single step of GD on a least-squares linear regression objective.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google