Дэн Фу рассказал о применении моделей пространства состояний в языковом моделировании

The TWIML AI Podcast 2,2 тыс. 35 мин 5 мин 22.05.2023
Главное

В эпоху стремительного развития генеративного искусственного интеллекта размер контекстного окна остается одним из главных ограничений современных больших языковых моделей. Пользователям часто хочется, чтобы нейросеть могла проанализировать массив их писем за несколько лет, прочесть огромную книгу или учесть все детали долгого разговора, однако современные архитектуры быстро упираются в вычислительный тупик. В рамках подкаста TWIML AI исследователь из Стэнфордского университета Дэн Фу (Dan Fu) подробно рассказал о своей научной работе «Hungry Hungry Hippos» (H3), представленной на конференции ICLR, и объяснил, как альтернативные математические примитивы могут избавить индустрию от диктатуры квадратичной сложности классического механизма внимания.

🧠 Проблема квадратичной сложности: почему ИИ «забывает» контекст 4:10

Современные языковые модели, такие как GPT-3 или GPT-4, совершают удивительные вещи, однако фундаментальные ограничения объема контекста долгое время не позволяли им обрабатывать документы длиннее нескольких страниц. Главная причина кроется в самом фундаменте архитектуры Трансформеров — механизме внимания (Attention), вычислительная сложность и требования к памяти которого растут квадратично относительно длины последовательности токенов.

Механизм внимания представляет собой своего рода метод «грубой силы» в обработке языка. Чтобы определить значение конкретного слова в предложении, алгоритм напрямую сравнивает его с каждым вторым словом в этой же структуре. Если бы нейросеть обрабатывала длинный диалог по такому принципу, то для генерации каждого последующего слова ей приходилось бы заново сопоставлять его со всей историей беседы. По мнению Дэна Фу, человеческая речь устроена иначе: люди не удерживают в памяти абсолютно каждое сказанное слово в явном виде, что и дает ученым надежду на существование более эффективных, субквадратичных подходов к моделированию языка.

⚡ Flash Attention: оптимизация на стыке ИИ и баз данных 8:06

До появления архитектурных альтернатив исследователи пытались решить проблему контекста с помощью системной оптимизации вычислений. Одним из прорывных решений в этой области стал алгоритм Flash Attention, разработанный Дэном Фу в соавторстве с коллегами по Стэнфорду примерно за год до текущего интервью.

В Flash Attention ученые применили классический подход из теории баз данных, известный как тайлинг (tiling). Вместо того чтобы проводить матричное умножение для всего объема данных целиком и сохранять гигантские промежуточные результаты, алгоритм осуществляет вычисления поблочно. Это позволило добиться следующих результатов:

Тем не менее, Дэн Фу подчеркивает, что системная оптимизация имеет свой предел: при переходе от 32 000 к 64 000 токенов вычислительные затраты все равно возрастают в четыре раза, что вынуждает искать новые базовые блоки для нейросетей.

🦛 Архитектура Hungry Hungry Hippos: от обработки сигналов к тексту 11:44

В поисках нового математического фундамента для языковых моделей авторы исследования под названием Hungry Hungry Hippos (H3) обратились к концепции моделей пространства состояний (State Space Models, SSM). Этот примитив на протяжении десятилетий успешно применяется в обработке сигналов для анализа длинных временных рядов (например, при фиксации сейсмической активности во время землетрясений), но ранее практически не использовался в лингвистических задачах.

Модели пространства состояний обладают уникальным математическим свойством: их можно одновременно рассматривать и как рекуррентные сети (подобные RNN или LSTM), и как сверточные нейросети (CNN). Это дает критически важные преимущества:

🧩 Преодоление качественного разрыва и тест на ассоциативное вспоминание 17:30

Простая замена механизма внимания на модель пространства состояний внутри стандартного Трансформера изначально приводила к резкому падению качества языкового моделирования. Разрыв составлял около 5 пунктов перплексии, что сопоставимо с разницей в качестве работы между миниатюрной моделью на 100 миллионов параметров и гигантской сетью на 10 миллиардов параметров.

