Выявление алгоритмов меса-оптимизации в трансформерах

Uncovering mesa-optimization algorithms in Transformers

September 11, 2023

Авторы: Johannes von Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Blaise Agüera y Arcas, Max Vladymyrov, Razvan Pascanu, João Sacramento

cs.AI

Аннотация

Трансформеры стали доминирующей моделью в глубоком обучении, однако причины их превосходной производительности остаются малоизученными. В данной работе мы выдвигаем гипотезу, что высокая эффективность Трансформеров обусловлена их архитектурной предрасположенностью к меза-оптимизации — процессу, который обучается в рамках прямого прохода модели и состоит из двух этапов: (i) построения внутренней цели обучения и (ii) нахождения соответствующего решения через оптимизацию. Чтобы проверить эту гипотезу, мы проводим обратный инжиниринг ряда авторегрессивных Трансформеров, обученных на простых задачах моделирования последовательностей, и обнаруживаем лежащие в их основе градиентные алгоритмы меза-оптимизации, управляющие генерацией предсказаний. Более того, мы показываем, что обученный алгоритм оптимизации в прямом проходе может быть немедленно адаптирован для решения задач обучения с малым количеством примеров, что позволяет предположить, что меза-оптимизация может лежать в основе способности крупных языковых моделей к обучению в контексте. Наконец, мы предлагаем новый слой self-attention, меза-слой, который явно и эффективно решает задачи оптимизации, заданные в контексте. Мы обнаруживаем, что этот слой может улучшить производительность в синтетических и предварительных экспериментах по языковому моделированию, что подтверждает нашу гипотезу о том, что меза-оптимизация является важной операцией, скрытой в весах обученных Трансформеров.

English

Transformers have become the dominant model in deep learning, but the reason for their superior performance is poorly understood. Here, we hypothesize that the strong performance of Transformers stems from an architectural bias towards mesa-optimization, a learned process running within the forward pass of a model consisting of the following two steps: (i) the construction of an internal learning objective, and (ii) its corresponding solution found through optimization. To test this hypothesis, we reverse-engineer a series of autoregressive Transformers trained on simple sequence modeling tasks, uncovering underlying gradient-based mesa-optimization algorithms driving the generation of predictions. Moreover, we show that the learned forward-pass optimization algorithm can be immediately repurposed to solve supervised few-shot tasks, suggesting that mesa-optimization might underlie the in-context learning capabilities of large language models. Finally, we propose a novel self-attention layer, the mesa-layer, that explicitly and efficiently solves optimization problems specified in context. We find that this layer can lead to improved performance in synthetic and preliminary language modeling experiments, adding weight to our hypothesis that mesa-optimization is an important operation hidden within the weights of trained Transformers.