Uncovering mesa-optimization algorithms in Transformers. Johannes von Oswald,...

Uncovering mesa-optimization algorithms in Transformers
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
Статья: https://arxiv.org/abs/2309.05858
Тред от соавтора: https://twitter.com/oswaldjoh/status/1701873029100241241

Интересная свежая работа с попыткой разобраться, почему трансформеры хорошо работают в режиме in-context learning.

На данный момент уже есть пул работ, показывающих, что при in-context learning (то есть при forward pass, без всякого бэкпропа) внутри трансформера реализуется что-то типа градиентного спуска. Например, относительно недавно это было показано в работе “Transformers learn in-context by gradient descent” (https://arxiv.org/abs/2212.07677). Там показали, что один слой линейного self-attention’а (SA) может имплементировать один шаг градиентного спуска.

Это в чём-то похоже на meta-learning, но в отличие от традиционного мета-обучения, которое находится на уровень выше базового оптимизатора, этот вариант обучения находится на уровень ниже, здесь в сети при работе внешнего, базового, оптимизатора выучивается другой внутренний оптимизатор.

Благодаря уже классической работе “Risks from Learned Optimization in Advanced Machine Learning Systems“ (https://arxiv.org/abs/1906.01820), для этого есть термин mesa-optimization, буду по-русски называть его меза-обучением. У Роберта Майлса есть популярный рассказ про меза-оптимизацию и проблемы выравнивания (https://www.youtube.com/watch?v=bJLcIBixGj8). А если кто любит популярно текстом, то есть это https://www.lesswrong.com/posts/fLAvmWHmpJEiw8KEp/mesa-optimization-explain-it-like-i-m-10-edition. Что забавно, идея про название с меза- пришла из статьи Joe Cheal “What is the opposite of meta?” (https://anlp.org/knowledge-base/acuity-volume-2), которая из журнала про NLP, но не того NLP, про которое вы скорее всего подумали :)

Текущая работа развивает работу про обнаружение градиентного спуска и исследует авторегрессионные трансформеры с causally-masked attention (то есть causalLM, а не prefixLM), обученные на последовательностях. Показывают, что минимизация общего авторегрессионного лосса ведёт к появлению алгоритма градиентной оптимизации в forward pass трансформера. Ну что, уже начинаете бояться LLMок?

Подробнее про то, как оно всё устроено внутре.

Цель авторегрессионного трансформера при обучении минимизировать лосс предсказания текущего элемента последовательности по предыдущим. Здесь это просто сумма L2 лоссов по последовательности.

Слоёв теперь много, соответственно потенциально можно делать много шагов градиентной оптимизации. При этом всё хитрее, чем “K слоёв могут делать K шагов” -- получается т.н. unconventional online gradient-based optimizer, который работает хуже ванильного градиентного спуска. А с учётом другой свежей работы “CausalLM is not optimal for in-context learning” (https://arxiv.org/abs/2308.06912) про то, что causalLM стабильно недоперформит по сравнению с prefixLM, явно есть место для более эффективных меза-оптимизаторов (видимо, это будет темой отдельных статей).

Рассматриваемый меза-оптимизатор работает в два этапа.

1). Первый этап состоит из 1+ слоёв SA, где алгоритм реализует итеративную preconditioning procedure. Результатом является регуляризованная mesa-objective с лучшим числом обусловленности по сравнению с авторегрессионным лоссом.

2) На втором этапе последний SA слой делает единственный градиентный шаг относительно полученной mesa-objective.

Это теоретически обоснованная конструкция (ей посвящены большие приложения к статье) и в работе показывают, что обучение авторегрессионного трансформера ведёт именно к таким решениям.

Отдельная интересная ветка работы заключается в том, что если считать, что меза-оптимизация -- это желательное свойство модели, то можно сделать архитектурные модификации, которые сделают такое свойство встроенным по умолчанию.