Как архитектура RMT масштабирует контекст Transformer до миллиона токенов

Yannic Kilcher 59,6 тыс. 24 мин 6 мин 27.04.2023
Главное

Популярный ИТ-блогер и исследователь Янник Килхер (Yannic Kilcher) разобрал громкий технический отчет исследователей Айдара Булатова, Юрия Куратова и Михаила Бурцева, пообещавших масштабировать контекст Transformer до 1 миллиона токенов и более с помощью технологии Recurrent Memory Transformer (RMT). В своем видео ведущий детально анализирует архитектурные особенности этого подхода, сопоставляет его с более ранними решениями вроде Transformer XL и объясняет, почему за громким маркетинговым заголовком скрывается хорошо знакомая специалистам концепция рекуррентных нейросетей. Писать об этом важно без отрыва от контекста развития больших языковых моделей, где длина контекста остается одним из главных технологических барьеров.

🧐 Маркетинг против реальности: суть технического отчета об RMT 0:00

Технический отчет под названием «Scaling Transformer to 1 million tokens and Beyond with RMT», авторами которого являются Айдар Булатов, Юрий Куратов и Михаил Бурцев, привлек широкое внимание ИИ-сообщества благодаря обещанию расширить контекст до 1 или даже 2 миллионов токенов. На первом же графике работы показано, что модель успешно справляется с задачами на удержание памяти и логику при гигантских объемах входных данных. Однако Янник Килхер сразу призывает снизить градус ожиданий, указывая на определенные нюансы в позиционировании работы. По словам ведущего, данная публикация представляет собой не принципиально новую архитектуру, а технический отчет, развивающий идеи их предыдущей статьи по модели RMT, которая была представлена на конференции NeurIPS в 2022 году.

Оригинальная работа 2022 года применяла рекуррентный механизм памяти к декодерным моделям (causally masked decoder-only Transformers), что требовало сложного разделения токенов на операции чтения и записи из-за однонаправленного внимания. Новый же отчет описывает применение той же технологии к энкодерным моделям, таким как BERT. Янник Килхер отмечает, что концептуально этот вариант даже проще, поскольку в нем отсутствует необходимость обрабатывать причинно-следственные маски (causal masking). По мнению блогера, этот метод не является ультимативным способом масштабирования стандартных трансформеров, как это может показаться на первый взгляд.

🧠 Проблема квадратичной сложности и исторический контекст 2:20

Классические трансформеры страдают от проблемы квадратичного масштабирования: вычислительная сложность и затраты памяти для механизма внимания (attention mechanism) растут пропорционально квадрату длины последовательности. Если увеличить количество токенов в два раза, потребуется в четыре раза больше памяти, а увеличение в четыре раза потребует уже шестнадцатикратного объема ресурсов, что приводит к быстрому исчерпанию памяти GPU. Чтобы обойти это ограничение, авторы RMT предлагают разбивать длинный текст на отдельные блоки или сегменты и обрабатывать их последовательно, связывая их между собой в процессе инференса.

Янник Килхер напоминает, что идея разбиения текста на сегменты далеко не нова: ИИ-сообщество разрабатывало подобные подходы еще в 2019–2020 годах, создавая такие архитектуры, как Transformer XL. Позже появились другие методы расширения длины контекста, включая Longformer и Big Bird. Как объясняет ведущий, эти модели обрабатывают всю последовательность целиком, но используют различные оптимизации внимания:

В отличие от Longformer, подход RMT полностью фокусируется на разделении текста на изолированные куски, обрабатывая каждый из них трансформером типа BERT независимо, но передавая накопленные знания дальше.

🎯 Тестовые задачи: как проверяли «память» модели 4:19

Для оценки эффективности RMT авторы использовали три типа синтетических задач, где полезная информация искусственно рассредоточена внутри огромного массива случайного «шума», взятого из Википедии.

Первая задача — простое запоминание (memorize task). В ней целевой факт гарантированно находится в самом первом сегменте текста, после чего идут тысячи токенов информационного шума. В стандартный BERT такой объем текста просто не поместился бы из-за жесткого ограничения на размер контекста.

Вторая задача — обнаружение и запоминание (detect and memorize task) — является более сложной. Здесь целевой факт может быть спрятан абсолютно в любом сегменте длинного текста, поэтому модели необходимо сначала локализовать его. Янник Килхер подчеркивает разницу подходов: стандартные длинноконтекстные трансформеры полагаются на прямое внимание от вопроса к месту хранения факта, в то время как RMT записывает этот факт в специальную выделенную память.

Третья задача — логическое рассуждение (reasoning task), основанное на наборе данных bAbI. Модели нужно сопоставить сразу два факта, разнесенных по тексту, чтобы ответить на финальный вопрос. Ведущий с иронией замечает, что пример, приведенный авторами в статье, выглядит ошибочным:

Факты в примере звучат так: «Коридор находится к востоку от ванной», «Спальня находится к западу от ванной». Вопрос: «К востоку от чего находится спальня?». Ответ: «Спальня».

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

