Un Passo di Discesa del Gradiente è Dimostrabilmente l'Apprenditore In-Contesto Ottimale con un Singolo Strato di Auto-Attenzione Lineare

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

July 7, 2023

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

cs.AI

Abstract

Recenti lavori hanno analizzato empiricamente l'apprendimento in contesto e dimostrato che i trasformatori addestrati su compiti sintetici di regressione lineare possono imparare a implementare la regressione ridge, che è il predittore Bayes-ottimale, data una capacità sufficiente [Akyürek et al., 2023], mentre trasformatori a un solo strato con auto-attenzione lineare e senza strato MLP imparano a implementare un singolo passo di discesa del gradiente (GD) su un obiettivo di regressione lineare ai minimi quadrati [von Oswald et al., 2022]. Tuttavia, la teoria alla base di queste osservazioni rimane poco compresa. Studiamo teoricamente i trasformatori con un singolo strato di auto-attenzione lineare, addestrati su dati sintetici di regressione lineare rumorosa. In primo luogo, dimostriamo matematicamente che quando le covariate sono estratte da una distribuzione gaussiana standard, il trasformatore a un solo strato che minimizza la perdita di pre-addestramento implementerà un singolo passo di GD sull'obiettivo di regressione lineare ai minimi quadrati. Successivamente, scopriamo che cambiare la distribuzione delle covariate e del vettore dei pesi in una distribuzione gaussiana non isotropica ha un forte impatto sull'algoritmo appreso: il minimizzatore globale della perdita di pre-addestramento ora implementa un singolo passo di GD pre-condizionato. Tuttavia, se viene modificata solo la distribuzione delle risposte, ciò non ha un grande effetto sull'algoritmo appreso: anche quando la risposta proviene da una famiglia più generale di funzioni non lineari, il minimizzatore globale della perdita di pre-addestramento implementa comunque un singolo passo di GD su un obiettivo di regressione lineare ai minimi quadrati.

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.