Достижение бесконечного контекста с уменьшением объема памяти в 114 раз.
Делиться
1. Введение
За последние два года мы стали свидетелями гонки за увеличением длины последовательностей в языковых моделях искусственного интеллекта. Мы постепенно эволюционировали от 4000 символов в контексте до 32000, затем до 128000, и, наконец, до огромного окна в 1 миллион токенов, впервые обещанного такими моделями, как Gemini 1.5 pro. Обещание было заманчивым: загрузить в модель целые кодовые базы или романы и позволить ей рассуждать на основе всего этого.
Но у этой практически «бесконечной» длины контекста есть скрытая цена, о которой редко упоминают: память.
В стандартной архитектуре Transformer запоминание и анализ всего входного запроса не обходится без затрат. По мере роста входной последовательности модель должна хранить состояния «ключ-значение» (KV) для каждого отдельного токена, чтобы вычислять показатели внимания. Для последовательности из 1 миллиона токенов этот кэш KV может быстро разрастись до сотен гигабайт, что, в свою очередь, требует больших кластеров графических процессоров в нескольких центрах обработки данных, и все это только для того, чтобы хранить весь диалог в памяти.
2. Мотивация
В стандартном механизме внимания (Vaswani et al., 2017)6 каждый новый токен, генерируемый моделью, должен «оглядываться» на каждый предыдущий токен в подсказке, чтобы полностью понять контекст. Для обеспечения эффективности на протяжении нескольких поколений модель кэширует векторы ключа (K) и значения (V) предыдущих токенов в видеопамяти графического процессора. Это называется кэшем KV.
Ловушка линейного роста
Хотя кэширование векторов «ключ-значение» (кэш KV) может быть эффективным по времени (поскольку нам не нужно пересчитывать прошлые значения для каждого нового токена), оно требует огромного объема памяти, который линейно растет с длиной входной последовательности.
Для сравнения: для хранения кэша ключ-значение для стандартной модели с 500 байтами параметров и контекстом всего в 20 000 токенов требуется около 126 ГБ памяти. Если же масштабировать это до количества параметров современных моделей LLM (более 1 ТБ) и обеспечить обслуживание миллионов пользователей одновременно, то общий объем используемой памяти станет астрономически большим.
Исторически сложилось так, что у нас есть два способа обработки последовательных данных, ни один из которых не является идеальным:
- Рекуррентные нейронные сети (РНН) обрабатывают входной запрос токен за токеном, обновляя единственное и фиксированное скрытое состояние. Хотя это может значительно снизить требования к памяти, им трудно сохранять информацию и детали на протяжении длительных запросов. Это приводит к тому, что модели в конечном итоге забывают начало входной последовательности к моменту, когда доходят до конца.
- Трансформеры: В отличие от рекуррентных нейронных сетей (RNN), трансформеры не страдают от этой проблемы, поскольку они идеально запоминают все, храня всю историю разговора в кэше ключ-значение (KV-кэш). У них идеальная память, но из-за большого KV-кэша они требуют больших объемов памяти.
Именно такой компромисс и призвана компенсировать технология Infini-attention.
3. Решение: Бесконечное внимание
Для решения парадокса памяти исследователи из Google сформулировали концепцию бесконечного внимания (Munkhdalai et al., 2024)1. Основной принцип этого подхода заключается в том, что вместо хранения всего разговора мы можем хранить его краткое изложение.
Концепция «бесконечного внимания» разделяет выходной сигнал внимания на два отдельных механизма, работающих одновременно:
- Локальное внимание: аналогично стандартному трансформеру. Оно учитывает непосредственный контекст и вычисляет матрицу внимания для каждого токена, чтобы зафиксировать детали с высокой точностью.
- Глобальное линейное внимание: сжатая память, которая хранит сводку всей прошлой истории в матрице фиксированного размера, чтобы модель могла к ней обращаться.
Давайте рассмотрим, как обрабатывается длинный входной сигнал.
Шаг 1: Сегментация
Во-первых, вся входная последовательность делится на более мелкие сегменты (например, N=2048 токенов). В каждом сегменте модель использует стандартный механизм внимания с скалярным произведением для понимания контекста. Это гарантирует, что для непосредственных задач разрешение остается идеальным.
Шаг 2: Сжатие (обновление памяти)
Для перехода к следующему сегменту модель сохраняет сжатые состояния ключа (K) и значения (V) текущего сегмента в матрицу памяти фиксированного размера (M). Это позволяет модели обращаться к матрице памяти (вместо большего кэша KV) для получения информации о предыдущих сегментах.
Однако добавление новых данных в матрицу памяти вслепую может быстро испортить ранее хранившуюся в ней информацию. Чтобы предотвратить это, авторы используют правило Дельта (Schlag et al., 2021)7. Суть его заключается в следующем: прежде чем добавлять какую-либо новую информацию, необходимо проверить, хранит ли она уже эту информацию в памяти. Это позволяет избежать избыточных обновлений. Весь процесс обновления описан ниже:
А. «Peek» (Вычисление полученного значения)
Во-первых, модель извлекает значения из существующей памяти, используя текущие ключи (K) так, как если бы они были запросами. Модель делает это для того, чтобы определить, какая информация (значения) уже связана с текущими ключами в памяти.
B. Этап обновления
Затем модель сравнивает фактические новые значения (V) с полученными значениями (Vretrieved). Она вычисляет разницу (остаток) и добавляет в память только её. Это позволяет избежать обновления памяти уже имеющимися данными.
Это означает, что если память уже содержит информацию о текущем сегменте в полном объеме, обновление не требуется. Это обеспечивает стабильность и «чистоту» памяти при многочисленных обновлениях.
Шаг 3: Глобальный поиск (линейное внимание)
Для генерации следующего токена модели необходима контекстная информация из всего запроса, то есть из всех сегментов. Чтобы получить необходимую информацию, модель запрашивает матрицу памяти, выполняя матричное умножение.
Полученная матрица Amem содержит необходимую информацию из всех предыдущих сегментов для генерации следующего токена.
Шаг 4: Агрегация («Микшер»)
Наконец, модель имеет два выходных параметра:
- Adot: Подробный местный контекст из текущего сегмента.
- Амем: Сжатая глобальная история всех предыдущих сегментов из матрицы памяти.
Для объединения этих двух методов используется обученный скалярный параметр управления, β (бета):
Параметр β выступает в качестве коэффициента смешивания, определяющего компромисс между долгосрочными (Amem) и краткосрочными (Adot) информационными потоками:
- Когда β низкое: сигмоидная функция приближается к 0. Это приводит к тому, что дополнительный весовой коэффициент (1−sigmoid(β)) становится доминирующим, в результате чего модель отдает приоритет локальному вниманию на основе скалярного произведения (Adot) больше, чем глобальной компрессионной памяти.
- Когда β высокое: сигмоидная функция приближается к 1. Модель отдает приоритет извлеченному содержимому памяти (Amem), позволяя глобальному контексту переопределять локальную информацию из текущего сегмента.
4. Результаты: Почему важна концепция бесконечного внимания
Авторы протестировали механизм Infini-attention на существующих моделях с длинным контекстом, таких как Transformer-XL (Dai et al., 2019)² и Memorising Transformers (Wu et al., 2022)³. Ниже представлены результаты:
1. Сжатие памяти «114x».
Наиболее значимым достижением этой работы является существенное сокращение используемых ресурсов памяти. Поскольку Infini-Attention хранит весь исторический контекст в матрице памяти фиксированного размера, а не в линейно растущем кэше ключ-значение, он может позволить себе хранить в 114 раз меньше параметров в видеопамяти графического процессора по сравнению с запоминающими трансформерами. Как показано в таблице ниже, при длине контекста в 65 000 токенов Infini-Attention достигает лучших показателей перплексии в таких бенчмарках, как PG19 и Arxiv-math, при этом ему требуется хранить всего 1,6 млн параметров (размер матрицы памяти), в отличие от конкурирующих архитектур.
2. Тест «ключа» на миллион токенов.
Для архитектуры с длинным контекстом задача поиска иголки в стоге сена является общепринятой. Авторы проверили это, спрятав случайный пароль в огромном корпусе текста и попросив модель его найти. Как показано в таблице ниже, в условиях нулевого обучения модель испытывает трудности с поиском ключа, достигая в основном точности менее 20%.
Затем авторы доработали модель в течение 400 шагов, используя последовательности длиной всего 5000 токенов. Примечательно, что модель смогла обобщить результаты доработки для работы с последовательностями длиной до 1 миллиона токенов, значительно улучшив точность поиска в целом.
3. Современные методы составления кратких обзоров книг (500 тыс. контекста)
Помимо синтетических тестов, авторы также протестировали модель на бенчмарке BookSum (Kryściński et al.)5, где от модели требуется сгенерировать краткое содержание длинного романа. Модель Infini-Attention с 8B параметрами установила новый рекорд производительности на этом бенчмарке, успешно сгенерировав краткие содержания книг длиной до 500 000 токенов.
Результаты также показывают четкую тенденцию к улучшению способности модели к суммаризации по мере ввода в нее более длинных контекстов. График, представленный ниже, подтверждает эту гипотезу: вместо того, чтобы забывать предыдущую информацию (распространенный тип ошибки, известный как «потеря информации в середине»), модель может эффективно использовать матрицу памяти для генерации точных резюме.
4. Визуализация скалярной величины стробирующего элемента.
В качестве дополнительного исследования методом абляции авторы визуализировали изученный скалярный параметр управления (β), чтобы увидеть, как модель использует свою новую память. Ниже представлена тепловая карта полученной визуализации. Механизмы внимания разделились на две отдельные роли:
- Специализированные руководители: Руководители, имеющие оценку, близкую к 1 или 0, что указывает на их предпочтение либо локальному контексту (внутри сегмента), либо глобальной истории (предыдущие сегменты).
- «Смесительные головки»: головки с показателями, близкими к 0,5, что указывает на то, что их основная роль заключается в эффективном объединении информации из обоих путей.
Это говорит о том, что модель способна научиться переключаться между краткосрочным и долгосрочным запоминанием, а также смешивать информацию на протяжении всей последовательности.
5. Заключение
Хотя это и не сможет полностью заменить внешние векторные базы данных и системы RAG для рассуждений на основе статических знаний, это, тем не менее, меняет способ обработки стандартных пользовательских запросов моделями. Интеграция таких архитектур может стать следующим шагом вперед, позволяющим раскрыть исследовательский потенциал, который ранее был ограничен развитием аппаратного обеспечения, что в конечном итоге ускорит прогресс в области языкового моделирования.
👉Если вам понравилась эта статья, я публикую более короткие и актуальные материалы на Substack.
👉А если вы хотите поддержать написание независимых научных работ, BuyMeACoffee поможет вам в этом.
6. Список литературы
- Бесконечное внимание (Основная статья): Мункхдалай, Т., Фаруки, М., и Гопал, С. (2024). Не оставляйте контекст позади: эффективные преобразователи бесконечного контекста с бесконечным вниманием. Препринт arXiv:2404.07143.
- Transformer-XL: Дай, З., Ян, З., Ян, Ю., Карбонелл, Дж., Ле, К.В., и Салахутдинов, Р. (2019). Transformer-XL: Внимательные языковые модели за пределами контекста фиксированной длины. Препринт arXiv:1901.02860.
- Запоминание трансформеров: Ву, Ю., Рабе, М.Н., Хатчинс, Д., и Сегеди, К. (2022). Запоминание трансформеров. Препринт arXiv:2203.08913.
- Линейное внимание (математические основы): Катаропулос, А., Вьяс, А., Паппас, Н., и Флёре, Ф. (2020). Трансформеры — это рекуррентные нейронные сети: быстрые авторегрессивные трансформеры с линейным вниманием. Международная конференция по машинному обучению.
- Тест BookSum: Крыщиньски В., Раджани Н., Агарвал Д., Сюн К. и Радев Д. (2021). BookSum: коллекция наборов данных для обобщения подробного повествования. Препринт arXiv arXiv:2105.08209.
- Стандартное внимание: Васвани, Ашиш и др. «Внимания достаточно». Достижения в системах обработки нейронной информации 30 (2017).
- Правило Дельта: Шлаг, Иманол, Казуки Ирие и Юрген Шмидхубер. «Линейные трансформеры — это скрытые быстрые программисты весов». Международная конференция по машинному обучению. PMLR, 2021.
Источник: towardsdatascience.com
