Как устроено пространство, в котором думают языковые модели?
Изображение сгенерировано Dalle-3
Хабр, привет! Меня зовут Антон Разжигаев, я аспирант Сколтеха и участник научной группы Fusion Brain в институте AIRI.
С момента выхода первой статьи «Attention is All You Need» я с жадностью и любопытством, присущими любому исследователю, пытаюсь углубиться во все особенности и свойства моделей на базе архитектуры трансформер. Но, если честно, я до сих пор не понимаю, как они работают и почему так хорошо обучаются. Очень хочу разобраться, в чём же причина такой эффективности этих моделей, и есть ли предел их возможностей?
Такому изучению трансформеров «под микроскопом» и посвящена наша научная работа, только что представленная на конференции EACL 2024, которая проходила на Мальте — «The Shape of Learning: Anisotropy and Intrinsic Dimensions in Transformer-Based Models». В этой работе мы сфокусировались на наблюдении за пространством эмбеддингов (активаций) на промежуточных слоях по мере обучения больших и маленьких языковых моделей (LM) и получили очень интересные результаты.
Итак, приступим!
Данные
Начнём с рассказа о данных — это нужно для того, чтобы было проще понять, что мы сделали и что обнаружили. Т.к. нас интересовало пространство контекстуализированных эмбеддингов (в т.ч. промежуточных), надо было их где-то добыть.
Мы взяли enwik8 — аккуратно очищенные статьи Википедии на английском языке. Эти тексты мы прогнали через изучаемые модели, сохраняя все промежуточные активации (для каждого токена и с каждого слоя). Так мы получили «пространство эмбеддингов» или, другими словами, многомерное облако точек, с которым и стали дальше работать.
Чтобы исключить зависимость наблюдений от выбранного датасета, мы повторили эксперименты на случайных последовательностях токенов, и все выводы повторились. Поэтому в дальнейшем не буду акцентировать на этом внимание, а лучше сразу перейду к результатам.
Анизотропия
Один из самых главных вопросов, которые мы себе задали в процессе исследования —, а какая вообще форма у этих облаков точек? Визуализировать их сложно — пространство эмбеддингов очень многомерно, —, а методы снижения размерности тут не сильно помогают. Поэтому мы решили использовать анизотропию в качестве нашего «микроскопа». Анизотропия — это мера, показывающая насколько облако точек вытянуто, насколько оно неоднородно. Чем выше это значение, тем сильнее вытягивается пространство.
К примеру, уже давно известно (см. статью Representation Degeneration Problem in Training Natural Language Generation Models), что эмбеддинги трансформеров-энкодеров лежат в «узком конусе» — из-за этого косинусы между текстовыми репрезентациями всегда очень высокие. Но тем не менее, если вычесть среднее значение и отцентрировать это облако точек, оно становится вполне изотропным, то есть похожим на многомерный шарик. Поэтому говорят, что эмбеддинги трансформеров-энкодеров (Bert, RoBERTa, Albert, …) локально изотропны.
А в случае с декодерами (GPT, Llama, Mistral, …) мы обнаружили, что это совершенно не так! Даже после центрирования и использования более устойчивых к смещению методов на базе сингулярных чисел мы видим, что на средних слоях языковых моделей анизотропия практически равна 1. Это означает, что облако точек там вытянуто вдоль прямой линии. Но почему? Это же так сильно снижает ёмкость модели, она из-за этого практически не использует ТЫСЯЧИ других размерностей.
Наша реакция на экстремально высокие значения анизотропии.
Откуда берётся эта неоднородность пространства репрезентаций в декодерах, мы пока не знаем, но предполагаем, что это связано с процессом их обучения, задачей предсказания следующего токена и треугольной маской внимания. Это одна из исследовательских задач, которая сейчас стоит перед нами.
Если посмотреть на профиль анизотропии по слоям, то становится видно, что в начале и в конце декодеров эмбеддинги гораздо более изотропны, а экстремально высокая анизотропия наблюдается только в середине, где и должен происходить весь мыслительный процесс.
Локальная анизотропия разительно отличается в трансформерах-декодерах и -энкодерах.
Мы проследили за тем, как меняется анизотропия от чекпоинта к чекпоинту по мере обучения моделей (мы взяли все модели с промежуточными весами из того, что было опубликовано на тот момент). Оказалось, что все модели класса трансформер-декодер постепенно сходятся к одной и той же форме пространства и одному и тому же куполообразному профилю анизотропии.
По мере обучения анизотропия декодеров постепенно приходит к одному и тому же профилю.
Внутренняя размерность
Следующий наш «микроскоп» для наблюдения за активациями — внутренняя размерность. Это довольно красивое математическое понятие, описывающее «сложность» фигуры (многообразия или манифолда), на котором располагаются точки в многомерном пространстве.
Чтобы было понятнее, рассмотрим трёхмерную фигуру в виде ленты, свёрнутой в спираль (см. картинку ниже). Если мы приблизимся к какому-либо её участку, то обнаружим, что в малой окрестности точки будто бы лежат на плоскости. Следовательно, локальная внутренняя размерность тут равна двум.
Двухмерная лента в трехмерном пространстве.
Самое главное, что внутреннюю размерность довольно легко оценить, так как она сильно связана со скоростью роста «объёма» многомерного шара (количества точек данных, попадающих внутрь шара) по мере увеличения радиуса. Измерение зависимости количества точек от радиуса позволяет определить внутреннюю размерность в локальной области облака данных.
Итак, что же мы обнаружили? Во-первых, размерность довольно низкая, но это не новость, т.к. это было обнаружено и до нас. Во-вторых, — и это гораздо интереснее — эта размерность изменяется одинаково для всех моделей по мере обучения! Этот процесс состоит из двух фаз — сначала рост, а затем падение (см. график).
Изменение внутренней размерности промежуточных эмбеддингов трансформеров-декодеров по мере их обучения.
Похоже, что первая часть обучения переводит фичи в более высокие измерения, чтобы «запомнить» как можно больше информации, а во второй фазе — фичи начинают сжиматься, позволяя выявлять больше закономерностей, усиливая обобщающие способности модели.
Ещё раз — у всех LLM во время обучения присутствуют две фазы: инфляция эмбеддингов и их последующая компрессия.
Что это даёт?
Мы верим, что, вооружившись новым знанием, мы сможем улучшить процесс обучения языковых моделей (и не только), сделать его эффективнее, а сами трансформеры — быстрее и компактнее. Ведь если эмбеддинги проходят стадию компрессии и вообще стремятся расположиться вдоль одной линии, то почему бы просто не повыкидывать неиспользуемые измерения? Или помочь модели с более быстрым преодолением первой фазы.
Также мы обнаружили, что незадолго до взрывов лосса во время обучения (больная тема всех, кто учит LLM) внутренняя размерность сильно подрастает. Возможно, у нас получится предсказывать взрывы лосса и не тратить вычислительные ресурсы впустую, или вообще победить эти нестабильности, поняв их природу.
Хотя кого я обманываю, всё это нужно только ради удовлетворения своего любопытства!
Подписывайтесь на каналы авторов в телеграме AbstractDL, CompleteAI, Dendi Math&AI, Ivan Oseledets
Habrahabr.ru прочитано 11230 раз
