선형 자기 주의력 단일 계층을 사용한 그래디언트 디센트의 한 단계가 최적의 인-컨텍스트 학습자임이 증명됨
One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention
July 7, 2023
저자: Arvind Mahankali, Tatsunori B. Hashimoto, Tengyu Ma
cs.AI
초록
최근 연구들은 컨텍스트 내 학습을 실증적으로 분석하며, 합성 선형 회귀 작업에 대해 훈련된 트랜스포머가 충분한 용량이 주어졌을 때 베이즈 최적 예측자인 릿지 회귀를 구현할 수 있음을 보여주었다[Akyürek et al., 2023]. 반면, 선형 자기 주의 메커니즘만을 갖고 MLP 층이 없는 단일 층 트랜스포머는 최소 제곱 선형 회귀 목표에 대한 경사 하강법(GD)의 한 단계를 학습한다[von Oswald et al., 2022]. 그러나 이러한 관찰 결과에 대한 이론적 배경은 아직 잘 이해되지 않고 있다. 우리는 합성 잡음 선형 회귀 데이터에 대해 훈련된 단일 층 선형 자기 주의 메커니즘을 가진 트랜스포머를 이론적으로 연구한다. 먼저, 공변량이 표준 가우시안 분포에서 추출되었을 때, 사전 훈련 손실을 최소화하는 단일 층 트랜스포머가 최소 제곱 선형 회귀 목표에 대한 GD의 한 단계를 구현함을 수학적으로 보인다. 그런 다음, 공변량과 가중치 벡터의 분포를 비등방성 가우시안 분포로 변경하면 학습된 알고리즘에 강한 영향을 미침을 발견한다: 이제 사전 훈련 손실의 전역 최소화자는 사전 조건화된 GD의 한 단계를 구현한다. 그러나 반응 변수의 분포만 변경된 경우, 이는 학습된 알고리즘에 큰 영향을 미치지 않는다: 반응 변수가 더 일반적인 비선형 함수군에서 나오더라도, 사전 훈련 손실의 전역 최소화자는 여전히 최소 제곱 선형 회귀 목표에 대한 GD의 한 단계를 구현한다.
English
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.