Чтобы понять причину этой деградации, ученые создали искусственный синтетический язык для тестирования базовых навыков ИИ, названный задачей на «ассоциативное вспоминание» (associative recall). Модели предлагались простые строки, связывающие буквы и цифры (например, «А связано с 3»), а в конце строки шла проверка: нейросеть должна была динамически считать букву «А» и вспомнить правильную цифру из начала текста. Трансформеры справлялись с этим тестом со 100% точностью, тогда как стандартные модели пространства состояний полностью проваливали задачу.

Оказалось, что классическая SSM-структура способна запомнить, какой токен находился на определенной фиксированной позиции, но не умеет гибко и динамически определять, куда именно в тексте нужно посмотреть для извлечения связи. Решение, предложенное в работе Hungry Hungry Hippos, оказалось изящным: авторы соединили две модели пространства состояний последовательно («в стопку») и перемножили их выходные значения. Мультипликативное взаимодействие сработало как оператор направления внимания: одна модель формировала глобальную память для всей последовательности, а вторая указывала, какую именно ассоциацию требуется извлечь.

📊 Результаты бенчмарков и конвейер «гиен» 25:12

При масштабировании обновленной архитектуры H3 до реальных текстовых задач исследователи применили гибридный подход. В рамках сети, состоящей примерно из 30 слоев, авторы заменили почти весь механизм внимания на H3-слои, сохранив лишь два классических слоя внимания — один в начале и один в середине архитектуры.

Эксперименты показали, что такой гибрид способен превзойти стандартные Трансформеры по качеству работы на классических NLP-бенчмарках, таких как SuperGLUE, а также продемонстрировать более низкие показатели перплексии. Академическая модель H3 была успешно масштабирована до 3 миллиардов параметров. По словам Дэн Фу, от пользователей начали поступать неофициальные отзывы о том, что обученные по этой методике сети выдают более связный и логичный текст, чем классические Трансформеры аналогичного размера.

В полноценных тренировочных запусках исследователи довели длину контекста моделей H3 до 8000 токенов, а на синтетических тестах подтвердили работоспособность алгоритма на дистанциях в 32 000 и 128 000 токенов. Развитием этой идеи стал проект Hyena («Гиена»), в котором авторы полностью отказались от оставшихся слоев внимания в пользу чисто сверточного подхода, что позволило еще немного поднять планку производительности нейросетей. Сам ученый признает, что прямая переносимость этих результатов на коммерческие модели ультрагигантского масштаба (уровня триллионов параметров) требует проведения дополнительных тестов на массивных вычислительных кластерах, однако математический потенциал моделей пространства состояний уже сейчас выглядит серьезным вызовом для доминирующей архитектуры Трансформеров.

💬 Цитаты

«В языке нам не нужно помнить абсолютно все вещи, и это дает нам надежду на существование субквадратичного алгоритма.»

«Мы смогли убрать почти все слои внимания и получить то, что работает даже лучше классического Трансформера.»

👥 Спикеры
📖 Термины
Перплексия
Метрика оценки языковой модели, характеризующая степень ее уверенности при предсказании следующего слова (чем ниже показатель, тем лучше модель).
Тайлинг
Метод оптимизации вычислений, при котором крупные матрицы разбиваются на блоки (тайлы) для поочередной обработки без перегрузки памяти.
Модель пространства состояний (SSM)
Математический примитив из теории обработки сигналов, описывающий систему через набор внутренних скрытых состояний.
Ассоциативное вспоминание
Синтетический тест для ИИ, проверяющий способность модели находить ключи и связанные с ними значения в длинном тексте.
📊 Цифры
🗓 Хронология
  1. Около года назад Разработка и публикация алгоритма Flash Attention для линейной оптимизации памяти GPU.
  2. Текущий момент Презентация архитектуры Hungry Hungry Hippos (H3) на базе моделей пространства состояний на конференции ICLR.
  3. Ближайшее будущее Развитие сверточного подхода в рамках новой архитектуры Hyena и масштабирование моделей на сверхдлинный контекст.
⚖️ Другая сторона
Искусственный интеллект Дэн Фу State Space Models Flash Attention Transformers контекстное окно