Un Paso de Descenso de Gradiente es Demostrablemente el Aprendizaje en Contexto Óptimo con una Capa de Autoatención Lineal

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

Resumen

Trabajos recientes han analizado empíricamente el aprendizaje en contexto y han demostrado que los transformadores entrenados en tareas sintéticas de regresión lineal pueden aprender a implementar la regresión de ridge, que es el predictor óptimo de Bayes, dada suficiente capacidad [Akyürek et al., 2023], mientras que los transformadores de una sola capa con autoatención lineal y sin capa MLP aprenderán a implementar un paso de descenso de gradiente (GD) sobre un objetivo de regresión lineal de mínimos cuadrados [von Oswald et al., 2022]. Sin embargo, la teoría detrás de estas observaciones sigue siendo poco comprendida. Estudiamos teóricamente transformadores con una sola capa de autoatención lineal, entrenados en datos sintéticos de regresión lineal ruidosa. Primero, demostramos matemáticamente que cuando las covariables se extraen de una distribución gaussiana estándar, el transformador de una capa que minimiza la pérdida de preentrenamiento implementará un paso de GD sobre el objetivo de regresión lineal de mínimos cuadrados. Luego, encontramos que cambiar la distribución de las covariables y el vector de pesos a una distribución gaussiana no isotrópica tiene un fuerte impacto en el algoritmo aprendido: el minimizador global de la pérdida de preentrenamiento ahora implementa un paso de GD precondicionado. Sin embargo, si solo se cambia la distribución de las respuestas, esto no tiene un gran efecto en el algoritmo aprendido: incluso cuando la respuesta proviene de una familia más general de funciones no lineales, el minimizador global de la pérdida de preentrenamiento sigue implementando un paso de GD sobre un objetivo de regresión lineal de mínimos cuadrados.

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.