title: One Step of Gradient Descent is Provably the Optimal In-Context Learner
    with One Layer of Linear Self-Attention
summary: 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.
paper_authors._id: 64ab98caeaf3ada4ec9c795b
paper_authors.name: Arvind Mahankali
paper_authors.status: extracted_pending
paper_authors.statusLastChangedAt: 2024-05-23T03:22:15.899Z
paper_authors._id: 664eb64bda286d1a6015796b
paper_authors.avatarUrl: /avatars/f5c6060a234683dd9f5b355bb566352d.svg
paper_authors.fullname: Arvind Mahankali
paper_authors.type: user
paper_authors.user: amahanka
paper_authors._id: 64ab98caeaf3ada4ec9c795c
paper_authors.name: Tatsunori B. Hashimoto
paper_authors.status: extracted_pending
paper_authors.statusLastChangedAt: 2024-04-09T19:24:39.405Z
paper_authors._id: 661595d1b3d0b21da55cde7d
paper_authors.avatarUrl: /avatars/ba3fa065536518637d21a5c46cee5dd1.svg
paper_authors.fullname: Tatsu Hashimoto
paper_authors.type: user
paper_authors.user: thashim
paper_authors._id: 64ab98caeaf3ada4ec9c795d
paper_authors.name: Tengyu Ma
paper_title: One Step of Gradient Descent is Provably the Optimal In-Context Learner
    with One Layer of Linear Self-Attention
paper_summary: 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.
paper_ai_keywords: transformers linear self-attention ridge regression Bayes-optimal predictor gradient descent (GD) least-squares linear regression pre-training loss Gaussian distribution pre-conditioned GD nonlinear functions