🛠️ Как это устроено: токены памяти и рекуррентный механизм 7:56

Чтобы передавать информацию от сегмента к сегменту, авторы RMT добавляют к обычным текстовым токенам фиксированное количество специальных токенов памяти (memory tokens). Процесс обработки выглядит следующим образом: на вход слоя трансформера подаются токены текущего сегмента текста вместе со специальными слотами памяти.

Модель, инициализированная на базе предобученного BERT, настраивается (fine-tune) непосредственно под целевую задачу без повторного прохождения MLM-этапа. На выходе из трансформера генерируются обновленные токены памяти.

Эти выходные векторы берутся и перенаправляются на вход следующего сегмента. На финальном сегменте, где содержится вопрос, трансформер обращается к этим же токенам памяти через self-attention для извлечения сохраненной информации. Если же текущий сегмент не содержит важных данных, модель может просто транслировать токены памяти на выход в неизменном виде.

🔄 RMT против Transformer XL: сила обратного распространения ошибки во времени 10:38

Янник Килхер подробно сопоставляет RMT с Transformer XL, чтобы объяснить ключевое технологическое отличие. Архитектура Transformer XL сохраняет скрытые состояния (hidden states) предыдущего сегмента, позволяя текущему слою обращаться к ним через кросс-внимание. Однако в Transformer XL принудительно обрывается градиент (stop_gradient) при переходе между сегментами. Из-за этого при обратном распространении ошибки сеть не может передать сигнал вглубь истории. В результате прошлые слои физически не могут научиться эффективно выбирать и сохранять информацию специально для будущих шагов.

RMT решает эту проблему за счет полноценного обратного распространения ошибки во времени (Backpropagation Through Time, BPTT). Чтобы упростить понимание архитектуры, Килхер предлагает отбросить сложную терминологию и представить процесс абстрактно:

  1. Есть входной сегмент текста и начальное состояние памяти, заполненное нулями.
  2. Они подаются в нейросетевой блок, который выдает целевой результат и новое состояние памяти.
  3. Полученная память вместе со следующим сегментом передается в точно такой же блок с общими весами.

По заключению Янника Килхера, эта схема представляет собой классическую рекуррентную нейросеть (RNN), где в качестве базового вычислительного блока используется трансформер. Таким образом, утверждение о «масштабировании трансформеров до миллиона токенов» — это скорее маркетинговый ход: перед нами гибрид RNN и Transformer.

📈 Результаты обучения и реальные ограничения модели 19:48

Поскольку BPTT требует удержания в памяти GPU всех промежуточных состояний для расчета градиентов, развернуть сеть на тысячи сегментов во время обучения невозможно, поэтому его ограничивают несколькими шагами. Авторы RMT смогли обучать модель максимум на 7 сегментах текста — это предел, который вместился в память современных видеокарт. Главный прорыв RMT заключается в том, что обучившись всего на 7 сегментах, модель успешно выполняет инференс на длине в миллионы токенов. При инференсе градиенты вычислять не нужно, старые скрытые состояния можно удалять, поэтому требования к памяти растут линейно, а не квадратично.

Эксперименты показывают важную закономерность:

Янник Килхер подчеркивает, что этот подход фундаментально отличается от честного трансформера с полноценным контекстом. RMT эффективна лишь там, где нужно выудить редкие одиночные факты из гигантского объема нерелевантного шума. Визуализация карт внимания подтверждает этот вывод: при обработке вопроса модель целенаправленно считывает данные из сохраненных слотов памяти.

💬 Цитаты

«Концептуально этот вариант даже проще, поскольку в нем отсутствует необходимость обрабатывать причинно-следственные маски.»

Янник Килхер 01:53

«Эта схема представляет собой классическую рекуррентную нейросеть (RNN), где в качестве базового вычислительного блока используется трансформер.»

Янник Килхер 15:58

«RMT не способна читать книгу целиком и находить сложные взаимосвязи между сотнями разрозненных деталей.»

Янник Килхер 21:22
👥 Спикер
🔗 Упомянутые сайты и проекты
📖 Термины
Токены памяти (Memory tokens)
Специальные выделенные векторы, добавляемые в начало и конец сегмента текста для сохранения и передачи контекста между шагами рекурсии.
Обратное распространение ошибки во времени (BPTT)
Метод обучения рекуррентных сетей, при котором градиенты вычисляются путем развертывания сети на несколько шагов назад по истории.
Механизм внимания (Attention mechanism)
Компонент архитектуры Transformer, определяющий степень взаимосвязи и важности одних токенов относительно других в последовательности.
Инференс (Inference)
Этап работы уже обученной нейросети, на котором она выполняет предсказания для новых входных данных без обновления своих весов.
📊 Цифры
🗓 Хронология
  1. 2019-2020 годы Разработка первых подходов к сегментации контекста трансформеров, включая модель Transformer XL.
  2. 2022 год Презентация оригинальной архитектуры Recurrent Memory Transformer (RMT) на конференции NeurIPS.
⚖️ Другая сторона
Искусственный интеллект Transformer RMT Янник Килхер длина контекста BERT