Ein Schritt des Gradientenabstiegs ist nachweislich der optimale In-Context-Lerner mit einer Schicht linearer Selbstaufmerksamkeit.

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

July 7, 2023

Autoren: Arvind Mahankali, Tatsunori B. Hashimoto, Tengyu Ma

cs.AI

Zusammenfassung

Aktuelle Arbeiten haben das In-Context-Lernen empirisch analysiert und gezeigt, dass Transformer, die auf synthetischen linearen Regressionsaufgaben trainiert wurden, Ridge-Regression implementieren können, was der Bayes-optimale Prädiktor ist, sofern ausreichend Kapazität vorhanden ist [Akyürek et al., 2023], während einlagige Transformer mit linearem Self-Attention und ohne MLP-Schicht einen Schritt des Gradientenabstiegs (GD) auf ein kleinstes-Quadrate-lineares Regressionsziel lernen [von Oswald et al., 2022]. Die Theorie hinter diesen Beobachtungen bleibt jedoch weitgehend unverstanden. Wir untersuchen theoretisch Transformer mit einer einzelnen Schicht von linearem Self-Attention, die auf synthetischen verrauschten linearen Regressionsdaten trainiert wurden. Zunächst zeigen wir mathematisch, dass, wenn die Kovariaten aus einer Standard-Gauß-Verteilung gezogen werden, der einlagige Transformer, der den Pre-Training-Verlust minimiert, einen Schritt des GD auf das kleinstes-Quadrate-lineare Regressionsziel implementiert. Dann stellen wir fest, dass eine Änderung der Verteilung der Kovariaten und des Gewichtsvektors zu einer nicht-isotropen Gauß-Verteilung einen starken Einfluss auf den gelernten Algorithmus hat: Der globale Minimierer des Pre-Training-Verlusts implementiert nun einen Schritt des vorkonditionierten GD. Wenn jedoch nur die Verteilung der Antwortvariablen geändert wird, hat dies keinen großen Einfluss auf den gelernten Algorithmus: Selbst wenn die Antwortvariable aus einer allgemeineren Familie nichtlinearer Funktionen stammt, implementiert der globale Minimierer des Pre-Training-Verlusts weiterhin einen Schritt des GD auf ein kleinstes-Quadrate-lineares Regressionsziel.

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.