One Step of Gradient Descent is Provably the Optimal In-Context Learner
  with One Layer of Linear Self-Attention
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 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 nonlinear functions, the global
minimizer of the pre-training loss still implements a single step of GD on a
least-squares linear regression objective.
64ab98caeaf3ada4ec9c795b
Arvind Mahankali
extracted_pending
2024-05-23T03:22:15.899Z
664eb64bda286d1a6015796b
/avatars/f5c6060a234683dd9f5b355bb566352d.svg
Arvind Mahankali
user
amahanka
64ab98caeaf3ada4ec9c795c
Tatsunori B. Hashimoto
extracted_pending
2024-04-09T19:24:39.405Z
661595d1b3d0b21da55cde7d
/avatars/ba3fa065536518637d21a5c46cee5dd1.svg
Tatsu Hashimoto
user
thashim
64ab98caeaf3ada4ec9c795d
Tengyu Ma
One Step of Gradient Descent is Provably the Optimal In-Context Learner
  with One Layer of Linear Self-Attention
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 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 nonlinear functions, the global
minimizer of the pre-training loss still implements a single step of GD on a
least-squares linear regression objective.
Transformers with linear self-attention trained on synthetic linear regression tasks learn to implement gradient descent, with distributional changes affecting preconditioned gradient descent or nonlinear functions.
transformers linear self-attention in-context learning ridge regression gradient descent pre-training loss Gaussian distribution preconditioned GD nonlinear functions least-squares linear regression objective