Une étape de descente de gradient est prouvée comme étant l'apprenant en contexte optimal avec une couche d'auto-attention linéaire

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

Résumé

Des travaux récents ont analysé empiriquement l'apprentissage en contexte et ont montré que les transformeurs entraînés sur des tâches synthétiques de régression linéaire peuvent apprendre à implémenter la régression ridge, qui est le prédicteur optimal au sens de Bayes, à condition d'avoir une capacité suffisante [Akyürek et al., 2023], tandis que les transformeurs à une seule couche avec une auto-attention linéaire et sans couche MLP apprendront à implémenter une étape de descente de gradient (GD) sur un objectif de régression linéaire des moindres carrés [von Oswald et al., 2022]. Cependant, la théorie derrière ces observations reste mal comprise. Nous étudions théoriquement les transformeurs avec une seule couche d'auto-attention linéaire, entraînés sur des données synthétiques de régression linéaire bruitées. Tout d'abord, nous montrons mathématiquement que lorsque les covariables sont tirées d'une distribution gaussienne standard, le transformeur à une couche qui minimise la perte de pré-entraînement implémentera une étape de GD sur l'objectif de régression linéaire des moindres carrés. Ensuite, nous constatons que modifier la distribution des covariables et du vecteur de poids vers une distribution gaussienne non isotrope a un impact significatif sur l'algorithme appris : le minimiseur global de la perte de pré-entraînement implémente désormais une étape de GD pré-conditionnée. Cependant, si seule la distribution des réponses est modifiée, cela n'a pas un grand effet sur l'algorithme appris : même lorsque la réponse provient d'une famille plus générale de fonctions non linéaires, le minimiseur global de la perte de pré-entraînement implémente toujours une étape de GD sur un objectif de régression linéaire des moindres carrés.

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.