Um Passo de Descida de Gradiente é Comprovadamente o Aprendiz In-Contexto Ótimo com uma Camada de Auto-Atenção Linear

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

July 7, 2023

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

cs.AI

Resumo

Trabalhos recentes analisaram empiricamente o aprendizado em contexto e demonstraram que transformadores treinados em tarefas sintéticas de regressão linear podem aprender a implementar a regressão ridge, que é o preditor ótimo de Bayes, dada capacidade suficiente [Akyürek et al., 2023], enquanto transformadores de uma camada com auto-atenção linear e sem camada MLP aprenderão a implementar um passo de gradiente descendente (GD) em um objetivo de regressão linear de mínimos quadrados [von Oswald et al., 2022]. No entanto, a teoria por trás dessas observações ainda é pouco compreendida. Estudamos teoricamente transformadores com uma única camada de auto-atenção linear, treinados em dados sintéticos de regressão linear ruidosa. Primeiro, mostramos matematicamente que, quando as covariáveis são extraídas de uma distribuição gaussiana padrão, o transformador de uma camada que minimiza a perda de pré-treinamento implementará um único passo de GD no objetivo de regressão linear de mínimos quadrados. Em seguida, descobrimos que alterar a distribuição das covariáveis e do vetor de pesos para uma distribuição gaussiana não isotrópica tem um forte impacto no algoritmo aprendido: o minimizador global da perda de pré-treinamento agora implementa um único passo de GD pré-condicionado. No entanto, se apenas a distribuição das respostas for alterada, isso não terá um grande efeito no algoritmo aprendido: mesmo quando a resposta vem de uma família mais geral de funções não lineares, o minimizador global da perda de pré-treinamento ainda implementa um único passo de GD em um objetivo de regressão linear de mínimos quadrados.

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.