Один шаг градиентного спуска является оптимальным методом обучения в контексте с одним слоем линейного самовнимания

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

July 7, 2023

Авторы: Arvind Mahankali, Tatsunori B. Hashimoto, Tengyu Ma

cs.AI

Аннотация

В последних работах эмпирически анализировалось обучение в контексте, и было показано, что трансформеры, обученные на синтетических задачах линейной регрессии, могут научиться реализовывать ридж-регрессию, которая является байесово-оптимальным предсказателем, при наличии достаточной емкости [Akyürek et al., 2023], в то время как однослойные трансформеры с линейным self-attention и без слоя MLP будут обучаться реализации одного шага градиентного спуска (GD) для задачи наименьших квадратов линейной регрессии [von Oswald et al., 2022]. Однако теория, лежащая в основе этих наблюдений, остается слабо изученной. Мы теоретически исследуем трансформеры с одним слоем линейного self-attention, обученные на синтетических данных зашумленной линейной регрессии. Во-первых, мы математически показываем, что когда ковариаты извлекаются из стандартного гауссовского распределения, однослойный трансформер, минимизирующий потери предобучения, будет реализовывать один шаг GD для задачи наименьших квадратов линейной регрессии. Затем мы обнаруживаем, что изменение распределения ковариат и вектора весов на не изотропное гауссовское распределение оказывает сильное влияние на изучаемый алгоритм: глобальный минимизатор потерь предобучения теперь реализует один шаг предобусловленного GD. Однако если изменяется только распределение откликов, это не оказывает значительного влияния на изучаемый алгоритм: даже когда отклик происходит из более общего семейства нелинейных функций, глобальный минимизатор потерь предобучения по-прежнему реализует один шаг GD для задачи наименьших квадратов линейной регрессии.

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.