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

TL;DR

This paper provides a theoretical analysis showing that a single-layer linear self-attention transformer trained on linear regression data learns to perform one step of gradient descent, with the learned algorithm adapting to data distribution changes.

Contribution

It offers the first rigorous proof that such transformers implement gradient descent steps and explores how data distribution affects the learned algorithm.

Findings

Transformer with linear self-attention learns one step of GD.

Distribution of covariates influences whether GD is pre-conditioned.

Response distribution changes have limited impact on the learned algorithm.

Abstract

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.