Één stap van Gradient Descent is bewezen de optimale in-context leerder met één laag lineaire zelf-attentie.

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

July 7, 2023

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

cs.AI

Samenvatting

Recente studies hebben in-context learning empirisch geanalyseerd en aangetoond dat transformatoren die getraind zijn op synthetische lineaire regressietaken, kunnen leren om ridge-regressie te implementeren, wat de Bayes-optimale voorspeller is, mits voldoende capaciteit [Akyürek et al., 2023], terwijl éénlaagse transformatoren met lineaire self-attention en zonder MLP-laag leren om één stap van gradient descent (GD) te implementeren op een kleinste-kwadraten lineaire regressiedoelstelling [von Oswald et al., 2022]. De theorie achter deze observaties blijft echter slecht begrepen. Wij bestuderen theoretisch transformatoren met één laag van lineaire self-attention, getraind op synthetische ruisachtige lineaire regressiegegevens. Eerst tonen we wiskundig aan dat wanneer de covariaten getrokken worden uit een standaard Gaussische verdeling, de éénlaagse transformator die het pre-trainingsverlies minimaliseert, één stap van GD zal implementeren op de kleinste-kwadraten lineaire regressiedoelstelling. Vervolgens ontdekken we dat het veranderen van de verdeling van de covariaten en de gewichtsvector naar een niet-isotrope Gaussische verdeling een sterke impact heeft op het geleerde algoritme: de globale minimaliseerder van het pre-trainingsverlies implementeert nu één stap van voorwaardelijke GD. Als echter alleen de verdeling van de responsen wordt veranderd, dan heeft dit geen groot effect op het geleerde algoritme: zelfs wanneer de respons afkomstig is uit een meer algemene familie van niet-lineaire functies, implementeert de globale minimaliseerder van het pre-trainingsverlies nog steeds één stap van GD op een kleinste-kwadraten lineaire regressiedoelstelling.

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.